Fix(parser)!: require AS token in CTEs for all dialects except spark, databricks (#4546)

Author: georgesittas

sqlglot/dialects/clickhouse.py

@@ -633,7 +633,7 @@ def _parse_position(self, haystack_first: bool = False) -> exp.StrPosition:
             return super()._parse_position(haystack_first=True)
 
         # https://clickhouse.com/docs/en/sql-reference/statements/select/with/
-        def _parse_cte(self) -> exp.CTE:
+        def _parse_cte(self) -> t.Optional[exp.CTE]:
             # WITH <identifier> AS <subquery expression>
             cte: t.Optional[exp.CTE] = self._try_parse(super()._parse_cte)
 

sqlglot/dialects/spark.py

@@ -102,6 +102,8 @@ class Tokenizer(Spark2.Tokenizer):
         ]
 
     class Parser(Spark2.Parser):
+        OPTIONAL_ALIAS_TOKEN_CTE = True
+
         FUNCTIONS = {
             **Spark2.Parser.FUNCTIONS,
             "ANY_VALUE": _build_with_ignore_nulls(exp.AnyValue),

sqlglot/parser.py

@@ -1379,6 +1379,9 @@ class Parser(metaclass=_Parser):
     # Whether the `name AS expr` schema/column constraint requires parentheses around `expr`
     WRAPPED_TRANSFORM_COLUMN_CONSTRAINT = True
 
+    # Whether the 'AS' keyword is optional in the CTE definition syntax
+    OPTIONAL_ALIAS_TOKEN_CTE = False
+
     __slots__ = (
         "error_level",
         "error_message_context",
@@ -3131,12 +3134,17 @@ def _parse_with(self, skip_with_token: bool = False) -> t.Optional[exp.With]:
             exp.With, comments=comments, expressions=expressions, recursive=recursive
         )
 
-    def _parse_cte(self) -> exp.CTE:
+    def _parse_cte(self) -> t.Optional[exp.CTE]:
+        index = self._index
+
         alias = self._parse_table_alias(self.ID_VAR_TOKENS)
         if not alias or not alias.this:
             self.raise_error("Expected CTE to have alias")
 
-        self._match(TokenType.ALIAS)
+        if not self._match(TokenType.ALIAS) and not self.OPTIONAL_ALIAS_TOKEN_CTE:
+            self._retreat(index)
+            return None
+
         comments = self._prev_comments
 
         if self._match_text_seq("NOT", "MATERIALIZED"):

tests/dialects/test_clickhouse.py

@@ -29,6 +29,7 @@ def test_clickhouse(self):
         self.assertEqual(expr.sql(dialect="clickhouse"), "COUNT(x)")
         self.assertIsNone(expr._meta)
 
+        self.validate_identity("WITH arrayJoin([(1, [2, 3])]) AS arr SELECT arr")
         self.validate_identity("CAST(1 AS Bool)")
         self.validate_identity("SELECT toString(CHAR(104.1, 101, 108.9, 108.9, 111, 32))")
         self.validate_identity("@macro").assert_is(exp.Parameter).this.assert_is(exp.Var)