Fix!: enable loading of Snowflake models that use staged file paths (#2973)

Author: georgesittas

sqlmesh/core/dialect.py

@@ -101,9 +101,11 @@ def output_name(self) -> str:
         return self.this.name
 
 
-class StagedFilePath(exp.Table):
+class StagedFilePath(exp.Expression):
     """Represents paths to "staged files" in Snowflake."""
 
+    arg_types = exp.Table.arg_types.copy()
+
 
 def _parse_statement(self: Parser) -> t.Optional[exp.Expression]:
     if self._curr is None:
@@ -409,7 +411,7 @@ def _parse_types(
 # See: https://docs.snowflake.com/en/user-guide/querying-stage
 def _parse_table_parts(
     self: Parser, schema: bool = False, is_db_reference: bool = False
-) -> exp.Table:
+) -> exp.Table | StagedFilePath:
     index = self._index
     table = self.__parse_table_parts(schema=schema, is_db_reference=is_db_reference)  # type: ignore
 

tests/core/test_model.py

@@ -5371,3 +5371,16 @@ def test_trailing_comments():
     model = load_sql_based_model(expressions, path=Path("./examples/sushi/models/test_model.sql"))
     assert not model.render_pre_statements()
     assert not model.render_post_statements()
+
+
+def test_staged_file_path():
+    expressions = d.parse(
+        """
+        MODEL (name test, dialect snowflake);
+
+        SELECT * FROM @a.b/c/d.csv(FILE_FORMAT => 'b.ff')
+        """
+    )
+    model = load_sql_based_model(expressions)
+    query = model.render_query()
+    assert query.sql(dialect="snowflake") == "SELECT * FROM @a.b/c/d.csv (FILE_FORMAT => 'b.ff')"