Fix(clickhouse)!: preserve EXTRACT(date_part FROM datetime) calls (#3729)

georgesittas · web-flow · commit 8335ba10e60c · 2024-07-02T20:09:54.000+03:00
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -174,7 +174,6 @@ class Parser(parser.Parser):
             "DATE_FORMAT": _build_date_format,
             "DATE_SUB": build_date_delta(exp.DateSub, default_unit=None),
             "DATESUB": build_date_delta(exp.DateSub, default_unit=None),
-            "EXTRACT": exp.RegexpExtract.from_arg_list,
             "FORMATDATETIME": _build_date_format,
             "JSONEXTRACTSTRING": build_json_extract_path(
                 exp.JSONExtractScalar, zero_based_indexing=False
@@ -347,7 +346,6 @@ class Parser(parser.Parser):
             "QUANTILE": lambda self: self._parse_quantile(),
         }
 
-        FUNCTION_PARSERS.pop("EXTRACT")
         FUNCTION_PARSERS.pop("MATCH")
 
         NO_PAREN_FUNCTION_PARSERS = parser.Parser.NO_PAREN_FUNCTION_PARSERS.copy()
@@ -410,6 +408,23 @@ class Parser(parser.Parser):
             "INDEX",
         }
 
+        def _parse_extract(self) -> exp.Extract | exp.Anonymous:
+            index = self._index
+            this = self._parse_bitwise()
+            if self._match(TokenType.FROM):
+                self._retreat(index)
+                return super()._parse_extract()
+
+            # We return Anonymous here because extract and regexpExtract have different semantics,
+            # so parsing extract(foo, bar) into RegexpExtract can potentially break queries. E.g.,
+            # `extract('foobar', 'b')` works, but CH crashes for `regexpExtract('foobar', 'b')`.
+            #
+            # TODO: can we somehow convert the former into an equivalent `regexpExtract` call?
+            self._match(TokenType.COMMA)
+            return self.expression(
+                exp.Anonymous, this="extract", expressions=[this, self._parse_bitwise()]
+            )
+
         def _parse_assignment(self) -> t.Optional[exp.Expression]:
             this = super()._parse_assignment()
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -1477,9 +1477,9 @@ def _parse_command(self) -> exp.Command:
 
     def _try_parse(self, parse_method: t.Callable[[], T], retreat: bool = False) -> t.Optional[T]:
         """
-        Attemps to backtrack if a parse function that contains a try/catch internally raises an error. This behavior can
-        be different depending on the uset-set ErrorLevel, so _try_parse aims to solve this by setting & resetting
-        the parser state accordingly
+        Attemps to backtrack if a parse function that contains a try/catch internally raises an error.
+        This behavior can be different depending on the uset-set ErrorLevel, so _try_parse aims to
+        solve this by setting & resetting the parser state accordingly
         """
         index = self._index
         error_level = self.error_level
@@ -5345,7 +5345,7 @@ def _parse_next_value_for(self) -> t.Optional[exp.Expression]:
             order=self._match(TokenType.OVER) and self._parse_wrapped(self._parse_order),
         )
 
-    def _parse_extract(self) -> exp.Extract:
+    def _parse_extract(self) -> exp.Extract | exp.Anonymous:
         this = self._parse_function() or self._parse_var_or_string(upper=True)
 
         if self._match(TokenType.FROM):
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -25,6 +25,7 @@ def test_clickhouse(self):
         self.assertEqual(expr.sql(dialect="clickhouse"), "COUNT(x)")
         self.assertIsNone(expr._meta)
 
+        self.validate_identity("SELECT EXTRACT(YEAR FROM toDateTime('2023-02-01'))")
         self.validate_identity("extract(haystack, pattern)")
         self.validate_identity("SELECT * FROM x LIMIT 1 UNION ALL SELECT * FROM y")
         self.validate_identity("SELECT CAST(x AS Tuple(String, Array(Nullable(Float64))))")
