Fix: Resolution of this_model within macros of on_virtual_update (#3962)

Author: themisvaltinos

sqlmesh/core/model/definition.py

@@ -452,8 +452,6 @@ def render_on_virtual_update(
         engine_adapter: t.Optional[EngineAdapter] = None,
         **kwargs: t.Any,
     ) -> t.List[exp.Expression]:
-        if "this_model" not in kwargs:
-            kwargs["this_model"] = self.fully_qualified_table
         return self._render_statements(
             self.on_virtual_update,
             start=start,

sqlmesh/core/renderer.py

@@ -386,6 +386,7 @@ def render(
                 execution_time=execution_time,
                 snapshots=snapshots,
                 deployability_index=deployability_index,
+                table_mapping=table_mapping,
                 **kwargs,
             )
         except ParsetimeAdapterCallError:

tests/core/test_context.py

@@ -1813,3 +1813,75 @@ def test_plan_selector_expression_no_match(sushi_context: Context) -> None:
         match="Selector did not return any models. Please check your model selection and try again.",
     ):
         sushi_context.plan("prod", restate_models=["*missing*"])
+
+
+def test_plan_on_virtual_update_this_model_in_macro(tmp_path: pathlib.Path):
+    models_dir = pathlib.Path("models")
+    macros_dir = pathlib.Path("macros")
+    dialect = "duckdb"
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect=dialect),
+    )
+
+    model_file = """
+MODEL(
+  name db.test_view_macro_this_model,
+  kind full,
+);
+
+
+SELECT 1 AS cola;
+
+ON_VIRTUAL_UPDATE_BEGIN;
+CREATE OR REPLACE TABLE log_schema AS SELECT @resolve_template('@{schema_name}') as my_schema;
+@create_log_view(@this_model);
+ON_VIRTUAL_UPDATE_END;
+
+    """
+
+    create_temp_file(
+        tmp_path,
+        pathlib.Path(models_dir, "db", "test_view_macro_this_model.sql"),
+        model_file,
+    )
+
+    create_temp_file(
+        tmp_path,
+        pathlib.Path(macros_dir, "create_log_view.py"),
+        """
+from sqlmesh.core.macros import macro
+
+@macro()
+def create_log_view(evaluator, view_name):
+    return f"CREATE OR REPLACE TABLE log_view AS SELECT '{view_name}' as fqn_this_model,  '{evaluator.this_model}' as evaluator_this_model;"
+""",
+    )
+
+    context = Context(paths=tmp_path, config=config)
+    context.plan(environment="dev", auto_apply=True, no_prompts=True)
+
+    model = context.get_model("db.test_view_macro_this_model")
+    assert (
+        model.on_virtual_update[0].sql(dialect=dialect)
+        == "CREATE OR REPLACE TABLE log_schema AS SELECT @resolve_template('@{schema_name}') AS my_schema"
+    )
+    assert model.on_virtual_update[1].sql(dialect=dialect) == "@create_log_view(@this_model)"
+
+    snapshot = context.get_snapshot("db.test_view_macro_this_model")
+    assert snapshot and snapshot.version
+
+    log_view = context.fetchdf("select * from log_view").to_dict()
+    log_schema = context.fetchdf("select * from log_schema").to_dict()
+
+    # Validate that within macro for this_model we resolve to the environment-specific view
+    assert (
+        log_view["fqn_this_model"][0]
+        == '"db__dev"."test_view_macro_this_model" /* memory.db.test_view_macro_this_model */'
+    )
+
+    # Validate that from the macro evaluator this_model we get the environment-specific fqn
+    assert log_view["evaluator_this_model"][0] == '"db__dev"."test_view_macro_this_model"'
+
+    # Validate the schema is retrieved using resolve_template for the environment-specific schema
+    assert log_schema["my_schema"][0] == "db__dev"

tests/core/test_model.py

@@ -13,6 +13,7 @@
 from sqlglot.errors import ParseError
 from sqlglot.schema import MappingSchema
 from sqlmesh.cli.example_project import init_example_project, ProjectTemplate
+from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.model.kind import TimeColumn, ModelKindName
 
 from sqlmesh import CustomMaterialization, CustomKind
@@ -7843,6 +7844,7 @@ def test_model_on_virtual_update(make_snapshot: t.Callable):
     def resolve_parent_name(evaluator, name):
         return evaluator.resolve_table(name.name)
 
