fix(macros): resolve_template fails in audits when this_model is a subquery

sqlmesh/core/macros.py

@@ -1379,15 +1379,17 @@ def resolve_template(
     """
     Generates either a String literal or an exp.Table representing a physical table location, based on rendering the provided template String literal.
 
-    Note: It relies on the @this_model variable being available in the evaluation context (@this_model resolves to an exp.Table object
-    representing the current physical table).
+    Note: It relies on the @this_model variable being available in the evaluation context. @this_model usually resolves to an
+    exp.Table object representing the current physical table, but in an audit on a model with a time column it resolves to a
+    subquery that selects from that table and filters it down to the audited time range. In that case the placeholders below
+    are resolved against the physical table the subquery selects from.
     Therefore, the @resolve_template macro must be used at creation or evaluation time and not at load time.
 
     Args:
         template: Template string literal. Can contain the following placeholders:
-            @{catalog_name} -> replaced with the catalog of the exp.Table returned from @this_model
-            @{schema_name} -> replaced with the schema of the exp.Table returned from @this_model
-            @{table_name} -> replaced with the name of the exp.Table returned from @this_model
+            @{catalog_name} -> replaced with the catalog of the physical table @this_model refers to
+            @{schema_name} -> replaced with the schema of the physical table @this_model refers to
+            @{table_name} -> replaced with the name of the physical table @this_model refers to
         mode: What to return.
             'literal' -> return an exp.Literal string
             'table' -> return an exp.Table
@@ -1400,9 +1402,26 @@ def resolve_template(
         >>> evaluator.locals.update({"this_model": exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505")})
         >>> evaluator.transform(parse_one(sql)).sql()
         "'s3://data-bucket/prod/test_catalog/sqlmesh__test/test__test_model__2517971505'"
+
+        The same template resolves to the same location when @this_model is the time-filtered
+        subquery that audits on models with a time column receive:
+
+        >>> table = exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505")
+        >>> subquery = exp.select("*").from_(table).where(exp.column("ds").eq("2020-01-01")).subquery()
+        >>> evaluator.locals.update({"this_model": subquery})
+        >>> evaluator.transform(parse_one(sql)).sql()
+        "'s3://data-bucket/prod/test_catalog/sqlmesh__test/test__test_model__2517971505'"
     """
     if "this_model" in evaluator.locals:
-        this_model = exp.to_table(evaluator.locals["this_model"], dialect=evaluator.dialect)
+        this_model_expr = evaluator.locals["this_model"]
+        if isinstance(this_model_expr, exp.Subquery):
+            # Audits on models with a time column render @this_model as a subquery that filters the
+            # physical table on the audited time range, so resolve against the table it selects from
+            from_ = this_model_expr.unnest().args.get("from_")
+            if from_ is not None and isinstance(from_.this, exp.Table):
+                this_model_expr = from_.this
+
+        this_model = exp.to_table(this_model_expr, dialect=evaluator.dialect)
         template_str: str = template.this
         result = (
             template_str.replace("@{catalog_name}", this_model.catalog)

tests/core/test_audit.py

@@ -413,6 +413,30 @@ def test_macro(model: Model):
     assert model.render_audit_query(audit_jinja).sql() == expected_query
 
 
+def test_resolve_template(model_default_catalog: Model):
+    # @this_model is rendered as a subquery for models with a time column, but @resolve_template
+    # should still resolve to the underlying physical table
+    audit = ModelAudit(
+        name="test_audit",
+        query="SELECT * FROM @resolve_template('@{catalog_name}.@{schema_name}.@{table_name}$partitions', mode := 'table') WHERE a IS NULL",
+    )
+
+    assert (
+        model_default_catalog.render_audit_query(audit).sql()
+        == """SELECT * FROM "test_catalog"."db"."test_model$partitions" AS "test_model$partitions" WHERE "a" IS NULL"""
+    )
+
+    literal_audit = ModelAudit(
+        name="test_audit",
+        query="SELECT @resolve_template('s3://bucket/@{catalog_name}/@{schema_name}/@{table_name}') AS path FROM @this_model",
+    )
+
+    assert (
+        model_default_catalog.render_audit_query(literal_audit).sql()
+        == """SELECT 's3://bucket/test_catalog/db/test_model' AS "path" FROM (SELECT * FROM "test_catalog"."db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0\""""
+    )
+
+
 def test_load_with_defaults(model, assert_exp_eq):
     expressions = parse(
         """

tests/core/test_macros.py

@@ -1138,6 +1138,49 @@ def test_resolve_template_table():
     )
 
 
+def test_resolve_template_subquery():
+    # Audits on models with a time column render @this_model as a subquery that filters the
+    # physical table on the time range, so the underlying table needs to be extracted from it
+    parsed_sql = parse_one(
+        "SELECT * FROM @resolve_template('@{catalog_name}.@{schema_name}.@{table_name}$partitions', mode := 'table')"
+    )
+
+    evaluator = MacroEvaluator(runtime_stage=RuntimeStage.EVALUATING)
+    evaluator.locals.update(
+        {
+            "this_model": exp.select("*")
+            .from_(exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505"))
+            .where(exp.column("ds").between("2020-01-01", "2020-01-02"))
+            .subquery()
+        }
+    )
+
+    assert (
+        evaluator.transform(parsed_sql).sql(identify=True)
+        == 'SELECT * FROM "test_catalog"."sqlmesh__test"."test__test_model__2517971505$partitions"'
+    )
+
+    # the table is taken from the subquery's FROM clause, so other tables referenced elsewhere
+    # in it (here, in a WHERE subquery) can't be picked up instead
+    evaluator.locals.update(
+        {
+            "this_model": exp.select("*")
+            .from_(exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505"))
+            .where(
+                exp.column("ds").isin(
+                    query=exp.select("ds").from_(exp.to_table("other_catalog.other_schema.other"))
+                )
+            )
+            .subquery()
+        }
+    )
+
+    assert (
+        evaluator.transform(parsed_sql).sql(identify=True)
+        == 'SELECT * FROM "test_catalog"."sqlmesh__test"."test__test_model__2517971505$partitions"'
+    )
+
+
 def test_macro_with_spaces():
     evaluator = MacroEvaluator()
     evaluator.evaluate(d.parse_one(""" @DEF(x, "a b") """))