Repository: TobikoData/sqlmesh
Pull request #2980 (2024-08-02 18:43:08)

Fix: ensure Context.close doesn't unnecessarily create a state sync db
@treysp noticed that `sqlmesh info` triggers all migrations in a fresh project (e.g. `sqlmesh init duckdb`), which is weird. This happens because the CLI calls `close()` on the instantiated context and in turn that results in doing `self.state_sync` within that method, which creates a new state sync and populates it. We only need to do that if `_state_sync` actually exists.


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 85de1811..19833e24 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -11,7 +11,7 @@ For more information regarding what a context can do, see `sqlmesh.core.context.
 Creating and applying a plan against the staging environment.
 ```python
 from sqlmesh.core.context import Context
-context = Context(path="example", config="local_config")
+context = Context(paths="example", config="local_config")
 plan = context.plan("staging")
 context.apply(plan)
 ```
@@ -19,14 +19,14 @@ context.apply(plan)
 Running audits on your data.
 ```python
 from sqlmesh.core.context import Context
-context = Context(path="example", config="local_config")
+context = Context(paths="example", config="local_config")
 context.audit("yesterday", "now")
 ```
 
 Running tests on your models.
 ```python
 from sqlmesh.core.context import Context
-context = Context(path="example")
+context = Context(paths="example")
 context.test()
 ```
 """
@@ -1720,8 +1720,10 @@ class GenericContext(BaseContext, t.Generic[C]):
 
     def close(self) -> None:
         """Releases all resources allocated by this context."""
-        self.snapshot_evaluator.close()
-        self.state_sync.close()
+        if self._snapshot_evaluator:
+            self._snapshot_evaluator.close()
+        if self._state_sync:
+            self._state_sync.close()
 
     def _run(
         self,
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
index 4a793980..a77e5b4e 100644
--- a/sqlmesh/core/loader.py
+++ b/sqlmesh/core/loader.py
@@ -123,8 +123,9 @@ class Loader(abc.ABC):
         self._config_mtimes = {path: max(mtimes) for path, mtimes in config_mtimes.items()}
 
         macros, jinja_macros = self._load_scripts()
+        audits = self._load_audits(macros=macros, jinja_macros=jinja_macros)
         models = self._load_models(
-            macros, jinja_macros, context.gateway or context.config.default_gateway
+            macros, jinja_macros, context.gateway or context.config.default_gateway, audits or None
         )
 
         for model in models.values():
@@ -146,7 +147,7 @@ class Loader(abc.ABC):
             macros=macros,
             jinja_macros=jinja_macros,
             models=models,
-            audits=self._load_audits(macros=macros, jinja_macros=jinja_macros),
+            audits=audits,
             metrics=expand_metrics(metrics),
             dag=self._dag,
         )
@@ -171,7 +172,11 @@ class Loader(abc.ABC):
 
     @abc.abstractmethod
     def _load_models(
-        self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry, gateway: t.Optional[str]
+        self,
+        macros: MacroRegistry,
+        jinja_macros: JinjaMacroRegistry,
+        gateway: t.Optional[str],
+        audits: t.Optional[t.Dict[str, Audit]],
     ) -> UniqueKeyDict[str, Model]:
         """Loads all models."""
 
@@ -312,20 +317,27 @@ class SqlMeshLoader(Loader):
         return macros, jinja_macros
 
     def _load_models(
-        self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry, gateway: t.Optional[str]
+        self,
+        macros: MacroRegistry,
+        jinja_macros: JinjaMacroRegistry,
+        gateway: t.Optional[str],
+        audits: t.Optional[t.Dict[str, Audit]] = None,
     ) -> UniqueKeyDict[str, Model]:
         """
         Loads all of the models within the model directory with their associated
         audits into a Dict and creates the dag
         """
-        models = self._load_sql_models(macros, jinja_macros)
+        models = self._load_sql_models(macros, jinja_macros, audits)
         models.update(self._load_external_models(gateway))
         models.update(self._load_python_models())
 
         return models
 
     def _load_sql_models(
-        self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry
+        self,
+        macros: MacroRegistry,
+        jinja_macros: JinjaMacroRegistry,
+        audits: t.Optional[t.Dict[str, Audit]] = None,
     ) -> UniqueKeyDict[str, Model]:
         """Loads the sql models into a Dict"""
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
@@ -357,6 +369,8 @@ class SqlMeshLoader(Loader):
                         defaults=config.model_defaults.dict(),
                         macros=macros,
                         jinja_macros=jinja_macros,
+                        audits=audits,
+                        default_audits=config.model_defaults.audits,
                         path=Path(path).absolute(),
                         module_path=context_path,
                         dialect=config.model_defaults.dialect,
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 6fef5baf..5d9e6037 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -20,6 +20,7 @@ from sqlglot.optimizer.simplify import gen
 from sqlglot.schema import MappingSchema, nested_set
 from sqlglot.time import format_time
 
+from sqlmesh.core.dialect import extract_audit
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
 from sqlmesh.core.macros import MacroRegistry, MacroStrTemplate, macro
@@ -47,7 +48,7 @@ from sqlmesh.utils.metaprogramming import (
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
-    from sqlmesh.core.audit import ModelAudit
+    from sqlmesh.core.audit import ModelAudit, Audit
     from sqlmesh.core.context import ExecutionContext
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.engine_adapter._typing import QueryOrDF
@@ -1518,6 +1519,8 @@ def load_sql_based_model(
     time_column_format: str = c.DEFAULT_TIME_COLUMN_FORMAT,
     macros: t.Optional[MacroRegistry] = None,
     jinja_macros: t.Optional[JinjaMacroRegistry] = None,
+    audits: t.Optional[t.Dict[str, Audit]] = None,
+    default_audits: t.List[AuditReference] = [],
     python_env: t.Optional[t.Dict[str, Executable]] = None,
     dialect: t.Optional[str] = None,
     physical_schema_override: t.Optional[t.Dict[str, str]] = None,
@@ -1557,10 +1560,13 @@ def load_sql_based_model(
         )
 
     unrendered_signals = None
+    model_audits = None
     for prop in meta.expressions:
         if prop.name.lower() == "signals":
             unrendered_signals = prop.args.get("value")
 
+        if prop.name.lower() == "audits":
+            model_audits = prop.args.get("value")
     meta_python_env = _python_env(
         expressions=meta,
         jinja_macro_references=None,
@@ -1637,6 +1643,9 @@ def load_sql_based_model(
     common_kwargs = dict(
         pre_statements=pre_statements,
         post_statements=post_statements,
+        audit_expressions=_extract_audit_expressions(
+            audits, inline_audits, model_audits, default_audits
+        ),
         defaults=defaults,
         path=path,
         module_path=module_path,
@@ -1688,6 +1697,7 @@ def create_sql_model(
     *,
     pre_statements: t.Optional[t.List[exp.Expression]] = None,
     post_statements: t.Optional[t.List[exp.Expression]] = None,
+    audit_expressions: t.Optional[t.List[exp.Expression]] = None,
     defaults: t.Optional[t.Dict[str, t.Any]] = None,
     path: Path = Path(),
     module_path: Path = Path(),
@@ -1736,10 +1746,11 @@ def create_sql_model(
 
     pre_statements = pre_statements or []
     post_statements = post_statements or []
+    audit_expressions = audit_expressions or []
 
     if not python_env:
         python_env = _python_env(
-            [*pre_statements, query, *post_statements],
+            [*pre_statements, query, *post_statements, *audit_expressions],
             jinja_macro_references,
             module_path,
             macros or macro.get_registry(),
@@ -1774,6 +1785,7 @@ def create_seed_model(
     dialect: t.Optional[str] = None,
     pre_statements: t.Optional[t.List[exp.Expression]] = None,
     post_statements: t.Optional[t.List[exp.Expression]] = None,
+    audit_expressions: t.Optional[t.List[exp.Expression]] = None,
     defaults: t.Optional[t.Dict[str, t.Any]] = None,
     path: Path = Path(),
     module_path: Path = Path(),
@@ -1818,10 +1830,11 @@ def create_seed_model(
 
     pre_statements = pre_statements or []
     post_statements = post_statements or []
+    audit_expressions = audit_expressions or []
 
     if not python_env:
         python_env = _python_env(
-            [*pre_statements, *post_statements],
+            [*pre_statements, *post_statements, *audit_expressions],
             jinja_macro_references,
             module_path,
             macros or macro.get_registry(),
@@ -2236,6 +2249,38 @@ def _list_of_calls_to_exp(value: t.List[t.Tuple[str, t.Dict[str, t.Any]]]) -> ex
     )
 
 
+def _extract_audit_expressions(
+    audits: t.Optional[dict[str, Audit]] = None,
+    inline_audits: t.Optional[dict[str, Audit]] = None,
+    model_audits: t.Optional[t.Any] = None,
+    default_audits: t.Optional[t.List[AuditReference]] = None,
+) -> list:
+    audit_names = []
+    if model_audits:
+        if isinstance(model_audits, (exp.Tuple, exp.Array)):
+            audit_names = [extract_audit(i)[0] for i in model_audits.expressions]
+        elif isinstance(model_audits, exp.Paren):
+            audit_names = [extract_audit(model_audits.this)[0]]
+        elif isinstance(model_audits, exp.Expression):
+            audit_names = [extract_audit(model_audits)[0]]
+
+    if default_audits:
+        for audit_name, _ in default_audits:
+            audit_names.append(audit_name)
+
+    audit_expressions = []
+    if audit_names and audits:
+        for audit_name in audit_names:
+            audit = audits.get(audit_name)
+            if audit:
+                audit_expressions.append(audit.query)
+    if inline_audits:
+        for audit_name in inline_audits:
+            audit_expressions.append(inline_audits[audit_name].query)
+
+    return audit_expressions
+
+
 def _is_projection(expr: exp.Expression) -> bool:
     parent = expr.parent
     return isinstance(parent, exp.Select) and expr.arg_key == "expressions"
diff --git a/sqlmesh/dbt/loader.py b/sqlmesh/dbt/loader.py
index 90db2b7d..6ef0dca7 100644
--- a/sqlmesh/dbt/loader.py
+++ b/sqlmesh/dbt/loader.py
@@ -89,7 +89,11 @@ class DbtLoader(Loader):
         )
 
     def _load_models(
-        self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry, gateway: t.Optional[str]
+        self,
+        macros: MacroRegistry,
+        jinja_macros: JinjaMacroRegistry,
+        gateway: t.Optional[str],
+        audits: t.Optional[t.Dict[str, Audit]],
     ) -> UniqueKeyDict[str, Model]:
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
 


Test patch:
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index 769518f6..309727c0 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -9,6 +9,7 @@ from freezegun import freeze_time
 
 from sqlmesh.cli.example_project import init_example_project
 from sqlmesh.cli.main import cli
+from sqlmesh.core.context import Context
 
 FREEZE_TIME = "2023-01-01 00:00:00"
 
@@ -642,3 +643,20 @@ def test_table_name(runner, tmp_path):
         )
     assert result.exit_code == 0
     assert result.output.startswith("db.sqlmesh__sqlmesh_example.sqlmesh_example__full_model__")
+
+
+def test_info_on_new_project_does_not_create_state_sync(runner, tmp_path):
+    create_example_project(tmp_path)
+
+    # Invoke the info command
+    result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "info"])
+    assert result.exit_code == 0
+
+    context = Context(paths=tmp_path)
+
+    # Confirm that the state sync tables haven't been created
+    assert not context.engine_adapter.table_exists("sqlmesh._snapshots")
+    assert not context.engine_adapter.table_exists("sqlmesh._environments")
+    assert not context.engine_adapter.table_exists("sqlmesh._intervals")
+    assert not context.engine_adapter.table_exists("sqlmesh._plan_dags")
+    assert not context.engine_adapter.table_exists("sqlmesh._versions")
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 1a186cf2..7d582727 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -15,7 +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.audit import ModelAudit, load_audit
 from sqlmesh.core.config import (
     Config,
     NameInferenceConfig,
@@ -4541,6 +4541,88 @@ def test_macros_in_model_statement(sushi_context, assert_exp_eq):
     assert model.session_properties == {"foo": exp.column("bar_baz", quoted=False)}
 
 
+def test_macro_references_in_audits():
+    @macro()
+    def zero_value(evaluator: MacroEvaluator) -> int:
+        return 0
+
+    @macro()
+    def min_value(evaluator: MacroEvaluator) -> int:
+        return 1
+
+    @macro()
+    def not_loaded_macro(evaluator: MacroEvaluator) -> int:
+        return 10
+
+    @macro()
+    def max_value(evaluator: MacroEvaluator) -> int:
+        return 1000
+
+    audit_expression = parse(
+        """
+        AUDIT (
+    name assert_max_value,
+    );
+    SELECT *
+    FROM @this_model
+    WHERE
+    id > @max_value;
+    """
+    )
+
+    not_zero_audit = parse(
+        """
+        AUDIT (
+    name assert_not_zero,
+    );
+    SELECT *
+    FROM @this_model
+    WHERE
+    id = @zero_value;
+    """
+    )
+
+    model_expression = d.parse(
+        """
+        MODEL (
+            name db.audit_model,
+            audits (assert_max_value, assert_positive_ids),
+        );
+        SELECT 1 as id;
+
+        AUDIT (
+    name assert_positive_ids,
+    );
+    SELECT *
+    FROM @this_model
+    WHERE
+    id < @min_value;
+    """
+    )
+
+    audits = {
+        "assert_max_value": load_audit(audit_expression, dialect="duckdb"),
+        "assert_not_zero": load_audit(not_zero_audit, dialect="duckdb"),
+    }
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb", audits=["assert_not_zero"])
+    )
+    model = load_sql_based_model(
+        model_expression, audits=audits, default_audits=config.model_defaults.audits
+    )
+
+    assert len(model.audits) == 2
+    assert len(model.inline_audits) == 1
+    assert len(model.python_env) == 3
+    assert config.model_defaults.audits == [("assert_not_zero", {})]
+    assert model.audits == [("assert_max_value", {}), ("assert_positive_ids", {})]
+    assert isinstance(model.inline_audits["assert_positive_ids"], ModelAudit)
+    assert isinstance(model.python_env["min_value"], Executable)
+    assert isinstance(model.python_env["max_value"], Executable)
+    assert isinstance(model.python_env["zero_value"], Executable)
+    assert "not_loaded_macro" not in model.python_env
+
+
 def test_python_model_dialect():
     model._dialect = "snowflake"