Repository: TobikoData/sqlmesh
Pull request #2812 (2024-06-24 11:20:19)

Feat: Add support for audits within model definition
Τhis update introduces the ability to define audits directly within a model definition instead of a separate file, fixes:  #546 

The syntax for defining audits remains the same ([Audits Docs](https://sqlmesh.readthedocs.io/en/stable/concepts/audits/)), while multiple audits can be specified within the same SQL file:

```
MODEL (
    name sushi.items,
    audits(does_not_exceed_threshold(column := id, threshold := 1000), price_is_not_null)
);
SELECT  ...;

AUDIT (name does_not_exceed_threshold);
SELECT * FROM @this_model
WHERE @column >= @threshold;

AUDIT (name price_is_not_null);
SELECT * FROM @this_model
WHERE price IS NULL;
```

Themiscodes: Addressed comments and adapted logic and tests for snapshot with inline_audits
georgesittas: @Themiscodes I think the direction here is good, just need to clean some parts up a bit. Let's also define an inline audit within a model under examples/sushi, so that it's tested end-to-end.
Themiscodes: Fixed Pydantic issues with @georgesittas in the validator, moved audit loading in the parsing stage, added docs and tests.

Patch:
diff --git a/docs/concepts/audits.md b/docs/concepts/audits.md
index b8b6d73a..61643803 100644
--- a/docs/concepts/audits.md
+++ b/docs/concepts/audits.md
@@ -10,7 +10,7 @@ A comprehensive suite of audits can identify data issues upstream, whether they
 **NOTE**: For incremental models, audits are only applied to intervals being processed - not for the entire underlying table.
 
 ## User-Defined Audits
-In SQLMesh, user-defined audits are defined in `.sql` files in an `audit` directory in your SQLMesh project. Multiple audits can be defined in a single file, so you can organize them to your liking.
+In SQLMesh, user-defined audits are defined in `.sql` files in an `audit` directory in your SQLMesh project. Multiple audits can be defined in a single file, so you can organize them to your liking. Alternatively, audits can be defined inline within the model definition itself.
 
 Audits are SQL queries that should not return any rows; in other words, they query for bad data, so returned rows indicates that something is wrong.
 
@@ -89,6 +89,28 @@ MODEL (
 )
 ```
 
+### Inline audits
+
+You can also define audits directly within a model definition using the same syntax. Multiple audits can be specified within the same SQL model file:
+
+
+```sql linenums="1"
+MODEL (
+    name sushi.items,
+    audits(does_not_exceed_threshold(column := id, threshold := 1000), price_is_not_null)
+);
+SELECT id, price 
+FROM sushi.seed;
+
+AUDIT (name does_not_exceed_threshold);
+SELECT * FROM @this_model
+WHERE @column >= @threshold;
+
+AUDIT (name price_is_not_null);
+SELECT * FROM @this_model
+WHERE price IS NULL;
+```
+
 ## Built-in audits
 SQLMesh comes with a suite of built-in generic audits that cover a broad set of common use cases. Built-in audits are blocking by default, but they all have non-blocking counterparts which you can use by appending `_non_blocking` - see [Non-blocking audits](#non-blocking-audits).
 
diff --git a/examples/sushi/models/waiter_names.sql b/examples/sushi/models/waiter_names.sql
index 22749e25..9f3ac04f 100644
--- a/examples/sushi/models/waiter_names.sql
+++ b/examples/sushi/models/waiter_names.sql
@@ -6,5 +6,25 @@ MODEL (
   ),
   owner jen,
   grain id,
-  description 'List of waiter names'
-)
+  description 'List of waiter names',
+  audits (
+    assert_positive_id,
+    does_not_exceed_threshold(column := id, threshold := 200),
+    assert_valid_name,
+  )
+);
+
+AUDIT (
+  name does_not_exceed_threshold,
+);
+SELECT * FROM @this_model WHERE @column >= @threshold;
+
+AUDIT (
+  name assert_positive_id,
+);
+SELECT * FROM @this_model WHERE id < 0;
+
+AUDIT (
+  name assert_valid_name,
+);
+SELECT * FROM @this_model WHERE LENGTH(name) < 1;
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index bbd656fb..bb514e6f 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -28,7 +28,7 @@ from sqlmesh.core.model.kind import ModelKindName, SeedKind, create_model_kind
 from sqlmesh.core.model.meta import ModelMeta
 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
+from sqlmesh.utils import columns_to_types_all_known, str_to_bool, UniqueKeyDict
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_datetime, to_time_column
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, raise_config_error
 from sqlmesh.utils.hashing import hash_data
@@ -180,7 +180,12 @@ class _Model(ModelMeta, frozen=True):
                             value=exp.to_table(field_value, dialect=self.dialect),
                         )
                     )
-                elif field_name not in ("column_descriptions_", "default_catalog", "enabled"):
+                elif field_name not in (
+                    "column_descriptions_",
+                    "default_catalog",
+                    "enabled",
+                    "inline_audits",
+                ):
                     expressions.append(
                         exp.Property(
                             this=field_info.alias or field_name,
@@ -462,8 +467,11 @@ class _Model(ModelMeta, frozen=True):
         from sqlmesh.core.audit import BUILT_IN_AUDITS
 
         referenced_audits = []
+
         for audit_name, _ in self.audits:
-            if audit_name in 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(
@@ -762,13 +770,19 @@ class _Model(ModelMeta, frozen=True):
 
         for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):
             metadata.append(audit_name)
-
+            audit = None
             if audit_name in BUILT_IN_AUDITS:
                 for arg_name, arg_value in audit_args.items():
                     metadata.append(arg_name)
                     metadata.append(gen(arg_value))
+            elif audit_name in self.inline_audits:
+                audit = self.inline_audits[audit_name]
             elif audit_name in audits:
                 audit = audits[audit_name]
+            else:
+                raise SQLMeshError(f"Unexpected audit name '{audit_name}'.")
+
+            if audit:
                 query = (
                     audit.render_query(self, **t.cast(t.Dict[str, t.Any], audit_args))
                     or audit.query
@@ -781,8 +795,6 @@ class _Model(ModelMeta, frozen=True):
                         str(audit.blocking),
                     ]
                 )
-            else:
-                raise SQLMeshError(f"Unexpected audit name '{audit_name}'.")
 
         for key, value in (self.virtual_properties or {}).items():
             metadata.append(key)
@@ -1035,7 +1047,6 @@ class SqlModel(_SqlBasedModel):
 
     def validate_definition(self) -> None:
         query = self._query_renderer.render()
-
         if query is None:
             if self.depends_on_ is None:
                 raise_config_error(
@@ -1494,8 +1505,8 @@ def load_sql_based_model(
     rendered_meta = rendered_meta_exprs[0]
 
     # Extract the query and any pre/post statements
-    query_or_seed_insert, pre_statements, post_statements = _split_sql_model_statements(
-        expressions[1:], path
+    query_or_seed_insert, pre_statements, post_statements, inline_audits = (
+        _split_sql_model_statements(expressions[1:], path, dialect)
     )
 
     meta_fields: t.Dict[str, t.Any] = {
@@ -1552,6 +1563,7 @@ def load_sql_based_model(
         default_catalog=default_catalog,
         variables=variables,
         used_variables=used_variables,
+        inline_audits=inline_audits,
         **meta_fields,
     )
 
@@ -1847,6 +1859,7 @@ def _create_model(
     depends_on: t.Optional[t.Set[str]] = None,
     dialect: t.Optional[str] = None,
     physical_schema_override: t.Optional[t.Dict[str, str]] = None,
+    inline_audits: t.Optional[t.Dict[str, ModelAudit]] = None,
     **kwargs: t.Any,
 ) -> Model:
     _validate_model_fields(klass, {"name", *kwargs} - {"grain", "table_properties"}, path)
@@ -1867,6 +1880,7 @@ def _create_model(
     try:
         model = klass(
             name=name,
+            inline_audits=inline_audits,
             **{
                 **(defaults or {}),
                 "jinja_macros": jinja_macros or JinjaMacroRegistry(),
@@ -1892,36 +1906,59 @@ INSERT_SEED_MACRO_CALL = d.parse_one("@INSERT_SEED()")
 
 
 def _split_sql_model_statements(
-    expressions: t.List[exp.Expression], path: Path
-) -> t.Tuple[t.Optional[exp.Expression], t.List[exp.Expression], t.List[exp.Expression]]:
+    expressions: t.List[exp.Expression], path: Path, dialect: t.Optional[str] = None
+) -> t.Tuple[
+    t.Optional[exp.Expression],
+    t.List[exp.Expression],
+    t.List[exp.Expression],
+    UniqueKeyDict[str, ModelAudit],
+]:
     """Extracts the SELECT query from a sequence of expressions.
 
     Args:
         expressions: The list of all SQL statements in the model definition.
 
     Returns:
-        A tuple containing the extracted SELECT query or the `@INSERT_SEED()` call, the statements before the it, and
-        the statements after it.
+        A tuple containing the extracted SELECT query or the `@INSERT_SEED()` call, the statements before the it,
+        the statements after it, and the inline audit definitions.
 
     Raises:
         ConfigError: If the model definition contains more than one SELECT query or `@INSERT_SEED()` call.
     """
+    from sqlmesh.core.audit import ModelAudit, load_audit
+
     query_positions = []
-    for idx, expression in enumerate(expressions):
-        if (
-            isinstance(expression, (exp.Query, d.JinjaQuery))
-            or expression == INSERT_SEED_MACRO_CALL
-            or (isinstance(expression, d.MacroFunc) and expression.this.name.lower() == "union")
-        ):
-            query_positions.append((expression, idx))
+    sql_statements = []
+    inline_audits: UniqueKeyDict[str, ModelAudit] = UniqueKeyDict("inline_audits")
+
+    idx = 0
+    length = len(expressions)
+    while idx < length:
+        expr = expressions[idx]
+
+        if isinstance(expr, d.Audit):
+            loaded_audit = load_audit([expr, expressions[idx + 1]], dialect=dialect)
+            assert isinstance(loaded_audit, ModelAudit)
+            inline_audits[loaded_audit.name] = loaded_audit
+            idx += 2
+        else:
+            if (
+                isinstance(expr, (exp.Query, d.JinjaQuery))
+                or expr == INSERT_SEED_MACRO_CALL
+                or (isinstance(expr, d.MacroFunc) and expr.this.name.lower() == "union")
+            ):
+                query_positions.append((expr, idx))
+            sql_statements.append(expr)
+            idx += 1
 
     if not query_positions:
-        return None, expressions, []
+        return None, sql_statements, [], inline_audits
+
     elif len(query_positions) > 1:
         raise_config_error("Only one SELECT query is allowed per model", path)
 
     query, pos = query_positions[0]
-    return query, expressions[:pos], expressions[pos + 1 :]
+    return query, sql_statements[:pos], sql_statements[pos + 1 :], inline_audits
 
 
 def _resolve_session_properties(
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index 86a2e6e9..3ab76c79 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -67,6 +67,7 @@ class ModelMeta(_Node):
         default=None, alias="column_descriptions"
     )
     audits: t.List[AuditReference] = []
+    inline_audits: t.Dict[str, t.Any] = {}
     grains: t.List[exp.Expression] = []
     references: t.List[exp.Expression] = []
     physical_schema_override: t.Optional[str] = None
@@ -83,6 +84,24 @@ class ModelMeta(_Node):
     _default_catalog_validator = default_catalog_validator
     _depends_on_validator = depends_on_validator
 
+    @field_validator("inline_audits", mode="before")
+    @field_validator_v1_args
+    def _inline_audits_validator(cls, v: t.Any, values: t.Dict[str, t.Any]) -> t.Any:
+        if not isinstance(v, dict):
+            return {}
+
+        from sqlmesh.core.audit import ModelAudit
+
+        inline_audits = {}
+
+        for name, audit in v.items():
+            if isinstance(audit, ModelAudit):
+                inline_audits[name] = audit
+            elif isinstance(audit, dict):
+                inline_audits[name] = ModelAudit.parse_obj(audit)
+
+        return inline_audits
+
     @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]]:


Test patch:
diff --git a/tests/core/test_audit.py b/tests/core/test_audit.py
index 725a89a4..03877c03 100644
--- a/tests/core/test_audit.py
+++ b/tests/core/test_audit.py
@@ -2,6 +2,7 @@ import pytest
 from sqlglot import exp, parse_one
 
 from sqlmesh.core import constants as c
+from sqlmesh.core.context import Context
 from sqlmesh.core.audit import (
     BUILT_IN_AUDITS,
     ModelAudit,
@@ -11,7 +12,13 @@ from sqlmesh.core.audit import (
     load_multiple_audits,
 )
 from sqlmesh.core.dialect import parse
-from sqlmesh.core.model import IncrementalByTimeRangeKind, Model, create_sql_model
+from sqlmesh.core.model import (
+    IncrementalByTimeRangeKind,
+    Model,
+    SeedModel,
+    create_sql_model,
+    load_sql_based_model,
+)
 from sqlmesh.utils.errors import AuditConfigError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
 from sqlmesh.utils.metaprogramming import Executable
@@ -742,3 +749,48 @@ def test_variables(assert_exp_eq):
         audit.render_query(audit).sql(dialect="bigquery")
         == "SELECT * FROM `db`.`table` AS `table` WHERE `col` = 'test_val'"
     )
+
+
+def test_load_inline_audits(assert_exp_eq):
+    expressions = parse(
+        """
+        MODEL (
+            name db.table,
+            dialect spark,
+            audits(does_not_exceed_threshold)
+        );
+
+        SELECT id FROM tbl;
+
+        AUDIT (
+        name does_not_exceed_threshold,
+        );
+        SELECT * FROM @this_model
+        WHERE @column >= @threshold;
+
+        AUDIT (
+        name assert_positive_id,
+        );
+        SELECT *
+        FROM @this_model
+        WHERE
+        id < 0;
+    """
+    )
+
+    model = load_sql_based_model(expressions)
+    assert len(model.audits) == 1
+    assert len(model.inline_audits) == 2
+    assert isinstance(model.inline_audits["assert_positive_id"], ModelAudit)
+    assert isinstance(model.inline_audits["does_not_exceed_threshold"], ModelAudit)
+
+
+def test_model_inline_audits(sushi_context: Context):
+    model_name = "sushi.waiter_names"
+    expected_query = 'SELECT * FROM (SELECT * FROM "memory"."sushi"."waiter_names" AS "waiter_names") AS "_q_0" WHERE "id" < 0'
+    model = sushi_context.get_snapshot(model_name, raise_if_missing=True).node
+
+    assert isinstance(model, SeedModel)
+    assert len(model.inline_audits) == 3
+    assert isinstance(model.inline_audits["assert_valid_name"], ModelAudit)
+    assert model.inline_audits["assert_positive_id"].render_query(model).sql() == expected_query
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 85faa75f..050f3662 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -114,6 +114,7 @@ def test_json(snapshot: Snapshot):
                 "dialect": "spark",
             },
             "mapping_schema": {},
+            "inline_audits": {},
             "start": "2020-01-01",
             "dialect": "spark",
             "name": "name",
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index 5384dc74..e30723dd 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -101,6 +101,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
                         "dialect": "spark",
                     },
                     "mapping_schema": {},
+                    "inline_audits": {},
                     "name": "test_model",
                     "partitioned_by": ["`a`"],
                     "query": "SELECT a, ds FROM tbl",