Repository: TobikoData/sqlmesh
Pull request #2444 (2024-04-15 17:31:21)

Chore: Add a script for testing the migration
Test steps:
1. Check out the previous release.
2. Initialize the SQLMesh state by running the `plan` command against the prod environment.
3. Switch back to the original state of the git repo.
4. Migrate
5. Make sure there are no differences when comparing the local project state with prod.


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/examples/sushi/config.py b/examples/sushi/config.py
index 7da0dc56..386c2ae9 100644
--- a/examples/sushi/config.py
+++ b/examples/sushi/config.py
@@ -11,7 +11,6 @@ from sqlmesh.core.config import (
     GatewayConfig,
     ModelDefaultsConfig,
     PlanConfig,
-    PostgresConnectionConfig,
     SparkConnectionConfig,
 )
 from sqlmesh.core.notification_target import (
@@ -152,20 +151,3 @@ environment_catalog_mapping_config = Config(
         ".*": "dev_catalog",
     },
 )
-
-postgresql_config = Config(
-    # default_connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/local.duckdb"),
-    default_connection=DuckDBConnectionConfig(database="gclouddata.duckdb"),
-    model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-    gateways={
-        "my_gateway": GatewayConfig(
-            state_connection=PostgresConnectionConfig(
-                host="35.222.243.76",
-                port=5432,
-                user="unique-glider",
-                password="1d1dWu[NRYV5XUY3",
-                database="sushi-demo",
-            ),
-        ),
-    },
-)
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..6f07580a 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1138,17 +1138,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 +1161,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             self.default_catalog,
             no_diff=not detailed,
         )
+        return context_diff.has_changes
 
     def table_diff(
         self,
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index ac185026..19f1b8aa 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -189,7 +189,7 @@ class MacroEvaluator:
                 return exp.convert(
                     self.transform(value) if isinstance(value, exp.Expression) else value
                 )
-            if isinstance(node, exp.Identifier):
+            if isinstance(node, exp.Identifier) and "@" in node.this:
                 text = self.template(node.this, self.locals)
                 if node.this != text:
                     changed = True
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 85c58d4c..d57a1a4c 100644
--- a/sqlmesh/utils/__init__.py
+++ b/sqlmesh/utils/__init__.py
@@ -316,20 +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 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_macros.py b/tests/core/test_macros.py
index 0c57f487..588854c7 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -348,3 +348,9 @@ def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):
 
 def test_macro_returns_none(macro_evaluator):
     assert macro_evaluator.transform(parse_one("@NOOP()")) is None
+
+
+def test_transform_no_template(mocker, macro_evaluator):
+    spy = mocker.spy(macro_evaluator, "template")
+    macro_evaluator.transform(parse_one("select a, b"))
+    spy.assert_not_called()
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": [],
     }
diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py
index dc41b430..48008064 100644
--- a/tests/utils/test_helpers.py
+++ b/tests/utils/test_helpers.py
@@ -1,5 +1,6 @@
 import pytest
 from sqlglot import expressions
+from sqlglot.optimizer.annotate_types import annotate_types
 
 from sqlmesh.utils import columns_to_types_all_known
 
@@ -63,6 +64,10 @@ from sqlmesh.utils import columns_to_types_all_known
             },
             False,
         ),
+        (
+            {"a": annotate_types(expressions.DataType.build("VARCHAR(MAX)", dialect="redshift"))},
+            True,
+        ),
     ],
 )
 def test_columns_to_types_all_known(columns_to_types, expected) -> None: