Fix: treat Snowflake's staged file paths as MacroVars (#1500)

* Fix: treat Snowflake's staged file paths as MacroVars

* Rephrase 'does not exist' to 'is undefined' in error msg

* Increase test coverage

Author: georgesittas

sqlmesh/core/dialect.py

@@ -10,6 +10,7 @@
 import pandas as pd
 from sqlglot import Dialect, Generator, Parser, Tokenizer, TokenType, exp
 from sqlglot.dialects.dialect import DialectType
+from sqlglot.dialects.snowflake import Snowflake
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.scope import traverse_scope
 from sqlglot.tokens import Token
@@ -89,6 +90,10 @@ def output_name(self) -> str:
         return self.this.name
 
 
+class StagedFilePath(exp.Table):
+    """Represents paths to "staged files" in Snowflake."""
+
+
 def _scan_var(self: Tokenizer) -> None:
     param = False
     bracket = False
@@ -315,6 +320,22 @@ def _parse_types(
     return parsed_type
 
 
+# Only needed for Snowflake: its "staged file" syntax (@<path>) clashes with our macro
+# var syntax. By converting the Var representation to a MacroVar, we should be able to
+# handle both use cases: if there's no value in the MacroEvaluator's context for that
+# MacroVar, it'll render into @<path>, so it won't break staged file path references.
+#
+# See: https://docs.snowflake.com/en/user-guide/querying-stage
+def _parse_table_parts(self: Parser, schema: bool = False) -> exp.Table:
+    table = self.__parse_table_parts(schema=schema)  # type: ignore
+    table_arg = table.this
+
+    if isinstance(table_arg, exp.Var) and table_arg.name.startswith("@"):
+        return StagedFilePath(this=MacroVar(this=table_arg.name[1:]))
+
+    return table
+
+
 def _create_parser(parser_type: t.Type[exp.Expression], table_keys: t.List[str]) -> t.Callable:
     def parse(self: Parser) -> t.Optional[exp.Expression]:
         from sqlmesh.core.model.kind import ModelKindName
@@ -634,18 +655,19 @@ def extend_sqlglot() -> None:
                 {
                     Audit: lambda self, e: _sqlmesh_ddl_sql(self, e, "Audit"),
                     DColonCast: lambda self, e: f"{self.sql(e, 'this')}::{self.sql(e, 'to')}",
+                    Jinja: lambda self, e: e.name,
+                    JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
+                    JinjaStatement: lambda self, e: f"{JINJA_STATEMENT_BEGIN};\n{e.name.strip()}\n{JINJA_END};",
                     MacroDef: lambda self, e: f"@DEF({self.sql(e.this)}, {self.sql(e.expression)})",
                     MacroFunc: _macro_func_sql,
                     MacroStrReplace: lambda self, e: f"@{self.sql(e.this)}",
                     MacroSQL: lambda self, e: f"@SQL({self.sql(e.this)})",
                     MacroVar: lambda self, e: f"@{e.name}",
                     Metric: lambda self, e: _sqlmesh_ddl_sql(self, e, "METRIC"),
                     Model: lambda self, e: _sqlmesh_ddl_sql(self, e, "MODEL"),
-                    Jinja: lambda self, e: e.name,
-                    JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
-                    JinjaStatement: lambda self, e: f"{JINJA_STATEMENT_BEGIN};\n{e.name.strip()}\n{JINJA_END};",
                     ModelKind: _model_kind_sql,
                     PythonCode: lambda self, e: self.expressions(e, sep="\n", indent=False),
+                    StagedFilePath: lambda self, e: self.table_sql(e),
                 }
             )
 
@@ -661,6 +683,7 @@ def extend_sqlglot() -> None:
     _override(Parser, _parse_having)
     _override(Parser, _parse_lambda)
     _override(Parser, _parse_types)
+    _override(Snowflake.Parser, _parse_table_parts)
 
 
 def select_from_values(

sqlmesh/core/macros.py

@@ -17,6 +17,7 @@
     MacroSQL,
     MacroStrReplace,
     MacroVar,
+    StagedFilePath,
 )
 from sqlmesh.utils import DECORATOR_RETURN_TYPE, UniqueKeyDict, registry_decorator
 from sqlmesh.utils.errors import MacroEvalError, SQLMeshError
@@ -137,6 +138,12 @@ def _transform_node(node: exp.Expression) -> exp.Expression:
 
             if isinstance(node, MacroVar):
                 changed = True
+                if node.name not in self.locals:
+                    if not isinstance(node.parent, StagedFilePath):
+                        raise SQLMeshError(f"Macro variable '{node.name}' is undefined.")
+
+                    return node
+
                 return exp.convert(_norm_env_value(self.locals[node.name]))
             if node.is_string:
                 text = node.this

tests/core/test_macros.py

@@ -5,7 +5,9 @@
 import pytest
 from sqlglot import exp, parse_one
 
+from sqlmesh.core.dialect import StagedFilePath
 from sqlmesh.core.macros import MacroEvaluator, macro
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.metaprogramming import Executable, ExecutableKind
 
 
@@ -76,6 +78,21 @@ def test_macro_var(macro_evaluator):
         macro_evaluator.locals = {"x": k}
         assert macro_evaluator.transform(expression).sql() == v
 
+    # Check Snowflake-specific StagedFilePath / MacroVar behavior
+    e = parse_one("select @x from @path, @y", dialect="snowflake")
+    macro_evaluator.locals = {"x": parse_one("a"), "y": parse_one("t2")}
+
+    assert e.find(StagedFilePath) is not None
+    assert macro_evaluator.transform(e).sql(dialect="snowflake") == "SELECT a FROM @path, t2"
+
+    # Referencing a var that doesn't exist in the evaluator's scope should raise
+    macro_evaluator.locals = {}
+    for dialect in ("", "snowflake"):
+        with pytest.raises(SQLMeshError) as ex:
+            macro_evaluator.transform(parse_one("SELECT @y", dialect=dialect))
+
+        assert "Macro variable 'y' is undefined." in str(ex.value)
+
 
 def test_macro_str_replace(macro_evaluator):
     expression = parse_one("""@'@val1, @val2'""")