+    dialect = "postgres"
     virtual_update_statements = """
         CREATE OR REPLACE VIEW test_view FROM demo_db.table;
         GRANT SELECT ON VIEW @this_model TO ROLE owner_name;
@@ -7869,7 +7871,8 @@ def resolve_parent_name(evaluator, name):
 
         on_virtual_update_end;
 
-    """
+    """,
+        default_dialect=dialect,
     )
 
     parent_expressions = d.parse(
@@ -7886,49 +7889,57 @@ def resolve_parent_name(evaluator, name):
         JINJA_END;
         ON_VIRTUAL_UPDATE_END;
 
-    """
+    """,
+        default_dialect=dialect,
     )
 
-    model = load_sql_based_model(expressions)
-    parent = load_sql_based_model(parent_expressions)
+    model = load_sql_based_model(expressions, dialect=dialect)
+    parent = load_sql_based_model(parent_expressions, dialect=dialect)
 
     parent_snapshot = make_snapshot(parent)
     parent_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
-    version = parent_snapshot.version
 
     model_snapshot = make_snapshot(model)
     model_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
-    assert model.on_virtual_update == d.parse(virtual_update_statements)
+    assert model.on_virtual_update == d.parse(virtual_update_statements, default_dialect=dialect)
 
     assert parent.on_virtual_update == d.parse(
-        "JINJA_STATEMENT_BEGIN; GRANT SELECT ON VIEW {{this_model}} TO ROLE admin; JINJA_END;"
+        "JINJA_STATEMENT_BEGIN; GRANT SELECT ON VIEW {{this_model}} TO ROLE admin; JINJA_END;",
+        default_dialect=dialect,
     )
 
+    environment_naming_info = EnvironmentNamingInfo(name="dev")
     table_mapping = {model.fqn: "demo_db__dev.table", parent.fqn: "default__dev.parent"}
     snapshots = {
         parent_snapshot.name: parent_snapshot,
         model_snapshot.name: model_snapshot,
     }
 
     rendered_on_virtual_update = model.render_on_virtual_update(
-        snapshots=snapshots, table_mapping=table_mapping
+        snapshots=snapshots,
+        table_mapping=table_mapping,
+        this_model=model_snapshot.qualified_view_name.table_for_environment(
+            environment_naming_info, dialect=dialect
+        ),
     )
 
     assert len(rendered_on_virtual_update) == 6
     assert (
         rendered_on_virtual_update[0].sql()
         == 'CREATE OR REPLACE VIEW "test_view" AS SELECT * FROM "demo_db__dev"."table" AS "table" /* demo_db.table */'
     )
+
     assert (
         rendered_on_virtual_update[1].sql()
-        == 'GRANT SELECT ON VIEW "demo_db__dev"."table" /* demo_db.table */ TO ROLE "owner_name"'
+        == 'GRANT SELECT ON VIEW "demo_db__dev"."table" TO ROLE "owner_name"'
     )
     assert (
         rendered_on_virtual_update[3].sql()
         == "GRANT REFERENCES, SELECT ON FUTURE VIEWS IN DATABASE demo_db TO ROLE owner_name"
     )
-    assert rendered_on_virtual_update[4].sql() == f'"sqlmesh__default"."parent__{version}"'
+
+    assert rendered_on_virtual_update[4].sql() == '"default__dev"."parent"'
 
     # When replace=false the table should remain as is
     assert (
@@ -7937,12 +7948,16 @@ def resolve_parent_name(evaluator, name):
     )
 
     rendered_parent_on_virtual_update = parent.render_on_virtual_update(
-        snapshots=snapshots, table_mapping=table_mapping
+        snapshots=snapshots,
+        table_mapping=table_mapping,
+        this_model=parent_snapshot.qualified_view_name.table_for_environment(
+            environment_naming_info, dialect=dialect
+        ),
     )
     assert len(rendered_parent_on_virtual_update) == 1
     assert (
         rendered_parent_on_virtual_update[0].sql()
-        == 'GRANT SELECT ON VIEW "default__dev"."parent" /* parent */ TO ROLE "admin"'
+        == 'GRANT SELECT ON VIEW "default__dev"."parent" TO ROLE "admin"'
     )
 
 

tests/core/test_snapshot_evaluator.py

@@ -2919,6 +2919,10 @@ def model_with_statements(context, **kwargs):
 def test_on_virtual_update_statements(mocker: MockerFixture, adapter_mock, make_snapshot):
     evaluator = SnapshotEvaluator(adapter_mock)
 
+    @macro()
+    def create_log_table(evaluator, view_name):
+        return f"CREATE OR REPLACE TABLE log_table AS SELECT '{view_name}' as fqn_this_model, '{evaluator.this_model}' as eval_this_model"
+
     model = load_sql_based_model(
         d.parse(
             """
@@ -2937,6 +2941,7 @@ def test_on_virtual_update_statements(mocker: MockerFixture, adapter_mock, make_
             GRANT SELECT ON VIEW test_schema.test_model TO ROLE admin;
             JINJA_END;
             GRANT REFERENCES, SELECT ON FUTURE VIEWS IN DATABASE demo_db TO ROLE owner_name;
+            @create_log_table(@this_model);
             ON_VIRTUAL_UPDATE_END;
 
             """
@@ -2988,6 +2993,12 @@ def test_on_virtual_update_statements(mocker: MockerFixture, adapter_mock, make_
         == "GRANT REFERENCES, SELECT ON FUTURE VIEWS IN DATABASE demo_db TO ROLE owner_name"
     )
 
+    # Validation that within the macro the environment specific view is used
+    assert (
+        on_virtual_update_calls[2].sql(dialect="postgres")
+        == 'CREATE OR REPLACE TABLE "log_table" AS SELECT \'"test_schema__test_env"."test_model" /* test_schema.test_model */\' AS "fqn_this_model", \'"test_schema__test_env"."test_model"\' AS "eval_this_model"'
+    )
+
 
 def test_on_virtual_update_python_model_macro(mocker: MockerFixture, adapter_mock, make_snapshot):
     evaluator = SnapshotEvaluator(adapter_mock)