Feat: return non-zero code when at least one audit fails in sqlmesh audit (#4258)

Author: georgesittas

sqlmesh/cli/main.py

@@ -690,7 +690,8 @@ def audit(
     execution_time: t.Optional[TimeLike] = None,
 ) -> None:
     """Run audits for the target model(s)."""
-    obj.audit(models=models, start=start, end=end, execution_time=execution_time)
+    if not obj.audit(models=models, start=start, end=end, execution_time=execution_time):
+        exit(1)
 
 
 @cli.command("check_intervals")

sqlmesh/core/context.py

@@ -1827,14 +1827,17 @@ def audit(
         *,
         models: t.Optional[t.Iterator[str]] = None,
         execution_time: t.Optional[TimeLike] = None,
-    ) -> None:
+    ) -> bool:
         """Audit models.
 
         Args:
             start: The start of the interval to audit.
             end: The end of the interval to audit.
             models: The models to audit. All models will be audited if not specified.
             execution_time: The date/time time reference to use for execution time. Defaults to now.
+
+        Returns:
+            False if any of the audits failed, True otherwise.
         """
 
         snapshots = (
@@ -1845,6 +1848,7 @@ def audit(
 
         num_audits = sum(len(snapshot.node.audits_with_args) for snapshot in snapshots)
         self.console.log_status_update(f"Found {num_audits} audit(s).")
+
         errors = []
         skipped_count = 0
         for snapshot in snapshots:
@@ -1884,6 +1888,7 @@ def audit(
                 )
 
         self.console.log_status_update("Done.")
+        return not errors
 
     @python_api_analytics
     def rewrite(self, sql: str, dialect: str = "") -> exp.Expression:

sqlmesh/magics.py

@@ -1000,10 +1000,10 @@ def run_test(self, context: Context, line: str) -> None:
     @argument("--execution-time", type=str, help="Execution time.")
     @line_magic
     @pass_sqlmesh_context
-    def audit(self, context: Context, line: str) -> None:
+    def audit(self, context: Context, line: str) -> bool:
         """Run audit(s)"""
         args = parse_argstring(self.audit, line)
-        context.audit(
+        return context.audit(
             models=args.models, start=args.start, end=args.end, execution_time=args.execution_time
         )
 

tests/core/test_context.py

@@ -2003,3 +2003,41 @@ def test_check_intervals(sushi_context, mocker):
         environment=None, no_signals=False, select_models=["*waiter_as_customer*"], end="next week"
     )
     assert tuple(intervals.values())[0].intervals
+
+
+def test_audit():
+    context = Context(config=Config())
+
+    parsed_model = parse(
+        """
+        MODEL (
+          name dummy,
+          audits (
+            not_null_non_blocking(columns=[c])
+          )
+        );
+
+        SELECT NULL AS c
+        """
+    )
+    context.upsert_model(load_sql_based_model(parsed_model))
+    context.plan(no_prompts=True, auto_apply=True)
+
+    assert context.audit(models=["dummy"], start="2020-01-01", end="2020-01-01") is False
+
+    parsed_model = parse(
+        """
+        MODEL (
+          name dummy,
+          audits (
+            not_null_non_blocking(columns=[c])
+          )
+        );
+
+        SELECT 1 AS c
+        """
+    )
+    context.upsert_model(load_sql_based_model(parsed_model))
+    context.plan(no_prompts=True, auto_apply=True)
+
+    assert context.audit(models=["dummy"], start="2020-01-01", end="2020-01-01") is True