Feat(tsql): if object_id is not null support closes #2044

tobymao · Aug 14, 2023 · 0746b6f · 0746b6f
1 parent 689956b
commit 0746b6f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -589,6 +589,17 @@ def _parse_create(self) -> exp.Create | exp.Command:
 
             return create
 
+        def _parse_if(self) -> t.Optional[exp.Expression]:
+            index = self._index
+
+            if self._match_text_seq("OBJECT_ID"):
+                self._parse_wrapped_csv(self._parse_string)
+                if self._match_text_seq("IS", "NOT", "NULL") and self._match(TokenType.DROP):
+                    return self._parse_drop(exists=True)
+                self._retreat(index)
+
+            return super()._parse_if()
+
     class Generator(generator.Generator):
         LOCKING_READS_SUPPORTED = True
         LIMIT_IS_TOP = True

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -1149,7 +1149,7 @@ def _parse_statement(self) -> t.Optional[exp.Expression]:
         expression = self._parse_set_operations(expression) if expression else self._parse_select()
         return self._parse_query_modifiers(expression)
 
-    def _parse_drop(self) -> exp.Drop | exp.Command:
+    def _parse_drop(self, exists: bool = False) -> exp.Drop | exp.Command:
         start = self._prev
         temporary = self._match(TokenType.TEMPORARY)
         materialized = self._match_text_seq("MATERIALIZED")
@@ -1161,7 +1161,7 @@ def _parse_drop(self) -> exp.Drop | exp.Command:
         return self.expression(
             exp.Drop,
             comments=start.comments,
-            exists=self._parse_exists(),
+            exists=exists or self._parse_exists(),
             this=self._parse_table(schema=True),
             kind=kind,
             temporary=temporary,

diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
@@ -10,6 +10,13 @@ def test_tsql(self):
         projection.assert_is(exp.Alias)
         projection.args["alias"].assert_is(exp.Identifier)
 
+        self.validate_all(
+            "IF OBJECT_ID('tempdb.dbo.#TempTableName', 'U') IS NOT NULL DROP TABLE #TempTableName",
+            write={
+                "tsql": "DROP TABLE IF EXISTS #TempTableName",
+                "spark": "DROP TABLE IF EXISTS TempTableName",
+            },
+        )
         self.validate_identity("UPDATE x SET y = 1 OUTPUT x.a, x.b INTO @y FROM y")
         self.validate_identity("UPDATE x SET y = 1 OUTPUT x.a, x.b FROM y")
         self.validate_identity("INSERT INTO x (y) OUTPUT x.a, x.b INTO l SELECT * FROM z")
@@ -898,6 +905,9 @@ def test_date_diff(self):
         )
 
     def test_iif(self):
+        self.validate_identity(
+            "SELECT IF(cond, 'True', 'False')", "SELECT IIF(cond, 'True', 'False')"
+        )
         self.validate_identity("SELECT IIF(cond, 'True', 'False')")
         self.validate_all(
             "SELECT IIF(cond, 'True', 'False');",