Fix(tsql): improve UDF parsing (#1973)

tobymao · Jul 28, 2023 · 2874ae5 · 2874ae5
1 parent 9d67283
commit 2874ae5
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 6 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1931,6 +1931,10 @@ class EngineProperty(Property):
     arg_types = {"this": True}
 
 
+class HeapProperty(Property):
+    arg_types = {}
+
+
 class ToTableProperty(Property):
     arg_types = {"this": True}
 

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -68,6 +68,7 @@ class Generator:
         exp.EncodeColumnConstraint: lambda self, e: f"ENCODE {self.sql(e, 'this')}",
         exp.ExecuteAsProperty: lambda self, e: self.naked_property(e),
         exp.ExternalProperty: lambda self, e: "EXTERNAL",
+        exp.HeapProperty: lambda self, e: "HEAP",
         exp.InlineLengthColumnConstraint: lambda self, e: f"INLINE LENGTH {self.sql(e, 'this')}",
         exp.LanguageProperty: lambda self, e: self.naked_property(e),
         exp.LocationProperty: lambda self, e: self.naked_property(e),
@@ -224,6 +225,7 @@ class Generator:
         exp.FallbackProperty: exp.Properties.Location.POST_NAME,
         exp.FileFormatProperty: exp.Properties.Location.POST_WITH,
         exp.FreespaceProperty: exp.Properties.Location.POST_NAME,
+        exp.HeapProperty: exp.Properties.Location.POST_WITH,
         exp.IsolatedLoadingProperty: exp.Properties.Location.POST_NAME,
         exp.JournalProperty: exp.Properties.Location.POST_NAME,
         exp.LanguageProperty: exp.Properties.Location.POST_SCHEMA,

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -605,6 +605,7 @@ class Parser(metaclass=_Parser):
         "FALLBACK": lambda self, **kwargs: self._parse_fallback(**kwargs),
         "FORMAT": lambda self: self._parse_property_assignment(exp.FileFormatProperty),
         "FREESPACE": lambda self: self._parse_freespace(),
+        "HEAP": lambda self: self.expression(exp.HeapProperty),
         "IMMUTABLE": lambda self: self.expression(
             exp.StabilityProperty, this=exp.Literal.string("IMMUTABLE")
         ),
@@ -1189,7 +1190,7 @@ def _parse_create(self) -> exp.Create | exp.Command:
 
         exists = self._parse_exists(not_=True)
         this = None
-        expression = None
+        expression: t.Optional[exp.Expression] = None
         indexes = None
         no_schema_binding = None
         begin = None
@@ -1209,12 +1210,16 @@ def extend_props(temp_props: t.Optional[exp.Properties]) -> None:
             extend_props(self._parse_properties())
 
             self._match(TokenType.ALIAS)
-            begin = self._match(TokenType.BEGIN)
-            return_ = self._match_text_seq("RETURN")
-            expression = self._parse_statement()
 
-            if return_:
-                expression = self.expression(exp.Return, this=expression)
+            if self._match(TokenType.COMMAND):
+                expression = self._parse_as_command(self._prev)
+            else:
+                begin = self._match(TokenType.BEGIN)
+                return_ = self._match_text_seq("RETURN")
+                expression = self._parse_statement()
+
+                if return_:
+                    expression = self.expression(exp.Return, this=expression)
         elif create_token.token_type == TokenType.INDEX:
             this = self._parse_index(index=self._parse_id_var())
         elif create_token.token_type in self.DB_CREATABLES:

diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
@@ -558,6 +558,30 @@ def test_fullproc(self):
         for expr, expected_sql in zip(parse(sql, read="tsql"), expected_sqls):
             self.assertEqual(expr.sql(dialect="tsql"), expected_sql)
 
+        sql = """
+            CREATE PROC [dbo].[transform_proc] AS
+
+            DECLARE @CurrentDate VARCHAR(20);
+            SET @CurrentDate = CONVERT(VARCHAR(20), GETDATE(), 120);
+
+            CREATE TABLE [target_schema].[target_table]
+            WITH (DISTRIBUTION = REPLICATE, HEAP)
+            AS
+
+            SELECT
+                @CurrentDate AS DWCreatedDate
+            FROM source_schema.sourcetable;
+        """
+
+        expected_sqls = [
+            'CREATE PROC "dbo"."transform_proc" AS DECLARE @CurrentDate VARCHAR(20)',
+            "SET @CurrentDate = CAST(FORMAT(GETDATE(), 'yyyy-MM-dd HH:mm:ss') AS VARCHAR(20))",
+            'CREATE TABLE "target_schema"."target_table" WITH (DISTRIBUTION=REPLICATE, HEAP) AS SELECT @CurrentDate AS DWCreatedDate FROM source_schema.sourcetable',
+        ]
+
+        for expr, expected_sql in zip(parse(sql, read="tsql"), expected_sqls):
+            self.assertEqual(expr.sql(dialect="tsql"), expected_sql)
+
     def test_charindex(self):
         self.validate_all(
             "CHARINDEX(x, y, 9)",