Repository: TobikoData/sqlmesh
Pull request #2454 (2024-04-16 04:51:01)

Chore: Move plan evaluator to private method


Patch:
diff --git a/.circleci/continue_config.yml b/.circleci/continue_config.yml
index 58ce8f55..a8a416af 100644
--- a/.circleci/continue_config.yml
+++ b/.circleci/continue_config.yml
@@ -107,6 +107,20 @@ jobs:
           name: Run slow tests
           command: make cicd-test
 
+  migration_test:
+    docker:
+      - image: cimg/python:3.10
+    resource_class: small
+    steps:
+      - halt_unless_core
+      - checkout
+      - run:
+          name: Install dependencies
+          command: make install-dev
+      - run:
+          name: Run the migration test
+          command: ./.circleci/test_migration.sh
+
   ui_style:
     docker:
       - image: cimg/python:3.8
@@ -207,9 +221,9 @@ jobs:
           name: Run tests
           command: make engine-docker-test
           no_output_timeout: 15m
-  
+
   trigger_private_tests:
-      docker: 
+      docker:
         - image: cimg/base:2021.11
       resource_class: small
       steps:
@@ -226,8 +240,8 @@ jobs:
                 --data '{
                   "branch":"main",
                   "parameters":{
-                    "run_main_pr":false, 
-                    "run_sqlmesh_commit":true, 
+                    "run_main_pr":false,
+                    "run_sqlmesh_commit":true,
                     "sqlmesh_branch":"'$CIRCLE_BRANCH'",
                     "sqlmesh_commit_author":"'$CIRCLE_USERNAME'",
                     "sqlmesh_commit_hash":"'$CIRCLE_SHA1'",
@@ -269,3 +283,4 @@ workflows:
                 - main
       - ui_style
       - ui_test
+      - migration_test
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index a7681619..8a9a3fe5 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -261,7 +261,8 @@ def format(ctx: click.Context, **kwargs: t.Any) -> None:
 @error_handler
 def diff(ctx: click.Context, environment: t.Optional[str] = None) -> None:
     """Show the diff between the local state and the target environment."""
-    ctx.obj.diff(environment)
+    if ctx.obj.diff(environment, detailed=True):
+        exit(1)
 
 
 @cli.command("plan")
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 9dc4a723..177b0406 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1110,9 +1110,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             NotificationEvent.APPLY_START, environment=plan.environment_naming_info.name
         )
         try:
