Repository: TobikoData/sqlmesh
Pull request #2235 (2024-03-08 00:30:32)

fix!: recursively check for unknown data types
Prior to this change you could have a struct with a unknown data type in it and we would consider this as having all the data types defined. This PR makes the data type check recursive so we will find unknown types in nested data types. 

`SchemaDiffer` was changed because I converted a classmethod it had into a utility function. No logic changes to that class or it's tests. 


Patch:
diff --git a/sqlmesh/core/schema_diff.py b/sqlmesh/core/schema_diff.py
index 23161c41..2544a150 100644
--- a/sqlmesh/core/schema_diff.py
+++ b/sqlmesh/core/schema_diff.py
@@ -8,6 +8,7 @@ from enum import Enum, auto
 from sqlglot import exp
 from sqlglot.helper import ensure_list, seq_get
 
+from sqlmesh.utils import columns_to_types_to_struct
 from sqlmesh.utils.pydantic import PydanticModel
 
 if t.TYPE_CHECKING:
@@ -303,16 +304,6 @@ class SchemaDiffer(PydanticModel):
 
     _coerceable_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = {}
 
-    @classmethod
-    def _dict_to_struct(cls, value: t.Dict[str, exp.DataType]) -> exp.DataType:
-        return exp.DataType(
-            this=exp.DataType.Type.STRUCT,
-            expressions=[
-                exp.ColumnDef(this=exp.to_identifier(k), kind=v) for k, v in value.items()
-            ],
-            nested=True,
-        )
-
     @property
     def coerceable_types(self) -> t.Dict[exp.DataType, t.Set[exp.DataType]]:
         if not self._coerceable_types:
@@ -588,7 +579,7 @@ class SchemaDiffer(PydanticModel):
             The list of schema deltas.
         """
         return self.compare_structs(
-            table_name, self._dict_to_struct(current), self._dict_to_struct(new)
+            table_name, columns_to_types_to_struct(current), columns_to_types_to_struct(new)
         )
 
 
diff --git a/sqlmesh/utils/__init__.py b/sqlmesh/utils/__init__.py
index ef74de47..8dc31ea1 100644
--- a/sqlmesh/utils/__init__.py
+++ b/sqlmesh/utils/__init__.py
@@ -308,11 +308,35 @@ def groupby(
     return grouped
 
 
+def columns_to_types_to_struct(
+    columns_to_types: t.Union[t.Dict[str, exp.DataType], t.Dict[str, str]]
+) -> exp.DataType:
+    """
+    Converts a dict of column names to types to a struct.
+    """
+    return exp.DataType(
+        this=exp.DataType.Type.STRUCT,
+        expressions=[
+            exp.ColumnDef(this=exp.to_identifier(k), kind=v) for k, v in columns_to_types.items()
+        ],
+        nested=True,
+    )
+
+
 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(
-        not column_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL)
-        for column_type in columns_to_types.values()
-    )
+
+    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):
+            return False
+        if d_type.expressions:
+            return all(check_expression_unknown(expression) for expression in d_type.expressions)
+        return True
+
+    return all(check_expression_unknown(expression) for expression in columns_to_types.values())


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index bea51f86..304cbb75 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -14,6 +14,7 @@ from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.engine_adapter import EngineAdapter, EngineAdapterWithIndexSupport
 from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy
 from sqlmesh.core.schema_diff import SchemaDiffer, TableAlterOperation
+from sqlmesh.utils import columns_to_types_to_struct
 from sqlmesh.utils.date import to_ds
 from sqlmesh.utils.errors import SQLMeshError, UnsupportedCatalogOperationError
 from tests.core.engine_adapter import to_sql_calls
@@ -807,7 +808,7 @@ def test_alter_table(
         operations = original_from_structs(current_struct, new_struct)
         assert (
             operations[-1].expected_table_struct.sql()
-            == SchemaDiffer._dict_to_struct(expected_final_structure).sql()
+            == columns_to_types_to_struct(expected_final_structure).sql()
         )
         return operations
 
diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py
index efe42535..dc41b430 100644
--- a/tests/utils/test_helpers.py
+++ b/tests/utils/test_helpers.py
@@ -1,24 +1,69 @@
+import pytest
 from sqlglot import expressions
 
 from sqlmesh.utils import columns_to_types_all_known
 
 
-def test_columns_to_types_all_known() -> None:
-    assert (
-        columns_to_types_all_known(
-            {"a": expressions.DataType.build("INT"), "b": expressions.DataType.build("INT")}
-        )
-        == True
-    )
-    assert (
-        columns_to_types_all_known(
-            {"a": expressions.DataType.build("UNKNOWN"), "b": expressions.DataType.build("INT")}
-        )
-        == False
-    )
-    assert (
-        columns_to_types_all_known(
-            {"a": expressions.DataType.build("NULL"), "b": expressions.DataType.build("INT")}
-        )
-        == False
-    )
+@pytest.mark.parametrize(
+    "columns_to_types, expected",
+    [
+        ({"a": expressions.DataType.build("INT"), "b": expressions.DataType.build("INT")}, True),
+        (
+            {"a": expressions.DataType.build("UNKNOWN"), "b": expressions.DataType.build("INT")},
+            False,
+        ),
+        ({"a": expressions.DataType.build("NULL"), "b": expressions.DataType.build("INT")}, False),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build(
+                    "STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d INT>"
+                ),
+            },
+            True,
+        ),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build(
+                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d INT>>"
+                ),
+            },
+            True,
+        ),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build(
+                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d UNKNOWN>>"
+                ),
+            },
+            False,
+        ),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build(
+                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d UNKNOWN>>"
+                ),
+            },
+            False,
+        ),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build("MAP<INT, STRING>"),
+            },
+            True,
+        ),
+        (
+            {
+                "a": expressions.DataType.build("INT"),
+                "b": expressions.DataType.build("MAP<INT, UNKNOWN>"),
+            },
+            False,
+        ),
+    ],
+)
+def test_columns_to_types_all_known(columns_to_types, expected) -> None:
+    assert columns_to_types_all_known(columns_to_types) == expected