allow tsql global temp

tobymao · Jan 6, 2023 · 52db84c · 52db84c
1 parent e19d7f7
commit 52db84c
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 11 deletions.
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -271,6 +271,9 @@ class Parser(parser.Parser):
             DataType.Type.NCHAR,
         }
 
+        # https://learn.microsoft.com/en-us/azure/synapse-analytics/sql-data-warehouse/sql-data-warehouse-tables-temporary#create-a-temporary-table
+        TABLE_PREFIX_TOKENS = {TokenType.HASH}
+
         def _parse_convert(self, strict):
             to = self._parse_types()
             self._match(TokenType.COMMA)

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -591,6 +591,9 @@ class Parser(metaclass=_Parser):
 
     TRANSACTION_KIND = {"DEFERRED", "IMMEDIATE", "EXCLUSIVE"}
 
+    # allows tables to have special tokens as prefixes
+    TABLE_PREFIX_TOKENS: t.Set[TokenType] = set()
+
     STRICT_CAST = True
 
     __slots__ = (
@@ -1427,14 +1430,9 @@ def _parse_table(self, schema=False, alias_tokens=None):
 
         catalog = None
         db = None
-        table = not schema and self._parse_function()
-        if not table:
-            # https://learn.microsoft.com/en-us/azure/synapse-analytics/sql-data-warehouse/sql-data-warehouse-tables-temporary#create-a-temporary-table
-            # if this needs to be transpiled, we should create a temp expression node instead
-            temp = self._match(TokenType.HASH)
-            table = self._parse_id_var(any_token=False)
-            if temp:
-                table.args["this"] = f"#{table.name}"
+        table = (not schema and self._parse_function()) or self._parse_id_var(
+            any_token=False, prefix_tokens=self.TABLE_PREFIX_TOKENS
+        )
 
         while self._match(TokenType.DOT):
             if catalog:
@@ -2601,14 +2599,20 @@ def _parse_alias(self, this, explicit=False):
 
         return this
 
-    def _parse_id_var(self, any_token=True, tokens=None):
+    def _parse_id_var(self, any_token=True, tokens=None, prefix_tokens=None):
         identifier = self._parse_identifier()
 
         if identifier:
             return identifier
 
+        prefix = ""
+
+        if prefix_tokens:
+            while self._match_set(prefix_tokens):
+                prefix += self._prev.text
+
         if (any_token and self._advance_any()) or self._match_set(tokens or self.ID_VAR_TOKENS):
-            return exp.Identifier(this=self._prev.text, quoted=False)
+            return exp.Identifier(this=prefix + self._prev.text, quoted=False)
         return None
 
     def _parse_string(self):

diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
@@ -6,6 +6,8 @@ class TestTSQL(Validator):
 
     def test_tsql(self):
         self.validate_identity('SELECT "x"."y" FROM foo')
+        self.validate_identity("SELECT * FROM #foo")
+        self.validate_identity("SELECT * FROM ##foo")
         self.validate_identity(
             "SELECT DISTINCT DepartmentName, PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY BaseRate) OVER (PARTITION BY DepartmentName) AS MedianCont FROM dbo.DimEmployee"
         )

diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
@@ -64,7 +64,6 @@ x IN ('a', 'a''a')
 x IN ((1))
 x BETWEEN -1 AND 1
 x BETWEEN 'a' || b AND 'c' || d
-SELECT * FROM #x
 NOT x IS NULL
 x IS TRUE
 x IS FALSE