-            self._scheduler.create_plan_evaluator(self).evaluate(
-                plan, circuit_breaker=circuit_breaker
-            )
+            self._apply(plan, circuit_breaker)
         except Exception as e:
             self.notification_target_manager.notify(
                 NotificationEvent.APPLY_FAILURE, traceback.format_exc()
@@ -1138,17 +1136,21 @@ class GenericContext(BaseContext, t.Generic[C]):
         else:
             self.console.log_success(f"Environment '{name}' has been invalidated.")
 
-    def diff(self, environment: t.Optional[str] = None, detailed: bool = False) -> None:
+    def diff(self, environment: t.Optional[str] = None, detailed: bool = False) -> bool:
         """Show a diff of the current context with a given environment.
 
         Args:
             environment: The environment to diff against.
             detailed: Show the actual SQL differences if True.
+
+        Returns:
+            True if there are changes, False otherwise.
         """
         environment = environment or self.config.default_target_environment
         environment = Environment.normalize_name(environment)
+        context_diff = self._context_diff(environment)
         self.console.show_model_difference_summary(
-            self._context_diff(environment),
+            context_diff,
             EnvironmentNamingInfo.from_environment_catalog_mapping(
                 self.config.environment_catalog_mapping,
                 name=environment,
@@ -1157,6 +1159,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             self.default_catalog,
             no_diff=not detailed,
         )
+        return context_diff.has_changes
 
     def table_diff(
         self,
@@ -1622,6 +1625,9 @@ class GenericContext(BaseContext, t.Generic[C]):
 
         return success
 
+    def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
+        self._scheduler.create_plan_evaluator(self).evaluate(plan, circuit_breaker=circuit_breaker)
+
     def table_name(self, model_name: str, dev: bool) -> str:
         """Returns the name of the pysical table for the given model name.
 
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index c166df20..b709c256 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -124,6 +124,16 @@ class SnapshotId(PydanticModel, frozen=True):
         """Helper method to return self."""
         return self
 
+    def __eq__(self, other: t.Any) -> bool:
+        return (
+            isinstance(other, self.__class__)
+            and self.name == other.name
+            and self.identifier == other.identifier
+        )
+
+    def __hash__(self) -> int:
+        return hash((self.__class__, self.name, self.identifier))
+
     def __lt__(self, other: SnapshotId) -> bool:
         return self.name < other.name
 
diff --git a/sqlmesh/utils/__init__.py b/sqlmesh/utils/__init__.py
index 8e6eb04e..d57a1a4c 100644
--- a/sqlmesh/utils/__init__.py
+++ b/sqlmesh/utils/__init__.py
@@ -316,22 +316,21 @@ def columns_to_types_to_struct(
     )
 
 
-def columns_to_types_all_known(columns_to_types: t.Dict[str, exp.DataType]) -> bool:
-    """
-    Checks that all column types are known and not NULL.
-    """
-
-    def check_expression_unknown(d_type: t.Union[exp.DataType, exp.ColumnDef]) -> bool:
-        if isinstance(d_type, exp.ColumnDef):
-            if not d_type.kind:
-                return False
-            d_type = d_type.kind
-        if isinstance(d_type, exp.DataTypeParam):
-            return True
-        if d_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL):
+def type_is_known(d_type: t.Union[exp.DataType, exp.ColumnDef]) -> bool:
+    """Checks that a given column type is known and not NULL."""
+    if isinstance(d_type, exp.ColumnDef):
+        if not d_type.kind:
             return False
-        if d_type.expressions:
-            return all(check_expression_unknown(expression) for expression in d_type.expressions)
+        d_type = d_type.kind
+    if isinstance(d_type, exp.DataTypeParam):
         return True
+    if d_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL):
+        return False
+    if d_type.expressions:
+        return all(type_is_known(expression) for expression in d_type.expressions)
+    return True
+
 
-    return all(check_expression_unknown(expression) for expression in columns_to_types.values())
+def columns_to_types_all_known(columns_to_types: t.Dict[str, exp.DataType]) -> bool:
+    """Checks that all column types are known and not NULL."""
+    return all(type_is_known(expression) for expression in columns_to_types.values())


Test patch:
diff --git a/.circleci/test_migration.sh b/.circleci/test_migration.sh
new file mode 100755
index 00000000..eb2ed9f2
--- /dev/null
+++ b/.circleci/test_migration.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -ex
+
+CONFIG_NAME="local_config"
+TMP_DIR=$(mktemp -d)
+SUSHI_DIR="$TMP_DIR/sushi"
+
+
+if [[ -z $(git tag --points-at HEAD) ]]; then
+    # If the current commit is not tagged, we need to find the last tag
+    LAST_TAG=$(git describe --tags --abbrev=0)
+else
+    # If the current commit is tagged, we need to find the previous tag
+    LAST_TAG=$(git tag --sort=-creatordate | head -n 2 | tail -n 1)
+fi
+
+git checkout $LAST_TAG
+
+cp -r ./examples/sushi $TMP_DIR
+
+# Run initial plan
+pushd $SUSHI_DIR
+rm -rf ./data/*
+sqlmesh --config $CONFIG_NAME plan --no-prompts --auto-apply
+popd
+
+# Switch back to the starting state of the repository
+git checkout -
+
+# Migrate and make sure the diff is empty
+pushd $SUSHI_DIR
+sqlmesh --config $CONFIG_NAME migrate
+sqlmesh --config $CONFIG_NAME diff prod
+popd
+
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index d6443023..7583a3a0 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -18,7 +18,7 @@ from sqlmesh.core import constants as c
 from sqlmesh.core.dialect import normalize_model_name, schema_
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.model import Model, PythonModel, SqlModel
-from sqlmesh.utils import UniqueKeyDict, random_id, yaml
+from sqlmesh.utils import UniqueKeyDict, random_id, type_is_known, yaml
 from sqlmesh.utils.date import pandas_timestamp_to_pydatetime
 from sqlmesh.utils.errors import ConfigError, TestError
 
@@ -94,20 +94,39 @@ class ModelTest(unittest.TestCase):
     def setUp(self) -> None:
         """Load all input tables"""
         for name, values in self.body.get("inputs", {}).items():
-            # Types specified in the test take precedence over the corresponding inferred ones
+            all_types_are_known = False
+            known_columns_to_types: t.Dict[str, exp.DataType] = {}
+
             model = self.models.get(name)
-            columns_to_types = {
-                **(model.columns_to_types if model and model.annotated else {}),
-                **values.get("columns", {}),
-            }
+            if model:
+                inferred_columns_to_types = model.columns_to_types or {}
+                known_columns_to_types = {
+                    c: t for c, t in inferred_columns_to_types.items() if type_is_known(t)
+                }
+                all_types_are_known = bool(inferred_columns_to_types) and (
+                    len(known_columns_to_types) == len(inferred_columns_to_types)
+                )
+
+            # Types specified in the test take precedence over the corresponding inferred ones
+            known_columns_to_types.update(values.get("columns", {}))
 
             rows = values["rows"]
-            if not columns_to_types and rows:
-                for i, v in rows[0].items():
-                    v_type = annotate_types(exp.convert(v)).type or type(v).__name__
-                    columns_to_types[i] = exp.maybe_parse(
-                        v_type, into=exp.DataType, dialect=self.dialect
-                    )
+            if not all_types_are_known and rows:
+                for col, value in rows[0].items():
+                    if col not in known_columns_to_types:
+                        v_type = annotate_types(exp.convert(value)).type or type(value).__name__
+                        v_type = exp.maybe_parse(v_type, into=exp.DataType, dialect=self.dialect)
+
+                        if not type_is_known(v_type):
+                            _raise_error(
+                                f"Failed to infer the data type of column '{col}' for '{name}'. This issue can be "
+                                "mitigated by casting the column in the model definition, setting its type in "
+                                "schema.yaml if it's an external model, setting the model's 'columns' property, "
+                                "or setting its 'columns' mapping in the test itself",
+                                self.path,
+                            )
+
+                        known_columns_to_types[col] = v_type
 
             test_fixture_table = self._test_fixture_table(name)
             if test_fixture_table.db:
@@ -115,8 +134,8 @@ class ModelTest(unittest.TestCase):
                     schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
                 )
 
-            df = _create_df(rows, columns=columns_to_types)
-            self.engine_adapter.create_view(test_fixture_table, df, columns_to_types)
+            df = _create_df(rows, columns=known_columns_to_types)
+            self.engine_adapter.create_view(test_fixture_table, df, known_columns_to_types)
 
     def tearDown(self) -> None:
         """Drop all fixture tables."""
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 216a27f3..1efb1a07 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -548,6 +548,84 @@ test_foo:
     )
 
 
+def test_partially_inferred_schemas(sushi_context: Context, mocker: MockerFixture) -> None:
+    parent = _create_model(
+        "SELECT a, b, s::ROW(d DATE) AS s FROM sushi.unknown",
+        meta="MODEL (name sushi.parent, kind FULL, dialect trino)",
+        default_catalog="memory",
+    )
+    parent = t.cast(SqlModel, sushi_context.upsert_model(parent))
+
+    child = _create_model(
+        "SELECT a, b, s.d AS d FROM sushi.parent",
+        meta="MODEL (name sushi.child, kind FULL)",
+        default_catalog="memory",
+    )
+    child = t.cast(SqlModel, sushi_context.upsert_model(child))
+
+    body = load_yaml(
+        """
+test_child:
+  model: sushi.child
+  inputs:
+    sushi.parent:
+      - a: 1
+        b: bla
+        s:
+          d: 2020-01-01
+  outputs:
+    query:
+      - a: 1
+        b: bla
+        d: 2020-01-01
+        """
+    )
+    test = _create_test(body, "test_child", child, sushi_context)
+
+    random_id = "jzngz56a"
+    test._test_id = random_id
+
+    spy_execute = mocker.spy(test.engine_adapter, "_execute")
+    _check_successful_or_raise(test.run())
+
+    spy_execute.assert_any_call(
+        f'CREATE OR REPLACE VIEW "memory"."sushi"."parent__fixture__{random_id}" ("s", "a", "b") AS '
+        "SELECT "
+        'CAST("s" AS STRUCT("d" DATE)) AS "s", '
+        'CAST("a" AS INT) AS "a", '
+        'CAST("b" AS TEXT) AS "b" '
+        """FROM (VALUES ({'d': CAST('2020-01-01' AS DATE)}, 1, 'bla')) AS "t"("s", "a", "b")"""
+    )
+
+
+def test_uninferrable_schema() -> None:
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
+test_foo:
+  model: sushi.foo
+  inputs:
+    raw:
+      - value: null
+  outputs:
+    query:
+      - value: null
+                """
+            ),
+            test_name="test_foo",
+            model=_create_model("SELECT value FROM raw"),
+            context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+        ).run(),
+        expected_msg=(
+            """Failed to infer the data type of column 'value' for '"raw"'. This issue can """
+            "be mitigated by casting the column in the model definition, setting its type in "
+            "schema.yaml if it's an external model, setting the model's 'columns' property, "
+            "or setting its 'columns' mapping in the test itself\n"
+        ),
+    )
+
+
 def test_missing_column_failure(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
     _check_successful_or_raise(
         _create_test(
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index 571d5aba..4bccc9b8 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -167,7 +167,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
         "models_to_backfill": ['"test_model"'],
         "end_bounded": False,
         "ensure_finalized_snapshots": False,
-        "directly_modified_snapshots": [snapshot.snapshot_id],
+        "directly_modified_snapshots": [{"identifier": "1736759018", "name": '"test_model"'}],
         "indirectly_modified_snapshots": {},
         "removed_snapshots": [],
     }