Fix(clickhouse): properly parse CREATE FUNCTION DDLs (#4282)

tobymao · Oct 22, 2024 · ac66d2f · ac66d2f
1 parent 0a8a71b
commit ac66d2f
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -503,6 +503,10 @@ class Parser(parser.Parser):
             TokenType.L_BRACE: lambda self: self._parse_query_parameter(),
         }
 
+        # https://clickhouse.com/docs/en/sql-reference/statements/create/function
+        def _parse_user_defined_function_expression(self) -> t.Optional[exp.Expression]:
+            return self._parse_lambda()
+
         def _parse_types(
             self, check_func: bool = False, schema: bool = False, allow_identifiers: bool = True
         ) -> t.Optional[exp.Expression]:

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -1844,7 +1844,7 @@ def extend_props(temp_props: t.Optional[exp.Properties]) -> None:
                         expression = self._parse_string()
                         extend_props(self._parse_properties())
                     else:
-                        expression = self._parse_statement()
+                        expression = self._parse_user_defined_function_expression()
 
                     end = self._match_text_seq("END")
 
@@ -5341,6 +5341,9 @@ def _kv_to_prop_eq(self, expressions: t.List[exp.Expression]) -> t.List[exp.Expr
 
         return transformed
 
+    def _parse_user_defined_function_expression(self) -> t.Optional[exp.Expression]:
+        return self._parse_statement()
+
     def _parse_function_parameter(self) -> t.Optional[exp.Expression]:
         return self._parse_column_def(self._parse_id_var())
 

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -684,6 +684,7 @@ def test_ddl(self):
             "CREATE TABLE foo ENGINE=Memory AS (SELECT * FROM db.other_table) COMMENT 'foo'",
         )
 
+        self.validate_identity("CREATE FUNCTION linear_equation AS (x, k, b) -> k * x + b")
         self.validate_identity("CREATE MATERIALIZED VIEW a.b TO a.c (c Int32) AS SELECT * FROM a.d")
         self.validate_identity("""CREATE TABLE ip_data (ip4 IPv4, ip6 IPv6) ENGINE=TinyLog()""")
         self.validate_identity("""CREATE TABLE dates (dt1 Date32) ENGINE=TinyLog()""")
@@ -707,6 +708,10 @@ def test_ddl(self):
         self.validate_identity(
             "CREATE TABLE foo (x UInt32) TTL time_column + INTERVAL '1' MONTH DELETE WHERE column = 'value'"
         )
+        self.validate_identity(
+            "CREATE FUNCTION parity_str AS (n) -> IF(n % 2, 'odd', 'even')",
+            "CREATE FUNCTION parity_str AS n -> CASE WHEN n % 2 THEN 'odd' ELSE 'even' END",
+        )
         self.validate_identity(
             "CREATE TABLE a ENGINE=Memory AS SELECT 1 AS c COMMENT 'foo'",
             "CREATE TABLE a ENGINE=Memory AS (SELECT 1 AS c) COMMENT 'foo'",