Repository: TobikoData/sqlmesh
Pull request #2979 (2024-08-01 23:00:57)

Fix!: Send notifications for failing non-blocking audits from the builtin scheduler
Currently, the builtin-scheduler only sends notifications when the `SnapshotEvaluator` raises an `AuditError` exception, which only happens for failing, blocking audits.

This PR addresses two things. First, it delays when the `AuditError` is raised, so that all of the snapshot's audits execute first and second, it sends notifications for all audits that failed, blocking or non-blocking.


Patch:
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
index 6364b842..0d1466da 100644
--- a/sqlmesh/core/scheduler.py
+++ b/sqlmesh/core/scheduler.py
@@ -6,6 +6,8 @@ import traceback
 import typing as t
 from datetime import datetime
 
+from sqlglot import exp
+
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.environment import EnvironmentNamingInfo
@@ -212,25 +214,38 @@ class Scheduler:
             batch_index=batch_index,
             **kwargs,
         )
-        try:
-            self.snapshot_evaluator.audit(
-                snapshot=snapshot,
-                start=start,
-                end=end,
-                execution_time=execution_time,
-                snapshots=snapshots,
-                deployability_index=deployability_index,
-                wap_id=wap_id,
-                **kwargs,
+        audit_results = self.snapshot_evaluator.audit(
+            snapshot=snapshot,
+            start=start,
+            end=end,
+            execution_time=execution_time,
+            raise_exception=False,
+            snapshots=snapshots,
+            deployability_index=deployability_index,
+            wap_id=wap_id,
+            **kwargs,
+        )
+
+        audit_error_to_raise: t.Optional[AuditError] = None
+        for audit_result in (result for result in audit_results if result.count):
+            error = AuditError(
+                audit_name=audit_result.audit.name,
+                model=snapshot.model_or_none,
+                count=t.cast(int, audit_result.count),
+                query=t.cast(exp.Query, audit_result.query),
+                adapter_dialect=self.snapshot_evaluator.adapter.dialect,
             )
-        except AuditError as e:
-            self.notification_target_manager.notify(NotificationEvent.AUDIT_FAILURE, e)
+            self.notification_target_manager.notify(NotificationEvent.AUDIT_FAILURE, error)
             if is_deployable and snapshot.node.owner:
                 self.notification_target_manager.notify_user(
-                    NotificationEvent.AUDIT_FAILURE, snapshot.node.owner, e
+                    NotificationEvent.AUDIT_FAILURE, snapshot.node.owner, error
                 )
+            if audit_result.audit.blocking:
+                audit_error_to_raise = error
+
+        if audit_error_to_raise:
             logger.error(f"Audit Failure: {traceback.format_exc()}")
-            raise e
+            raise audit_error_to_raise
 
         self.state_sync.add_interval(snapshot, start, end, is_dev=not is_deployable)
 


Test patch:
diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py
index a0ef5a89..92b62d68 100644
--- a/tests/core/test_scheduler.py
+++ b/tests/core/test_scheduler.py
@@ -4,6 +4,7 @@ import pytest
 from pytest_mock.plugin import MockerFixture
 from sqlglot import parse_one, parse
 
+from sqlmesh.core.audit import AuditResult
 from sqlmesh.core.context import Context
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.model import load_sql_based_model
@@ -15,9 +16,14 @@ from sqlmesh.core.model.kind import (
 )
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.scheduler import Scheduler, compute_interval_params
-from sqlmesh.core.snapshot import Snapshot, SnapshotEvaluator, SnapshotChangeCategory
+from sqlmesh.core.snapshot import (
+    Snapshot,
+    SnapshotEvaluator,
+    SnapshotChangeCategory,
+    DeployabilityIndex,
+)
 from sqlmesh.utils.date import to_datetime
-from sqlmesh.utils.errors import CircuitBreakerError
+from sqlmesh.utils.errors import CircuitBreakerError, AuditError
 
 
 @pytest.fixture
@@ -30,6 +36,11 @@ def orders(sushi_context_fixed_date: Context) -> Snapshot:
     return sushi_context_fixed_date.get_snapshot("sushi.orders", raise_if_missing=True)
 
 
+@pytest.fixture
+def waiter_names(sushi_context_fixed_date: Context) -> Snapshot:
+    return sushi_context_fixed_date.get_snapshot("sushi.waiter_names", raise_if_missing=True)
+
+
 @pytest.mark.slow
 def test_interval_params(scheduler: Scheduler, sushi_context_fixed_date: Context, orders: Snapshot):
     waiter_revenue = sushi_context_fixed_date.get_snapshot(
@@ -546,3 +557,65 @@ def test_check_ready_intervals(mocker: MockerFixture):
         [[(0, 1)], [(3, 4)]],
         [(0, 1), (3, 4)],
     )
+
+
+def test_audit_failure_notifications(
+    scheduler: Scheduler, waiter_names: Snapshot, mocker: MockerFixture
+):
+    evaluator_evaluate_mock = mocker.Mock()
+    mocker.patch("sqlmesh.core.scheduler.SnapshotEvaluator.evaluate", evaluator_evaluate_mock)
+    evaluator_audit_mock = mocker.Mock()
+    mocker.patch("sqlmesh.core.scheduler.SnapshotEvaluator.audit", evaluator_audit_mock)
+    notify_user_mock = mocker.Mock()
+    mocker.patch(
+        "sqlmesh.core.notification_target.NotificationTargetManager.notify_user", notify_user_mock
+    )
+    notify_mock = mocker.Mock()
+    mocker.patch("sqlmesh.core.notification_target.NotificationTargetManager.notify", notify_mock)
+
+    audit = next(iter(waiter_names.audits))
+
+    def _evaluate():
+        scheduler.evaluate(
+            waiter_names,
+            to_datetime("2022-01-01"),
+            to_datetime("2022-01-02"),
+            to_datetime("2022-01-03"),
+            DeployabilityIndex.all_deployable(),
+            0,
+        )
+        _, kwargs = evaluator_audit_mock.call_args_list[0]
+        assert not kwargs["raise_exception"]
+
+    evaluator_audit_mock.return_value = [
+        AuditResult(audit=audit, model=waiter_names.model, count=0, skipped=False)
+    ]
+    _evaluate()
+    assert notify_user_mock.call_count == 0
+    assert notify_mock.call_count == 0
+
+    evaluator_audit_mock.return_value = [
+        AuditResult(audit=audit, model=waiter_names.model, count=None, skipped=True)
+    ]
+    _evaluate()
+    assert notify_user_mock.call_count == 0
+    assert notify_mock.call_count == 0
+
+    audit = audit.copy(update={"blocking": False})
+    evaluator_audit_mock.return_value = [
+        AuditResult(audit=audit, model=waiter_names.model, count=1, skipped=False)
+    ]
+    _evaluate()
+    assert notify_user_mock.call_count == 1
+    assert notify_mock.call_count == 1
+    notify_user_mock.reset_mock()
+    notify_mock.reset_mock()
+
+    audit = audit.copy(update={"blocking": True})
+    evaluator_audit_mock.return_value = [
+        AuditResult(audit=audit, model=waiter_names.model, count=1, skipped=False)
+    ]
+    with pytest.raises(AuditError):
+        _evaluate()
+    assert notify_user_mock.call_count == 1
+    assert notify_mock.call_count == 1