Repository: TobikoData/sqlmesh
Pull request #2944 (2024-07-25 15:16:42)

feat: add GenericNotificationTarget to NotificationTargets
The GenericNotificationTarget object is a subclass of BaseNotificationTarget, and will offer users greater freedom in creating their own custom notifications.

This will be useful for users who do want to use the currently available methods (smtp, slack_webhook, slack_api).

vchan: Thank you for the PR, @kam-cheng-emis! How about adding a generic NotificationTarget class that subclasses BaseNotificationTarget?
```python
class GenericNotificationTarget(BaseNotificationTarget):
    type_: Literal["generic"] = Field(alias="type", default="generic")
```
You can call it simple or generic or whatever makes sense. Then we don't need to change the discriminator.

Patch:
diff --git a/docs/concepts/audits.md b/docs/concepts/audits.md
index de586e77..0762480b 100644
--- a/docs/concepts/audits.md
+++ b/docs/concepts/audits.md
@@ -88,6 +88,15 @@ SELECT * FROM @this_model
 WHERE @column >= @threshold;
 ```
 
+Alternatively, you can apply specific audits globally by including them in the model defaults configuration:
+
+```sql linenums="1"
+model_defaults:
+  audits: 
+    - assert_positive_order_ids
+    - does_not_exceed_threshold(column := id, threshold := 1000)
+```
+
 ### Naming
 We recommended avoiding SQL keywords when naming audit parameters. Quote any audit argument that is also a SQL keyword.
 
diff --git a/docs/concepts/overview.md b/docs/concepts/overview.md
index cabf0736..b1c7f288 100644
--- a/docs/concepts/overview.md
+++ b/docs/concepts/overview.md
@@ -61,7 +61,7 @@ You create audits by writing SQL queries that should return 0 rows. For example,
 
 Audits are flexible &mdash; they can be tied to a specific model's contents, or you can use [macros](./macros/overview.md) to create audits that are usable by multiple models. SQLMesh also includes pre-made audits for common use cases, such as detecting NULL or duplicated values.
 
-You specify which audits should run for a model by including them in the model's metadata properties.
+You specify which audits should run for a model by including them in the model's metadata properties. To apply them globally across your project, include them in the model defaults configuration.
 
 SQLMesh automatically runs audits when you apply a `plan` to an environment, or you can run them on demand with the [`audit` command](../reference/cli.md#audit).
 
diff --git a/docs/reference/model_configuration.md b/docs/reference/model_configuration.md
index 820f9f2b..93e00c56 100644
--- a/docs/reference/model_configuration.md
+++ b/docs/reference/model_configuration.md
@@ -51,6 +51,7 @@ The SQLMesh project-level `model_defaults` key supports the following options, d
 - storage_format
 - session_properties (on per key basis)
 - on_destructive_change (described [below](#incremental-models))
+- audits (described [here](../concepts/audits.md#generic-audits))
 
 
 ### Model Naming
diff --git a/sqlmesh/core/config/model.py b/sqlmesh/core/config/model.py
index 79673428..31f9a746 100644
--- a/sqlmesh/core/config/model.py
+++ b/sqlmesh/core/config/model.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import typing as t
 
+from sqlmesh.core.dialect import parse_one, extract_audit
 from sqlmesh.core.config.base import BaseConfig
 from sqlmesh.core.model.kind import (
     ModelKind,
@@ -10,6 +11,8 @@ from sqlmesh.core.model.kind import (
     on_destructive_change_validator,
 )
 from sqlmesh.utils.date import TimeLike
+from sqlmesh.core.model.meta import AuditReference
+from sqlmesh.utils.pydantic import field_validator
 
 
 class ModelDefaultsConfig(BaseConfig):
@@ -27,6 +30,7 @@ class ModelDefaultsConfig(BaseConfig):
         storage_format: The storage format used to store the physical table, only applicable in certain engines.
             (eg. 'parquet')
         on_destructive_change: What should happen when a forward-only model requires a destructive schema change.
+        audits: The audits to be applied globally to all models in the project.
     """
 
     kind: t.Optional[ModelKind] = None
@@ -37,6 +41,14 @@ class ModelDefaultsConfig(BaseConfig):
     storage_format: t.Optional[str] = None
     on_destructive_change: t.Optional[OnDestructiveChange] = None
     session_properties: t.Optional[t.Dict[str, t.Any]] = None
+    audits: t.Optional[t.List[AuditReference]] = None
 
     _model_kind_validator = model_kind_validator
     _on_destructive_change_validator = on_destructive_change_validator
+
+    @field_validator("audits", mode="before")
+    def _audits_validator(cls, v: t.Any) -> t.Any:
+        if isinstance(v, list):
+            return [extract_audit(parse_one(audit)) for audit in v]
+
+        return v
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 6fa14fb6..85de1811 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1907,6 +1907,7 @@ class GenericContext(BaseContext, t.Generic[C]):
                     audits=audits,
                     cache=fingerprint_cache,
                     ttl=ttl,
+                    config=self.config_for_node(node),
                 )
                 snapshots[snapshot.name] = snapshot
             return snapshots
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index ccefc616..68fd73ed 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -21,7 +21,7 @@ from sqlglot.schema import MappingSchema
 from sqlglot.tokens import Token
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
-from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.errors import SQLMeshError, ConfigError
 from sqlmesh.utils.pandas import columns_to_types_from_df
 
 if t.TYPE_CHECKING:
@@ -1099,3 +1099,24 @@ def interpret_key_value_pairs(
     e: exp.Tuple,
 ) -> t.Dict[str, exp.Expression | str | int | float | bool]:
     return {i.this.name: interpret_expression(i.expression) for i in e.expressions}
+
+
+def extract_audit(v: exp.Expression) -> t.Tuple[str, t.Dict[str, exp.Expression]]:
+    kwargs = {}
+
+    if isinstance(v, exp.Anonymous):
+        func = v.name
+        args = v.expressions
+    elif isinstance(v, exp.Func):
+        func = v.sql_name()
+        args = list(v.args.values())
+    else:
+        return v.name.lower(), {}
+
+    for arg in args:
+        if not isinstance(arg, (exp.PropertyEQ, exp.EQ)):
+            raise ConfigError(
+                f"Function '{func}' must be called with key-value arguments like {func}(arg := value)."
+            )
+        kwargs[arg.left.name.lower()] = arg.right
+    return func.lower(), kwargs
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 54b9febc..6fef5baf 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -25,7 +25,7 @@ from sqlmesh.core import dialect as d
 from sqlmesh.core.macros import MacroRegistry, MacroStrTemplate, macro
 from sqlmesh.core.model.common import expression_validator
 from sqlmesh.core.model.kind import ModelKindName, SeedKind, ModelKind, FullKind, create_model_kind
-from sqlmesh.core.model.meta import ModelMeta
+from sqlmesh.core.model.meta import ModelMeta, AuditReference
 from sqlmesh.core.model.seed import CsvSeedReader, Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
 from sqlmesh.utils import columns_to_types_all_known, str_to_bool, UniqueKeyDict
@@ -461,7 +461,11 @@ class _Model(ModelMeta, frozen=True):
             )
         return query
 
-    def referenced_audits(self, audits: t.Dict[str, ModelAudit]) -> t.List[ModelAudit]:
+    def referenced_audits(
+        self,
+        audits: t.Dict[str, ModelAudit],
+        default_audits: t.List[AuditReference] = [],
+    ) -> t.List[ModelAudit]:
         """Returns audits referenced in this model.
 
         Args:
@@ -471,7 +475,7 @@ class _Model(ModelMeta, frozen=True):
 
         referenced_audits = []
 
-        for audit_name, _ in self.audits:
+        for audit_name, _ 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:
diff --git a/sqlmesh/core/model/kind.py b/sqlmesh/core/model/kind.py
index b635cf84..d7a6e504 100644
--- a/sqlmesh/core/model/kind.py
+++ b/sqlmesh/core/model/kind.py
@@ -488,6 +488,14 @@ class IncrementalByPartitionKind(_Incremental):
     forward_only: Literal[True] = True
     disable_restatement: SQLGlotBool = True
 
+    @field_validator("forward_only", mode="before")
+    def _forward_only_validator(cls, v: t.Union[bool, exp.Expression]) -> Literal[True]:
+        if v is not True:
+            raise ConfigError(
+                "Do not specify the `forward_only` configuration key - INCREMENTAL_BY_PARTITION models are always forward_only."
+            )
+        return v
+
     @property
     def metadata_hash_values(self) -> t.List[t.Optional[str]]:
         return [
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index 86a2e6e9..d36a5f3d 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -10,7 +10,7 @@ from sqlglot.helper import ensure_collection, ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core import dialect as d
-from sqlmesh.core.dialect import normalize_model_name
+from sqlmesh.core.dialect import normalize_model_name, extract_audit
 from sqlmesh.core.model.common import (
     bool_validator,
     default_catalog_validator,
@@ -85,32 +85,12 @@ class ModelMeta(_Node):
 
     @field_validator("audits", mode="before")
     def _audits_validator(cls, v: t.Any) -> t.Any:
-        def extract(v: exp.Expression) -> t.Tuple[str, t.Dict[str, exp.Expression]]:
-            kwargs = {}
-
-            if isinstance(v, exp.Anonymous):
-                func = v.name
-                args = v.expressions
-            elif isinstance(v, exp.Func):
-                func = v.sql_name()
-                args = list(v.args.values())
-            else:
-                return v.name.lower(), {}
-
-            for arg in args:
-                if not isinstance(arg, (exp.PropertyEQ, exp.EQ)):
-                    raise ConfigError(
-                        f"Function '{func}' must be called with key-value arguments like {func}(arg := value)."
-                    )
-                kwargs[arg.left.name.lower()] = arg.right
-            return func.lower(), kwargs
-
         if isinstance(v, (exp.Tuple, exp.Array)):
-            return [extract(i) for i in v.expressions]
+            return [extract_audit(i) for i in v.expressions]
         if isinstance(v, exp.Paren):
-            return [extract(v.this)]
+            return [extract_audit(v.this)]
         if isinstance(v, exp.Expression):
-            return [extract(v)]
+            return [extract_audit(v)]
         if isinstance(v, list):
             audits = []
 
diff --git a/sqlmesh/core/notification_target.py b/sqlmesh/core/notification_target.py
index 7eaf1c3a..3fd68313 100644
--- a/sqlmesh/core/notification_target.py
+++ b/sqlmesh/core/notification_target.py
@@ -394,9 +394,36 @@ class BasicSMTPNotificationTarget(BaseTextBasedNotificationTarget):
         return all((self.host, self.user, self.password, self.sender))
 
 
+class GenericNotificationTarget(BaseNotificationTarget):
+    """A generic notification target that can be used to create custom notification targets.
+
+    This target is not meant to be used directly, but rather as a base class for custom notification targets.
+
+    The `send` method should be overridden to provide the actual notification functionality.
+
+    Example:
+    ```python
+    class MyCustomNotificationTarget(GenericNotificationTarget):
+        def send(self, notification_status: NotificationStatus, msg: str, audit_error: t.Optional[AuditError] = None, exc: t.Optional[str] = None, **kwargs: t.Any) -> None:
+            error = None
+            if audit_error:
+                error = str(audit_error)
+            elif exc:
+                error = exc
+
+            if error:
+                msg = f"{error} - {msg}"
+            print(f"Sending notification: {msg}")
+    ```
+    """
+
+    type_: Literal["generic"] = Field(alias="type", default="generic")
+
+
 NotificationTarget = Annotated[
     t.Union[
         BasicSMTPNotificationTarget,
+        GenericNotificationTarget,
         ConsoleNotificationTarget,
         SlackApiNotificationTarget,
         SlackWebhookNotificationTarget,
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 57143f1d..90adcabf 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -46,6 +46,7 @@ else:
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
     from sqlmesh.core.environment import EnvironmentNamingInfo
+    from sqlmesh.core.config import Config
 
 Interval = t.Tuple[int, int]
 Intervals = t.List[Interval]
@@ -617,6 +618,7 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         version: t.Optional[str] = None,
         audits: t.Optional[t.Dict[str, ModelAudit]] = None,
         cache: t.Optional[t.Dict[str, SnapshotFingerprint]] = None,
+        config: t.Optional[Config] = None,
     ) -> Snapshot:
         """Creates a new snapshot for a node.
 
@@ -634,8 +636,13 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         """
         created_ts = now_timestamp()
         kwargs = {}
+        default_audits = (
+            config.model_defaults.audits if (config and config.model_defaults.audits) else []
+        )
         if node.is_model:
-            kwargs["audits"] = tuple(t.cast(_Model, node).referenced_audits(audits or {}))
+            kwargs["audits"] = tuple(
+                t.cast(_Model, node).referenced_audits(audits or {}, default_audits)
+            )
 
         return cls(
             name=node.fqn,


Test patch:
diff --git a/.circleci/test_migration.sh b/.circleci/test_migration.sh
index fc869eb4..c5e0ea8d 100755
--- a/.circleci/test_migration.sh
+++ b/.circleci/test_migration.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -ex
 
-CONFIG_NAME="local_config"
+GATEWAY_NAME="duckdb_persistent"
 TMP_DIR=$(mktemp -d)
 SUSHI_DIR="$TMP_DIR/sushi"
 
@@ -24,7 +24,7 @@ cp -r ./examples/sushi $TMP_DIR
 # Run initial plan
 pushd $SUSHI_DIR
 rm -rf ./data/*
-sqlmesh --config $CONFIG_NAME plan --no-prompts --auto-apply
+sqlmesh --gateway $GATEWAY_NAME plan --no-prompts --auto-apply
 popd
 
 # Switch back to the starting state of the repository
@@ -35,7 +35,6 @@ make install-dev
 
 # Migrate and make sure the diff is empty
 pushd $SUSHI_DIR
-sqlmesh --config $CONFIG_NAME migrate
-sqlmesh --config $CONFIG_NAME diff prod
+sqlmesh --gateway $GATEWAY_NAME migrate
+sqlmesh --gateway $GATEWAY_NAME diff prod
 popd
-
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index d74efbd4..9e831731 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -557,3 +557,30 @@ model_defaults:
     assert attach_config_2.type == "postgres"
     assert attach_config_2.path == "dbname=postgres user=postgres host=127.0.0.1"
     assert attach_config_2.read_only is True
+
+
+def test_load_model_defaults_audits(tmp_path):
+    config_path = tmp_path / "config_model_defaults_audits.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+model_defaults:
+    dialect: ''
+    audits: 
+        - assert_positive_order_ids
+        - does_not_exceed_threshold(column := id, threshold := 1000)
+        """
+        )
+
+    config = load_config_from_paths(
+        Config,
+        project_paths=[config_path],
+    )
+
+    assert len(config.model_defaults.audits) == 2
+    assert config.model_defaults.audits[0] == ("assert_positive_order_ids", {})
+    assert config.model_defaults.audits[1][0] == "does_not_exceed_threshold"
+    assert type(config.model_defaults.audits[1][1]["column"]) == exp.Column
+    assert config.model_defaults.audits[1][1]["column"].this.this == "id"
+    assert type(config.model_defaults.audits[1][1]["threshold"]) == exp.Literal
+    assert config.model_defaults.audits[1][1]["threshold"].this == "1000"
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 0c00c790..ec9b45f1 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -15,6 +15,7 @@ from sqlmesh.cli.example_project import init_example_project
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
+from sqlmesh.core.audit import ModelAudit
 from sqlmesh.core.config import (
     Config,
     NameInferenceConfig,
@@ -969,6 +970,39 @@ def test_audits():
     assert model.tags == ["foo"]
 
 
+def test_enable_audits_from_model_defaults():
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.audit_model,
+        );
+        SELECT 1 as id;
+
+        AUDIT (
+    name assert_positive_order_ids,
+    );
+    SELECT *
+    FROM @this_model
+    WHERE
+    id < 0;
+    """
+    )
+
+    model = load_sql_based_model(expressions, path=Path("./examples/sushi/models/test_model.sql"))
+    assert len(model.audits) == 0
+    assert len(model.inline_audits) == 1
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb", audits=["assert_positive_order_ids"])
+    )
+    assert config.model_defaults.audits[0] == ("assert_positive_order_ids", {})
+
+    snapshot = Snapshot.from_node(model, nodes={}, config=config)
+    assert len(snapshot.audits) == 1
+    assert type(snapshot.audits[0]) == ModelAudit
+    assert snapshot.audits[0].query.sql() == "SELECT * FROM @this_model WHERE id < 0"
+
+
 def test_description(sushi_context):
     assert sushi_context.models['"memory"."sushi"."orders"'].description == "Table of sushi orders."
 
@@ -4721,6 +4755,24 @@ def test_incremental_by_partition(sushi_context, assert_exp_eq):
         )
         load_sql_based_model(expressions)
 
+    with pytest.raises(
+        ConfigError,
+        match=r".*Do not specify the `forward_only` configuration key.*",
+    ):
+        expressions = d.parse(
+            """
+            MODEL (
+                name db.table,
+                kind INCREMENTAL_BY_PARTITION (
+                    forward_only true
+                ),
+            );
+
+            SELECT a, b
+            """
+        )
+        load_sql_based_model(expressions)
+
 
 @pytest.mark.parametrize(
     ["model_def", "path", "expected_name"],