Support optional AS keyword in CTE definitions for Databricks

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()?;
 
@@ -14091,32 +14091,65 @@ impl<'a> Parser<'a> {
                 closing_paren_token: closing_paren_token.into(),
             }
         } else {
-            let columns = self.parse_table_alias_column_defs()?;
-            self.expect_keyword_is(Keyword::AS)?;
-            let mut is_materialized = None;
-            if dialect_of!(self is PostgreSqlDialect) {
-                if self.parse_keyword(Keyword::MATERIALIZED) {
-                    is_materialized = Some(CteAsMaterialized::Materialized);
-                } else if self.parse_keywords(&[Keyword::NOT, Keyword::MATERIALIZED]) {
-                    is_materialized = Some(CteAsMaterialized::NotMaterialized);
+            let as_optional = self.dialect.supports_cte_without_as();
+            let opt_query = if as_optional {
+                self.maybe_parse(|p| {
+                    p.expect_token(&Token::LParen)?;
+                    let query = p.parse_query()?;
+                    let closing_paren_token = p.expect_token(&Token::RParen)?;
+                    Ok((query, closing_paren_token))
+                })?
+            } else {
+                None
+            };
+            match opt_query {
+                Some((query, closing_paren_token)) => {
+                    let alias = TableAlias {
+                        explicit: false,
+                        name,
+                        columns: vec![],
+                    };
+                    Cte {
+                        alias,
+                        query,
+                        from: None,
+                        materialized: None,
+                        closing_paren_token: closing_paren_token.into(),
+                    }
                 }
-            }
-            self.expect_token(&Token::LParen)?;
+                None => {
+                    let columns = self.parse_table_alias_column_defs()?;
+                    if as_optional {
+                        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) {
+                            is_materialized = Some(CteAsMaterialized::Materialized);
+                        } else if self.parse_keywords(&[Keyword::NOT, Keyword::MATERIALIZED]) {
+                            is_materialized = Some(CteAsMaterialized::NotMaterialized);
+                        }
+                    }
+                    self.expect_token(&Token::LParen)?;
 
-            let query = self.parse_query()?;
-            let closing_paren_token = self.expect_token(&Token::RParen)?;
+                    let query = self.parse_query()?;
+                    let closing_paren_token = self.expect_token(&Token::RParen)?;
 
-            let alias = TableAlias {
-                explicit: false,
-                name,
-                columns,
-            };
-            Cte {
-                alias,
-                query,
-                from: None,
-                materialized: is_materialized,
-                closing_paren_token: closing_paren_token.into(),
+                    let alias = TableAlias {
+                        explicit: false,
+                        name,
+                        columns,
+                    };
+                    Cte {
+                        alias,
+                        query,
+                        from: None,
+                        materialized: is_materialized,
+                        closing_paren_token: closing_paren_token.into(),
+                    }
+                }
             }
         };
         if self.parse_keyword(Keyword::FROM) {

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());
+}