Fix: conditional statement parsing edge case (TobikoData#2967)

etonlels · Jul 30, 2024 · 6d1c54b · 6d1c54b
1 parent 1dc3b86
commit 6d1c54b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
@@ -454,7 +454,13 @@ def _parse_if(self: Parser) -> t.Optional[exp.Expression]:
         else:
             self.raise_error("Expecting )")
 
-        return exp.Anonymous(this="IF", expressions=[cond, self._parse_statement()])
+        index = self._index
+        stmt = self._parse_statement()
+        if self._curr:
+            self._retreat(index)
+            stmt = self._parse_as_command(self._tokens[index])
+
+        return exp.Anonymous(this="IF", expressions=[cond, stmt])
 
 
 def _create_parser(parser_type: t.Type[exp.Expression], table_keys: t.List[str]) -> t.Callable:

diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
@@ -606,3 +606,6 @@ def test_conditional_statement():
         q.sql("snowflake")
         == "@IF(TRUE, COPY INTO 's3://example/data.csv' FROM EXTRA.EXAMPLE.TABLE STORAGE_INTEGRATION = S3_INTEGRATION FILE_FORMAT = (TYPE=CSV COMPRESSION=NONE NULL_IF=('') FIELD_OPTIONALLY_ENCLOSED_BY='\"') HEADER = TRUE OVERWRITE = TRUE SINGLE = TRUE /* this is a comment */)"
     )
+
+    q = parse_one("@IF(cond, VACUUM ANALYZE);", read="postgres")
+    assert q.sql(dialect="postgres") == "@IF(cond, VACUUM ANALYZE)"