Fix: do not include a schema def in CTAS unless all column types are known (#2063)

* Do not include a schema def in CTAS unless all col types known

* Make columns_to_types all known helper

* Make tests cover all cases

Author: treysp

sqlmesh/core/engine_adapter/base.py

@@ -39,7 +39,7 @@
 )
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
-from sqlmesh.utils import double_escape, random_id
+from sqlmesh.utils import columns_to_types_all_known, double_escape, random_id
 from sqlmesh.utils.connection_pool import create_connection_pool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_ts
 from sqlmesh.utils.errors import SQLMeshError, UnsupportedCatalogOperationError
@@ -604,13 +604,14 @@ def _create_table_from_source_queries(
         # types, and for evaluation methods like `LogicalReplaceQueryMixin.replace_query()`
         # calls and SCD Type 2 model calls.
         schema = None
+        columns_to_types_known = columns_to_types and columns_to_types_all_known(columns_to_types)
         if (
             column_descriptions
-            and columns_to_types
+            and columns_to_types_known
             and self.COMMENT_CREATION_TABLE.is_in_schema_def_ctas
             and self.comments_enabled
         ):
-            schema = self._build_schema_exp(table, columns_to_types, column_descriptions)
+            schema = self._build_schema_exp(table, columns_to_types, column_descriptions)  # type: ignore
 
         with self.transaction(condition=len(source_queries) > 1):
             for i, source_query in enumerate(source_queries):

sqlmesh/core/model/definition.py

@@ -33,7 +33,7 @@
 from sqlmesh.core.model.meta import ModelMeta
 from sqlmesh.core.model.seed import CsvSeedReader, Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
-from sqlmesh.utils import str_to_bool
+from sqlmesh.utils import columns_to_types_all_known, str_to_bool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_datetime, to_ds, to_ts
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, raise_config_error
 from sqlmesh.utils.hashing import hash_data
@@ -575,10 +575,7 @@ def annotated(self) -> bool:
         }
         if not columns_to_types:
             return False
-        return all(
-            not column_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL)
-            for column_type in columns_to_types.values()
-        )
+        return columns_to_types_all_known(columns_to_types)
 
     @property
     def sorted_python_env(self) -> t.List[t.Tuple[str, Executable]]:

sqlmesh/utils/__init__.py

@@ -19,6 +19,7 @@
 from functools import lru_cache, reduce, wraps
 from pathlib import Path
 
+from sqlglot import exp
 from sqlglot.dialects.dialect import Dialects
 
 logger = logging.getLogger(__name__)
@@ -305,3 +306,13 @@ def groupby(
     for item in items:
         grouped[func(item)].append(item)
     return grouped
+
+
+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()
+    )

tests/core/engine_adapter/test_base.py

@@ -400,15 +400,24 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
 
     adapter.create_table(
         "test_table",
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
         table_description="test description",
         column_descriptions={"a": "a description"},
     )
 
     adapter.ctas(
         "test_table",
         parse_one("SELECT a, b FROM source_table"),
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
+        table_description="test description",
+        column_descriptions={"a": "a description"},
+    )
+
+    # CTAS call should not include schema definition if UNKNOWN data type present
+    adapter.ctas(
+        "test_table",
+        parse_one("SELECT a, b FROM source_table"),
+        {"a": exp.DataType.build("UNKNOWN"), "b": exp.DataType.build("INT")},
         table_description="test description",
         column_descriptions={"a": "a description"},
     )
