Repository: TobikoData/sqlmesh
Pull request #3106 (2024-09-06 21:02:45)

Fix: allow overriding blocking setting of builtin audits
Context: https://tobiko-data.slack.com/archives/C044BRE5W4S/p1725579711884559


Patch:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 96f5fbc9..b69eec2a 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -559,16 +559,28 @@ class _Model(ModelMeta, frozen=True):
 
         referenced_audits = []
 
-        for audit_name, _ in self.audits + default_audits:
+        for audit_name, audit_args in self.audits + default_audits:
             if audit_name in self.inline_audits:
                 referenced_audits.append(self.inline_audits[audit_name])
             elif audit_name in audits:
                 referenced_audits.append(audits[audit_name])
-            elif audit_name not in BUILT_IN_AUDITS:
-                raise_config_error(
-                    f"Unknown audit '{audit_name}' referenced in model '{self.name}'",
-                    self._path,
-                )
+            else:
+                audit = BUILT_IN_AUDITS.get(audit_name)
+                if not audit:
+                    raise_config_error(
+                        f"Unknown audit '{audit_name}' referenced in model '{self.name}'",
+                        self._path,
+                    )
+
+                # Builtin audits are generally not included in order to reduce fingerprint size,
+                # but those that override `blocking` need to be included because otherwise doing
+                # `audit.blocking` will always return the builtin audit's default value
+                blocking = audit_args.get("blocking")
+                if blocking:
+                    referenced_audits.append(
+                        audit.copy(update={"blocking": blocking == exp.true()})  # type: ignore
+                    )
+
         return referenced_audits
 
     def text_diff(self, other: Node) -> str:


Test patch:
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index 3f163ca1..14e31c37 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -292,6 +292,59 @@ def test_plan_execution_time():
     )
 
 
+def test_override_builtin_audit_blocking_mode():
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.x,
+                    kind FULL,
+                    audits (
+                        not_null(columns := [c], blocking := false),
+                        unique_values(columns := [c]),
+                    )
+                );
+
+                SELECT NULL AS c
+                """
+            )
+        )
+    )
+
+    plan = context.plan(auto_apply=True, no_prompts=True)
+    new_snapshot = next(iter(plan.context_diff.new_snapshots.values()))
+
+    # Even though there are two builtin audits referenced in the above definition, we only
+    # store the one that overrides `blocking` in the snapshot; the other one isn't needed
+    assert len(new_snapshot.audits) == 1
+    assert new_snapshot.audits[0].name == "not_null"
+    assert new_snapshot.audits[0].blocking is False
+
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.x,
+                    kind FULL,
+                    audits (
+                        not_null_non_blocking(columns := [c], blocking := true)
+                    )
+                );
+
+                SELECT NULL AS c
+                """
+            )
+        )
+    )
+
+    with pytest.raises(SQLMeshError):
+        context.plan(auto_apply=True, no_prompts=True)
+
+
 def test_python_model_empty_df_raises(sushi_context, capsys):
     @model(
         "memory.sushi.test_model",