Repository: TobikoData/sqlmesh
Pull request #2429 (2024-04-12 01:19:10)

Feat: allow partial schema inference in unit test
Prior to this change, if there was at least one unknown type in the schema of an input model we would try to infer all of its column types again using a combination of `convert` and `annotate_types`, thus ignoring the types that we'd previously correctly inferred.

However, relying solely on the `convert` and `annotate_types` approach is problematic, because the former is dialect/type-agnostic and thus it would always convert dictionaries into `MAP`s, which led to type mismatch errors (eg. when we would expect `STRUCT`s instead).

This PR ensures we leverage the inferred types, and only fall back to the above strategy for the columns that have unknown types. It also enforces that we'll only create an input fixture if its inferred schema contains known types.

This will open a path for users to fix invalid conversion issues by simply casting their nested objects into the appropriate types.


Patch:
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/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(