Repository: tobymao/sqlglot

Issue:
Parsing tsql stored proc with declare variable
**Error description**
I get the following sqlglot.errors.ParseError: Invalid expression / Unexpected token Line 4, Col: 28. which is the datatype varchar. 

I have look through the library and testet with and without the variable and it might be related to the way that the chunks is made when parsing. 

**Fully reproducible code snippet**
```sql
from sqlglot import parse_one, exp

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;
"""

for table in parse_one(sql, dialect='tsql').find_all(exp.Table):
    print(table.db + '.' + table.name)
```

**Official Documentation**
https://learn.microsoft.com/en-us/azure/synapse-analytics/sql-data-warehouse/sql-data-warehouse-develop-ctas


Discussion:

Patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index a32762ca..186226ae 100644
--- 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
index e728245d..2443c4f6 100644
--- 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
index 5b2e091f..4dfe2999 100644
--- 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 @@ class Parser(metaclass=_Parser):
 
         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 @@ class Parser(metaclass=_Parser):
             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:


Test patch:
diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
index c0f5c988..5266bd4b 100644
--- a/tests/dialects/test_tsql.py
+++ b/tests/dialects/test_tsql.py
@@ -558,6 +558,30 @@ WHERE
         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)",