Repository: TobikoData/sqlmesh
Pull request #2904 (2024-07-15 18:29:11)

Feat: Support the strict mode when creating external models
The strict mode ensures that the process fails if a referenced external table is missing in the target database.

Before this update, SQLMesh would emit a warning but ignore the missing table otherwise.


Patch:
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 4bc28920..3d0ed10a 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -702,12 +702,17 @@ def rollback(obj: Context) -> None:
 
 
 @cli.command("create_external_models")
+@click.option(
+    "--strict",
+    is_flag=True,
+    help="Raise an error if the external model is missing in the database",
+)
 @click.pass_obj
 @error_handler
 @cli_analytics
-def create_external_models(obj: Context) -> None:
+def create_external_models(obj: Context, **kwargs: t.Any) -> None:
     """Create a schema file containing external model schemas."""
-    obj.create_external_models()
+    obj.create_external_models(**kwargs)
 
 
 @cli.command("table_diff")
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 14783d00..06b86506 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1655,11 +1655,14 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._new_state_sync().rollback()
 
     @python_api_analytics
-    def create_external_models(self) -> None:
+    def create_external_models(self, strict: bool = False) -> None:
         """Create a file to document the schema of external models.
 
         The external models file contains all columns and types of external models, allowing for more
         robust lineage, validation, and optimizations.
+
+        Args:
+            strict: If True, raise an error if the external model is missing in the database.
         """
         if not self._models:
             self.load(update_schemas=False)
@@ -1685,6 +1688,7 @@ class GenericContext(BaseContext, t.Generic[C]):
                 dialect=config.model_defaults.dialect,
                 gateway=self.gateway,
                 max_workers=self.concurrent_tasks,
