Repository: TobikoData/sqlmesh
Pull request #3254 (2024-10-15 14:11:19)

Fix: refactor table part parsing for Snowflake
I found out that my recent [fix](https://github.com/TobikoData/sqlmesh/commit/d70092e69e5e1f8a6fdec41adc3e1d052176e6ee) wasn't complete. It works when someone's loading their SQLMesh project for the first time, but when the models are loaded from cache we fail to parse them because the global macro registry has already been reset and only contains the standard macros.

Since relying on the registry proved to be brittle, I decided to refactor the existing parsing logic altogether. For context, these are some past commits that aimed to address staged file path parsing edge cases:

- https://github.com/TobikoData/sqlmesh/commit/30071ccb009c2eebad585065b095e976ca1bf5cd
- https://github.com/TobikoData/sqlmesh/pull/2412/files
- https://github.com/TobikoData/sqlmesh/pull/2714
- https://github.com/tobikoData/sqlmesh/commit/932167933877a2e7f50f12e180c68d4ff0e4506a


Patch:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 9f44199a..db90e1ba 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -464,11 +464,23 @@ def _parse_table_parts(
     name = table_arg.name
 
     if isinstance(table_arg, exp.Var) and name.startswith(SQLMESH_MACRO_PREFIX):
-        # Macro functions do not clash with the staged file syntax, so we can safely parse them
-        from sqlmesh.core.macros import macro
-
-        macros = macro.get_registry()
-        if self._prev.token_type == TokenType.STRING or "{" in name or name[1:].lower() in macros:
+        # In these cases, we don't want to produce a `StagedFilePath` node:
+        #
+        # - @'...' needs to parsed as a string template
+        # - @{foo}.bar needs to be parsed as a table with a macro var part
+        # - @name(arg1 [, arg2 ...]) needs to be parsed as a macro function call
+        #
+        # These cases can unambiguously be parsed using the base `_parse_table_parts`, as there
+        # is no overlap with staged files https://docs.snowflake.com/en/user-guide/querying-stage
+        if (
+            self._prev.token_type == TokenType.STRING
+            or "{" in name
+            or (
+                self._curr
+                and self._prev.token_type == TokenType.L_PAREN
+                and self._curr.text.upper() not in ("FILE_FORMAT", "PATTERN")
+            )
+        ):
             self._retreat(index)
             return Parser._parse_table_parts(self, schema=schema, is_db_reference=is_db_reference)
 
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
index beb8d035..1b2f2632 100644
--- a/sqlmesh/core/loader.py
+++ b/sqlmesh/core/loader.py
@@ -93,18 +93,11 @@ class Loader(abc.ABC):
         self._config_mtimes = {path: max(mtimes) for path, mtimes in config_mtimes.items()}
 
         macros, jinja_macros = self._load_scripts()
-
-        # All macros need to be known at parse time
-        standard_macros = macro.get_registry()
-        macro.set_registry(macros)
-
         audits = self._load_audits(macros=macros, jinja_macros=jinja_macros)
         models = self._load_models(
             macros, jinja_macros, context.gateway or context.config.default_gateway, audits or None
         )
 
-        macro.set_registry(standard_macros)
-
         for model in models.values():
             self._add_model_to_dag(model)
 


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index ed1d254c..e2f144ce 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -5934,7 +5934,7 @@ from sqlmesh import macro
 
 @macro()
 def custom_macro(evaluator, arg1, arg2):
-    return f"{arg1}{arg2}"
+    return "SELECT 1 AS c"
     """)
 
     new_snowflake_model_file = tmp_path / "models/new_model.sql"
@@ -5948,7 +5948,7 @@ MODEL (
 @DEF(foo, foo);
 @DEF(bar, bar);
 
-SELECT * FROM @custom_macro(@foo, @bar)
+SELECT * FROM (@custom_macro(@foo, @bar)) AS q
     """)
 
     config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
@@ -5956,4 +5956,9 @@ SELECT * FROM @custom_macro(@foo, @bar)
 
     query = context.get_model("sqlmesh_example.test").render_query()
 
-    assert t.cast(exp.Query, query).sql("snowflake") == 'SELECT * FROM "FOOBAR" AS "FOOBAR"'
+    assert (
+        t.cast(exp.Query, query).sql("snowflake")
+        == 'SELECT "Q"."C" AS "C" FROM (SELECT 1 AS "C") AS "Q"'
+    )
+
+    context.plan(no_prompts=True, auto_apply=True)