From 0746b6f96d9b8fad0d8fbea3e23170e8d56eb3ee Mon Sep 17 00:00:00 2001
From: tobymao <toby.mao@gmail.com>
Date: Sun, 13 Aug 2023 21:52:08 -0700
Subject: [PATCH] Feat(tsql): if object_id is not null support closes #2044

---
 sqlglot/dialects/tsql.py    | 11 +++++++++++
 sqlglot/parser.py           |  4 ++--
 tests/dialects/test_tsql.py | 10 ++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
index 0e0d9ee9bf..a5089a1678 100644
--- 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
index 8623024694..10fecc9a27 100644
--- 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
index f43b41bd14..9decd35440 100644
--- 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');",