@@ -431,8 +440,10 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
 
     sql_calls = to_sql_calls(adapter)
     assert sql_calls == [
-        """CREATE TABLE IF NOT EXISTS "test_table" ("a" int COMMENT 'a description', "b" int) COMMENT='test description'""",
-        """CREATE TABLE IF NOT EXISTS "test_table" ("a" int COMMENT 'a description', "b" int) COMMENT='test description' AS SELECT "a", "b" FROM "source_table\"""",
+        """CREATE TABLE IF NOT EXISTS "test_table" ("a" INT COMMENT 'a description', "b" INT) COMMENT='test description'""",
+        """CREATE TABLE IF NOT EXISTS "test_table" ("a" INT COMMENT 'a description', "b" INT) COMMENT='test description' AS SELECT "a", "b" FROM "source_table\"""",
+        """CREATE TABLE IF NOT EXISTS "test_table" COMMENT='test description' AS SELECT "a", "b" FROM "source_table\"""",
+        """COMMENT ON COLUMN "test_table"."a" IS 'a description'""",
         """CREATE OR REPLACE VIEW "test_view" COMMENT='test description' AS SELECT "a", "b" FROM "source_table\"""",
         """COMMENT ON TABLE "test_table" IS 'test description'""",
         """COMMENT ON COLUMN "test_table"."a" IS 'a description'""",
@@ -1334,7 +1345,9 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable):
 
 def test_replace_query(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
-    adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
+    adapter.replace_query(
+        "test_table", parse_one("SELECT a FROM tbl"), {"a": exp.DataType.build("INT")}
+    )
 
     # TODO: Shouldn't we enforce that `a` is casted to an int?
     assert to_sql_calls(adapter) == [
@@ -1347,7 +1360,9 @@ def test_replace_query_pandas(make_mocked_engine_adapter: t.Callable):
     adapter.DEFAULT_BATCH_SIZE = 1
 
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
-    adapter.replace_query("test_table", df, {"a": "int", "b": "int"})
+    adapter.replace_query(
+        "test_table", df, {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}
+    )
 
     assert to_sql_calls(adapter) == [
         'CREATE OR REPLACE TABLE "test_table" AS SELECT CAST("a" AS INT) AS "a", CAST("b" AS INT) AS "b" FROM (VALUES (1, 4)) AS "t"("a", "b")',

tests/core/engine_adapter/test_bigquery.py

@@ -218,7 +218,9 @@ def test_replace_query(make_mocked_engine_adapter: t.Callable, mocker: MockerFix
     execute_mock = mocker.patch(
         "sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute"
     )
-    adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
+    adapter.replace_query(
+        "test_table", parse_one("SELECT a FROM tbl"), {"a": exp.DataType.build("INT")}
+    )
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == ["CREATE OR REPLACE TABLE `test_table` AS SELECT `a` FROM `tbl`"]
@@ -591,15 +593,15 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
 
     adapter.create_table(
         "test_table",
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
         table_description="test description",
         column_descriptions={"a": "a description"},
     )
 
     adapter.ctas(
         "test_table",
         parse_one("SELECT a, b FROM source_table"),
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
         table_description="test description",
         column_descriptions={"a": "a description"},
     )
@@ -617,8 +619,8 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int OPTIONS (description='a description'), `b` int) OPTIONS (description='test description')",
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int OPTIONS (description='a description'), `b` int) OPTIONS (description='test description') AS SELECT `a`, `b` FROM `source_table`",
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64 OPTIONS (description='a description'), `b` INT64) OPTIONS (description='test description')",
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64 OPTIONS (description='a description'), `b` INT64) OPTIONS (description='test description') AS SELECT `a`, `b` FROM `source_table`",
         "CREATE OR REPLACE VIEW `test_table` OPTIONS (description='test description') AS SELECT `a`, `b` FROM `source_table`",
         "ALTER TABLE `test_table` SET OPTIONS(description = 'test description')",
     ]

tests/core/engine_adapter/test_databricks.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import pytest
 from pytest_mock import MockFixture
-from sqlglot import parse_one
+from sqlglot import exp, parse_one
 
 from sqlmesh.core.engine_adapter import DatabricksEngineAdapter
 from tests.core.engine_adapter import to_sql_calls
@@ -18,7 +18,9 @@ def test_replace_query_not_exists(mocker: MockFixture, make_mocked_engine_adapte
         return_value=False,
     )
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter)
-    adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
+    adapter.replace_query(
+        "test_table", parse_one("SELECT a FROM tbl"), {"a": exp.DataType.build("INT")}
+    )
 
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `test_table` AS SELECT `a` FROM `tbl`",
@@ -47,7 +49,9 @@ def test_replace_query_pandas_not_exists(
     )
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter)
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
-    adapter.replace_query("test_table", df, {"a": "int", "b": "int"})
+    adapter.replace_query(
+        "test_table", df, {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}
+    )
 
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `test_table` AS SELECT CAST(`a` AS INT) AS `a`, CAST(`b` AS INT) AS `b` FROM VALUES (1, 4), (2, 5), (3, 6) AS `t`(`a`, `b`)",

tests/core/engine_adapter/test_spark.py

@@ -171,7 +171,9 @@ def test_replace_query_not_exists(mocker: MockerFixture, make_mocked_engine_adap
         return_value=False,
     )
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
-    adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
+    adapter.replace_query(
+        "test_table", parse_one("SELECT a FROM tbl"), {"a": exp.DataType.build("INT")}
+    )
 
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `test_table` AS SELECT `a` FROM `tbl`",
@@ -876,13 +878,13 @@ def test_replace_query_with_wap_self_reference(
     adapter.replace_query(
         "catalog.schema.table.branch_wap_12345",
         parse_one("SELECT 1 as a FROM catalog.schema.table.branch_wap_12345"),
-        columns_to_types={"a": "int"},
+        columns_to_types={"a": exp.DataType.build("INT")},
         storage_format="ICEBERG",
     )
 
     sql_calls = to_sql_calls(adapter)
     assert sql_calls == [
-        "CREATE TABLE IF NOT EXISTS `catalog`.`schema`.`table` (`a` int)",
+        "CREATE TABLE IF NOT EXISTS `catalog`.`schema`.`table` (`a` INT)",
         "CREATE SCHEMA IF NOT EXISTS `schema`",
         "CREATE TABLE IF NOT EXISTS `catalog`.`schema`.`temp_branch_wap_12345_abcdefgh` USING ICEBERG AS SELECT `a` FROM `catalog`.`schema`.`table`.`branch_wap_12345`",
         "INSERT OVERWRITE TABLE `catalog`.`schema`.`table`.`branch_wap_12345` (`a`) SELECT 1 AS `a` FROM `catalog`.`schema`.`temp_branch_wap_12345_abcdefgh`",

tests/utils/test_helpers.py

@@ -0,0 +1,24 @@
+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
+    )