+                strict=strict,
             )
 
     @python_api_analytics
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 33bd240b..19407487 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -241,7 +241,7 @@ def _parse_macro(self: Parser, keyword_macro: str = "") -> t.Optional[exp.Expres
             return field
         return self.expression(MacroVar, this=field.this, comments=comments)
 
-    if isinstance(field, exp.Window):
+    if isinstance(field, (exp.Window, exp.IgnoreNulls, exp.RespectNulls)):
         field.set("this", _build_macro(field.this))
     else:
         field = _build_macro(field)
diff --git a/sqlmesh/core/schema_loader.py b/sqlmesh/core/schema_loader.py
index c4704050..1b00c7bc 100644
--- a/sqlmesh/core/schema_loader.py
+++ b/sqlmesh/core/schema_loader.py
@@ -12,6 +12,7 @@ from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model.definition import Model
 from sqlmesh.core.state_sync import StateReader
 from sqlmesh.utils import UniqueKeyDict, yaml
+from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
 
@@ -24,6 +25,7 @@ def create_external_models_file(
     dialect: DialectType,
     gateway: t.Optional[str] = None,
     max_workers: int = 1,
+    strict: bool = False,
 ) -> None:
     """Create or replace a YAML file with column and types of all columns in all external models.
 
@@ -35,6 +37,7 @@ def create_external_models_file(
         dialect: The dialect to serialize the schema as.
         gateway: If the model should be associated with a specific gateway; the gateway key
         max_workers: The max concurrent workers to fetch columns.
+        strict: If True, raise an error if the external model is missing in the database.
     """
     external_model_fqns = set()
 
@@ -61,7 +64,10 @@ def create_external_models_file(
             try:
                 return adapter.columns(table, include_pseudo_columns=True)
             except Exception as e:
-                logger.warning(f"Unable to get schema for '{table}': '{e}'.")
+                msg = f"Unable to get schema for '{table}': '{e}'."
+                if strict:
+                    raise SQLMeshError(msg) from e
+                logger.warning(msg)
                 return None
 
         gateway_part = {"gateway": gateway} if gateway else {}
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
index e1c6bd66..3c4a844d 100644
--- a/sqlmesh/core/selector.py
+++ b/sqlmesh/core/selector.py
@@ -158,8 +158,9 @@ class Selector:
                     add_sub_results(
                         self._expand_model_tag(sub_selection[4:], models, models_by_tag)
                     )
-                elif sub_selection.startswith("git:"):
-                    add_sub_results(self._expand_git(sub_selection[4:]))
+                elif sub_selection.startswith(("git:", "+git:")):
+                    sub_selection = sub_selection.replace("git:", "")
+                    add_sub_results(self._expand_git(sub_selection, models))
                 else:
                     add_sub_results(self._expand_model_name(sub_selection, models))
 
@@ -170,7 +171,15 @@ class Selector:
 
         return results
 
-    def _expand_git(self, target_branch: str) -> t.Set[str]:
+    def _expand_git(self, target_branch: str, models: t.Dict[str, Model]) -> t.Set[str]:
+        results: t.Set[str] = set()
+
+        (
+            target_branch,
+            include_upstream,
+            include_downstream,
+        ) = self._get_value_and_dependency_inclusion(target_branch)
+
         git_modified_files = {
             *self._git_client.list_untracked_files(),
             *self._git_client.list_uncommitted_changed_files(),
@@ -180,8 +189,14 @@ class Selector:
 
         if not matched_models:
             logger.warning(f"Expression 'git:{target_branch}' doesn't match any models.")
+            return matched_models
 
-        return matched_models
+        for model_fqn in matched_models:
+            results.update(
+                self._get_models(model_fqn, include_upstream, include_downstream, models)
+            )
+
+        return results
 
     def _expand_model_name(self, selection: str, models: t.Dict[str, Model]) -> t.Set[str]:
         results = set()
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index e1c1c801..51652fc0 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -558,11 +558,17 @@ class SQLMeshMagics(Magics):
         context.console.log_success("Migration complete")
 
     @magic_arguments()
+    @argument(
+        "--strict",
+        action="store_true",
+        help="Raise an error if the external model is missing in the database",
+    )
     @line_magic
     @pass_sqlmesh_context
     def create_external_models(self, context: Context, line: str) -> None:
         """Create a schema file containing external model schemas."""
-        context.create_external_models()
+        args = parse_argstring(self.create_external_models, line)
+        context.create_external_models(strict=args.strict)
 
     @magic_arguments()
     @argument(


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 72bde9ce..6b1fba20 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -3,7 +3,7 @@ import typing as t
 import pytest
 from sqlglot import MappingSchema, ParseError, exp, parse_one
 
-from sqlmesh.core import constants as c
+from sqlmesh.core import constants as c, dialect as d
 from sqlmesh.core.dialect import StagedFilePath
 from sqlmesh.core.macros import SQL, MacroEvalError, MacroEvaluator, macro
 from sqlmesh.utils.errors import SQLMeshError
@@ -638,3 +638,24 @@ def test_macro_metadata_flag():
 
     assert getattr(noop_metadata_only, c.SQLMESH_METADATA) is True
     assert macro.get_registry()["noop_metadata_only"].metadata_only is True
+
+
+def test_macro_first_value_ignore_respect_nulls(assert_exp_eq) -> None:
+    schema = MappingSchema({}, dialect="duckdb")
+    evaluator = MacroEvaluator(schema=schema, dialect="duckdb")
+
+    evaluator.evaluate(t.cast(d.MacroDef, d.parse_one("@DEF(test, x -> x)")))
+
+    expected_sql = "SELECT FIRST_VALUE(x IGNORE NULLS) OVER (ORDER BY y NULLS FIRST) AS column_test"
+    actual_expr = d.parse_one(
+        "SELECT FIRST_VALUE(@test(x) IGNORE NULLS) OVER (ORDER BY y) AS column_test"
+    )
+    assert_exp_eq(evaluator.transform(actual_expr), expected_sql, dialect="duckdb")
+
+    expected_sql = (
+        "SELECT FIRST_VALUE(x RESPECT NULLS) OVER (ORDER BY y NULLS FIRST) AS column_test"
+    )
+    actual_expr = d.parse_one(
+        "SELECT FIRST_VALUE(@test(x) RESPECT NULLS) OVER (ORDER BY y) AS column_test"
+    )
+    assert_exp_eq(evaluator.transform(actual_expr), expected_sql, dialect="duckdb")
diff --git a/tests/core/test_schema_loader.py b/tests/core/test_schema_loader.py
index 2e970914..ada27dfd 100644
--- a/tests/core/test_schema_loader.py
+++ b/tests/core/test_schema_loader.py
@@ -1,3 +1,4 @@
+import pytest
 import logging
 import typing as t
 from pathlib import Path
@@ -16,6 +17,7 @@ from sqlmesh.core.model.definition import ExternalModel
 from sqlmesh.core.schema_loader import create_external_models_file
 from sqlmesh.core.snapshot import SnapshotChangeCategory
 from sqlmesh.utils.yaml import YAML
+from sqlmesh.utils.errors import SQLMeshError
 
 
 def test_create_external_models(tmpdir, assert_exp_eq):
@@ -354,3 +356,13 @@ def test_missing_table(tmp_path: Path):
     with open(filename, "r", encoding="utf8") as fd:
         schema = YAML().load(fd)
     assert len(schema) == 0
+
+    with pytest.raises(SQLMeshError, match=r"""Unable to get schema for '"tbl_source"'.*"""):
+        create_external_models_file(
+            filename,
+            {"a": model},  # type: ignore
+            context.engine_adapter,
+            context.state_reader,
+            "",
+            strict=True,
+        )
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index da8e2e61..f43e918b 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -392,6 +392,9 @@ def test_model_selection_normalized(mocker: MockerFixture, make_snapshot):
     [
         (["git:main"], {'"test_model_a"', '"test_model_c"'}),
         (["git:main & +*model_c"], {'"test_model_c"'}),
+        (["git:main+"], {'"test_model_a"', '"test_model_c"', '"test_model_d"'}),
+        (["+git:main"], {'"test_model_a"', '"test_model_c"', '"test_model_b"'}),
+        (["+git:main+"], {'"test_model_a"', '"test_model_c"', '"test_model_b"', '"test_model_d"'}),
     ],
 )
 def test_expand_git_selection(
@@ -407,7 +410,11 @@ def test_expand_git_selection(
     model_b._path = Path("/path/to/test_model_b.sql")
     models[model_b.fqn] = model_b
 
-    model_c = SqlModel(name="test_model_c", query=d.parse_one("SELECT 3 AS c"))
+    model_c = SqlModel(
+        name="test_model_c",
+        query=d.parse_one("SELECT b AS c FROM test_model_b"),
+        depends_on={"test_model_b"},
+    )
     model_c._path = Path("/path/to/test_model_c.sql")
     models[model_c.fqn] = model_c