Fix!: Add default audits in the model properly with their args

Author: themisvaltinos

sqlmesh/core/loader.py

@@ -594,7 +594,6 @@ def _load_sql_models(
                 macros=macros,
                 jinja_macros=jinja_macros,
                 audit_definitions=audits,
-                default_audits=self.config.model_defaults.audits,
                 module_path=self.config_path,
                 dialect=self.config.model_defaults.dialect,
                 time_column_format=self.config.time_column_format,

sqlmesh/core/model/definition.py

@@ -32,7 +32,7 @@
     sorted_python_env_payloads,
     validate_extra_and_required_fields,
 )
-from sqlmesh.core.model.meta import ModelMeta, FunctionCall
+from sqlmesh.core.model.meta import ModelMeta
 from sqlmesh.core.model.kind import (
     ModelKindName,
     SeedKind,
@@ -2038,7 +2038,6 @@ def load_sql_based_model(
     macros: t.Optional[MacroRegistry] = None,
     jinja_macros: t.Optional[JinjaMacroRegistry] = None,
     audits: t.Optional[t.Dict[str, ModelAudit]] = None,
-    default_audits: t.Optional[t.List[FunctionCall]] = None,
     python_env: t.Optional[t.Dict[str, Executable]] = None,
     dialect: t.Optional[str] = None,
     physical_schema_mapping: t.Optional[t.Dict[re.Pattern, str]] = None,
@@ -2211,7 +2210,6 @@ def load_sql_based_model(
         physical_schema_mapping=physical_schema_mapping,
         default_catalog=default_catalog,
         variables=variables,
-        default_audits=default_audits,
         inline_audits=inline_audits,
         blueprint_variables=blueprint_variables,
         **meta_fields,
@@ -2431,7 +2429,6 @@ def _create_model(
     physical_schema_mapping: t.Optional[t.Dict[re.Pattern, str]] = None,
     python_env: t.Optional[t.Dict[str, Executable]] = None,
     audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
-    default_audits: t.Optional[t.List[FunctionCall]] = None,
     inline_audits: t.Optional[t.Dict[str, ModelAudit]] = None,
     module_path: Path = Path(),
     macros: t.Optional[MacroRegistry] = None,
@@ -2541,6 +2538,8 @@ def _create_model(
         for jinja_macro in jinja_macros.root_macros.values():
             used_variables.update(extract_macro_references_and_variables(jinja_macro.definition)[1])
 
+    default_audits = defaults.get("audits", None) if kwargs.get("audits") else None
+
     model = klass(
         name=name,
         **{
@@ -2558,12 +2557,10 @@ def _create_model(
         **(inline_audits or {}),
     }
 
-    # TODO: default_audits needs to be merged with model.audits; the former's arguments
-    # are silently dropped today because we add them in audit_definitions. We also need
-    # to check for duplicates when we implement this merging logic.
-    used_audits: t.Set[str] = set()
-    used_audits.update(audit_name for audit_name, _ in default_audits or [])
-    used_audits.update(audit_name for audit_name, _ in model.audits)
+    if default_audits:
+        model = model.copy(update={"audits": default_audits + model.audits})
+
+    used_audits: t.Set[str] = {audit_name for audit_name, _ in model.audits}
 
     audit_definitions = {
         audit_name: audit_definitions[audit_name]

sqlmesh/migrations/v0088_include_default_audits_in_model.py

@@ -0,0 +1,5 @@
+"""Include the model defaults audits along with their args in the model."""
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    pass

tests/core/test_audit.py

@@ -3,6 +3,7 @@
 from sqlglot import exp, parse_one
 
 from sqlmesh.core import constants as c
+from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.context import Context
 from sqlmesh.core.audit import (
     ModelAudit,
@@ -962,6 +963,117 @@ def test_multiple_audits_with_same_name():
     assert model.audits[1][1] == model.audits[2][1]
 
 
+def test_default_audits_included_when_no_model_audits():
+    expressions = parse("""
+    MODEL (
+        name test.basic_model
+    );
+    SELECT 1 as id, 'test' as name;
+    """)
+
+    model_defaults = ModelDefaultsConfig(
+        dialect="duckdb", audits=["not_null(columns := ['id'])", "unique_values(columns := ['id'])"]
+    )
+    model = load_sql_based_model(expressions, defaults=model_defaults.dict())
+
+    assert len(model.audits) == 2
+    audit_names = [audit[0] for audit in model.audits]
+    assert "not_null" in audit_names
+    assert "unique_values" in audit_names
+
+    # Verify arguments are preserved
+    for audit_name, audit_args in model.audits:
+        if audit_name == "not_null":
+            assert "columns" in audit_args
+            assert audit_args["columns"].expressions[0].this == "id"
+        elif audit_name == "unique_values":
+            assert "columns" in audit_args
+            assert audit_args["columns"].expressions[0].this == "id"
+
+    for audit_name, audit_args in model.audits_with_args:
+        if audit_name == "not_null":
+            assert "columns" in audit_args
+            assert audit_args["columns"].expressions[0].this == "id"
+        elif audit_name == "unique_values":
+            assert "columns" in audit_args
+            assert audit_args["columns"].expressions[0].this == "id"
+
+
+def test_model_defaults_audits_with_same_name():
+    expressions = parse(
+        """
+        MODEL (
+            name db.table,
+            dialect spark,
+            audits(
+                does_not_exceed_threshold(column := id, threshold := 1000),
+                does_not_exceed_threshold(column := price, threshold := 100),
+                unique_values(columns := ['id'])
+            )
+        );
+
+        SELECT id, price FROM tbl;
+
+        AUDIT (
+            name does_not_exceed_threshold,
+        );
+        SELECT * FROM @this_model
+        WHERE @column >= @threshold;
+        """
+    )
+
+    model_defaults = ModelDefaultsConfig(
+        dialect="duckdb",
+        audits=[
+            "does_not_exceed_threshold(column := price, threshold := 33)",
+            "does_not_exceed_threshold(column := id, threshold := 65)",
+            "not_null(columns := ['id'])",
+        ],
+    )
+    model = load_sql_based_model(expressions, defaults=model_defaults.dict())
+    assert len(model.audits) == 6
+    assert len(model.audits_with_args) == 6
+    assert len(model.audit_definitions) == 1
+
+    expected_audits = [
+        (
+            "does_not_exceed_threshold",
+            {"column": exp.column("price"), "threshold": exp.Literal.number(33)},
+        ),
+        (
+            "does_not_exceed_threshold",
+            {"column": exp.column("id"), "threshold": exp.Literal.number(65)},
+        ),
+        ("not_null", {"columns": exp.convert(["id"])}),
+        (
+            "does_not_exceed_threshold",
+            {"column": exp.column("id"), "threshold": exp.Literal.number(1000)},
+        ),
+        (
+            "does_not_exceed_threshold",
+            {"column": exp.column("price"), "threshold": exp.Literal.number(100)},
+        ),
+        ("unique_values", {"columns": exp.convert(["id"])}),
+    ]
+
+    for (actual_name, actual_args), (expected_name, expected_args) in zip(
+        model.audits, expected_audits
+    ):
+        # Validate the audit names are preserved
+        assert actual_name == expected_name
+        for key in expected_args:
+            # comparing sql representaion is easier
+            assert actual_args[key].sql() == expected_args[key].sql()
+
+    # Validate audits with args as well along with their arguments
+    for (actual_audit, actual_args), (expected_name, expected_args) in zip(
+        model.audits_with_args, expected_audits
+    ):
+        assert actual_audit.name == expected_name
+        for key in expected_args:
+            assert actual_args[key].sql() == expected_args[key].sql()
+
+
 def test_audit_formatting_flag_serde():
     expressions = parse(
         """