Support optional AS keyword in CTE definitions for Databricks

Databricks allows omitting the AS keyword in CTE definitions:
`WITH cte (SELECT ...) SELECT * FROM cte`

Add `supports_cte_without_as()` dialect method and enable it for
Databricks and Generic dialects.

Author: funcpp

src/dialect/databricks.rs

@@ -90,4 +90,9 @@ impl Dialect for DatabricksDialect {
     fn supports_optimize_table(&self) -> bool {
         true
     }
+
+    /// See <https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-qry-select-cte>
+    fn supports_cte_without_as(&self) -> bool {
+        true
+    }
 }

src/dialect/generic.rs

@@ -288,4 +288,8 @@ impl Dialect for GenericDialect {
     fn supports_comma_separated_trim(&self) -> bool {
         true
     }
+
+    fn supports_cte_without_as(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -1664,6 +1664,17 @@ pub trait Dialect: Debug + Any {
     fn supports_comma_separated_trim(&self) -> bool {
         false
     }
+
+    /// Returns true if the dialect supports the `AS` keyword being
+    /// optional in a CTE definition. For example:
+    /// ```sql
+    /// WITH cte_name (SELECT ...)
+    /// ```
+    ///
+    /// [Databricks](https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-qry-select-cte)
+    fn supports_cte_without_as(&self) -> bool {
+        false
+    }
 }
 
 /// Operators for which precedence must be defined.

src/parser/mod.rs

@@ -14060,7 +14060,7 @@ impl<'a> Parser<'a> {
         })
     }
 
-    /// Parse a CTE (`alias [( col1, col2, ... )] AS (subquery)`)
+    /// Parse a CTE (`alias [( col1, col2, ... )] [AS] (subquery)`)
     pub fn parse_cte(&mut self) -> Result<Cte, ParserError> {
         let name = self.parse_identifier()?;
 
@@ -14090,9 +14090,33 @@ impl<'a> Parser<'a> {
                 materialized: is_materialized,
                 closing_paren_token: closing_paren_token.into(),
             }
+        } else if self.dialect.supports_cte_without_as()
+            && self.peek_token().token == Token::LParen
+            && self.is_query_start_token(&self.peek_nth_token(1))
+        {
+            self.expect_token(&Token::LParen)?;
+            let query = self.parse_query()?;
+            let closing_paren_token = self.expect_token(&Token::RParen)?;
+
+            let alias = TableAlias {
+                explicit: false,
+                name,
+                columns: vec![],
+            };
+            Cte {
+                alias,
+                query,
+                from: None,
+                materialized: None,
+                closing_paren_token: closing_paren_token.into(),
+            }
         } else {
             let columns = self.parse_table_alias_column_defs()?;
-            self.expect_keyword_is(Keyword::AS)?;
+            if self.dialect.supports_cte_without_as() {
+                let _ = self.parse_keyword(Keyword::AS);
+            } else {
+                self.expect_keyword_is(Keyword::AS)?;
+            }
             let mut is_materialized = None;
             if dialect_of!(self is PostgreSqlDialect) {
                 if self.parse_keyword(Keyword::MATERIALIZED) {
@@ -14125,6 +14149,23 @@ impl<'a> Parser<'a> {
         Ok(cte)
     }
 
+    fn is_query_start_token(&self, token: &TokenWithSpan) -> bool {
+        match &token.token {
+            Token::Word(w) => {
+                matches!(
+                    w.keyword,
+                    Keyword::SELECT
+                        | Keyword::WITH
+                        | Keyword::VALUES
+                        | Keyword::VALUE
+                        | Keyword::TABLE
+                ) || (w.keyword == Keyword::FROM && self.dialect.supports_from_first_select())
+            }
+            Token::LParen => true,
+            _ => false,
+        }
+    }
+
     /// Parse a "query body", which is an expression with roughly the
     /// following grammar:
     /// ```sql

tests/sqlparser_databricks.rs

@@ -644,3 +644,30 @@ fn parse_databricks_json_accessor() {
         "SELECT raw:store.bicycle.price::DOUBLE FROM store_data",
     );
 }
+
+#[test]
+fn parse_cte_without_as() {
+    databricks_and_generic().one_statement_parses_to(
+        "WITH cte (SELECT 1) SELECT * FROM cte",
+        "WITH cte AS (SELECT 1) SELECT * FROM cte",
+    );
+
+    databricks_and_generic().one_statement_parses_to(
+        "WITH a AS (SELECT 1), b (SELECT 2) SELECT * FROM a, b",
+        "WITH a AS (SELECT 1), b AS (SELECT 2) SELECT * FROM a, b",
+    );
+
+    databricks_and_generic().one_statement_parses_to(
+        "WITH cte (col1, col2) (SELECT 1, 2) SELECT * FROM cte",
+        "WITH cte (col1, col2) AS (SELECT 1, 2) SELECT * FROM cte",
+    );
+
+    databricks_and_generic().verified_query("WITH cte AS (SELECT 1) SELECT * FROM cte");
+
+    databricks_and_generic()
+        .verified_query("WITH cte (col1, col2) AS (SELECT 1, 2) SELECT * FROM cte");
+
+    assert!(all_dialects_where(|d| !d.supports_cte_without_as())
+        .parse_sql_statements("WITH cte (SELECT 1) SELECT * FROM cte")
+        .is_err());
+}