Repository: TobikoData/sqlmesh
Pull request #2075 (2024-02-02 00:15:24)

fix: replace query unknown self-reference
Prior to this change if you had a self-referencing query that is full refresh (SCD Type 2 is this) then we would always do a create table if not exists since we need a table to reference if one doesn't exist. If a column was unknown then we would get an error trying to create the table.

Now we check if column types are all known in create table and if they aren't then we raise an error. The exception is if the table already exists and we get IF NOT EXISTS then we just return since we aren't creating the table anyways. This addresses the error case here since even though we have an unknown it is ok since the table is already created. 


Patch:
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index f0f60533..d33bedcc 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -495,6 +495,17 @@ class EngineAdapter:
             kwargs: Optional create table properties.
         """
         table = exp.to_table(table_name)
+
+        if not columns_to_types_all_known(columns_to_types):
+            # It is ok if the columns types are not known if the table already exists and IF NOT EXISTS is set
+            if exists and self.table_exists(table_name):
+                return
+            raise SQLMeshError(
+                "Cannot create a table without knowing the column types. "
+                "Try casting the columns to an expected type or defining the columns in the model metadata. "
+                f"Columns to types: {columns_to_types}"
+            )
+
         primary_key_expression = (
             [exp.PrimaryKey(expressions=[exp.to_column(k) for k in primary_key])]
             if primary_key and self.SUPPORTS_INDEXES


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index 2cbb6641..b82a6746 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -15,7 +15,7 @@ from sqlmesh.core.engine_adapter import EngineAdapter, EngineAdapterWithIndexSup
 from sqlmesh.core.engine_adapter.shared import InsertOverwriteStrategy
 from sqlmesh.core.schema_diff import SchemaDiffer, TableAlterOperation
 from sqlmesh.utils.date import to_ds
-from sqlmesh.utils.errors import UnsupportedCatalogOperationError
+from sqlmesh.utils.errors import SQLMeshError, UnsupportedCatalogOperationError
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = pytest.mark.engine
@@ -454,14 +454,14 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
 
     adapter_no_comments.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"},
     )
 
     sql_calls = to_sql_calls(adapter_no_comments)
     assert sql_calls == [
-        """CREATE TABLE IF NOT EXISTS "test_table" ("a" int, "b" int)""",
+        """CREATE TABLE IF NOT EXISTS "test_table" ("a" INT, "b" INT)""",
     ]
 
 
@@ -1371,6 +1371,70 @@ def test_replace_query_pandas(make_mocked_engine_adapter: t.Callable):
     ]
 
 
+def test_replace_query_self_referencing_not_exists_unknown(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.EngineAdapter.table_exists",
+        return_value=False,
+    )
+
+    with pytest.raises(
+        SQLMeshError,
+        match="Cannot create a table without knowing the column types.*",
+    ):
+        adapter.replace_query(
+            "test",
+            parse_one("SELECT a FROM test"),
+            columns_to_types={"a": exp.DataType.build("UNKNOWN")},
+        )
+
+
+def test_replace_query_self_referencing_exists(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.EngineAdapter.table_exists",
+        return_value=True,
+    )
+
+    adapter.replace_query(
+        "test",
+        parse_one("SELECT a FROM test"),
+        columns_to_types={"a": exp.DataType.build("UNKNOWN")},
+    )
+
+    assert to_sql_calls(adapter) == [
+        'CREATE OR REPLACE TABLE "test" AS SELECT "a" FROM "test"',
+    ]
+
+
+def test_replace_query_self_referencing_not_exists_known(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.EngineAdapter.table_exists",
+        return_value=False,
+    )
+
+    adapter.replace_query(
+        "test",
+        parse_one("SELECT a FROM test"),
+        columns_to_types={"a": exp.DataType.build("INT")},
+    )
+
+    assert to_sql_calls(adapter) == [
+        'CREATE TABLE IF NOT EXISTS "test" ("a" INT)',
+        'CREATE OR REPLACE TABLE "test" AS SELECT "a" FROM "test"',
+    ]
+
+
 def test_create_table_like(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
diff --git a/tests/core/engine_adapter/test_bigquery.py b/tests/core/engine_adapter/test_bigquery.py
index a9019dd4..dd003c71 100644
--- a/tests/core/engine_adapter/test_bigquery.py
+++ b/tests/core/engine_adapter/test_bigquery.py
@@ -317,7 +317,7 @@ def test_create_table_date_partition(
     )
     adapter.create_table(
         "test_table",
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")},
         partitioned_by=partition_by_cols,
         partition_interval_unit=IntervalUnit.DAY,
         clustered_by=["b"],
@@ -325,7 +325,7 @@ def test_create_table_date_partition(
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
-        f"CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITION BY {partition_by_statement} CLUSTER BY `b`"
+        f"CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64, `b` INT64) PARTITION BY {partition_by_statement} CLUSTER BY `b`"
     ]
 
 
@@ -572,7 +572,7 @@ def test_create_table_table_options(make_mocked_engine_adapter: t.Callable, mock
 
     adapter.create_table(
         "test_table",
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")},
         table_properties={
             "partition_expiration_days": exp.convert(7),
         },
@@ -580,7 +580,7 @@ def test_create_table_table_options(make_mocked_engine_adapter: t.Callable, mock
 
     sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) OPTIONS (partition_expiration_days=7)"
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64, `b` INT64) OPTIONS (partition_expiration_days=7)"
     ]
 
 
diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py
index 5097a963..c4216b8b 100644
--- a/tests/core/engine_adapter/test_spark.py
+++ b/tests/core/engine_adapter/test_spark.py
@@ -238,10 +238,6 @@ def test_replace_query_pandas_exists(make_mocked_engine_adapter: t.Callable, moc
 def test_replace_query_self_ref_not_exists(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
 ):
-    mocker.patch(
-        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
-        return_value=False,
-    )
     mocker.patch(
         "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.get_current_catalog",
         lambda self: "spark_catalog",
@@ -266,13 +262,25 @@ def test_replace_query_self_ref_not_exists(
         return_value={"col": exp.DataType(this=exp.DataType.Type.INT)},
     )
 
+    def check_table_exists(table_name: exp.Table) -> bool:
+        for sql in to_sql_calls(adapter):
+            if f"CREATE TABLE IF NOT EXISTS {table_name.sql(dialect=adapter.dialect)}" in sql:
+                return True
+        return False
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        side_effect=check_table_exists,
+    )
+
     adapter.replace_query(table_name, parse_one(f"SELECT col + 1 AS col FROM {table_name}"))
 
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `db`.`table` (`col` INT)",
-        # This second query doesn't do anything but we can't just run this because we need the table
-        # to exist since it is self-referencing
-        "CREATE TABLE IF NOT EXISTS `db`.`table` AS SELECT `col` + 1 AS `col` FROM `db`.`table`",
+        "CREATE SCHEMA IF NOT EXISTS `db`",
+        "CREATE TABLE IF NOT EXISTS `db`.`temp_table_abcdefgh` AS SELECT `col` FROM `db`.`table`",
+        "INSERT OVERWRITE TABLE `db`.`table` (`col`) SELECT `col` + 1 AS `col` FROM `db`.`temp_table_abcdefgh`",
+        "DROP TABLE IF EXISTS `db`.`temp_table_abcdefgh`",
     ]
 
 
@@ -325,24 +333,28 @@ def test_create_table_table_options(make_mocked_engine_adapter: t.Callable):
 
     adapter.create_table(
         "test_table",
-        {"a": "int", "b": "int"},
+        {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")},
         table_properties={
             "test.conf.key": exp.convert("value"),
         },
     )
 
     assert to_sql_calls(adapter) == [
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) TBLPROPERTIES ('test.conf.key'='value')",
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) TBLPROPERTIES ('test.conf.key'='value')",
     ]
 
 
 def test_create_state_table(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
 
-    adapter.create_state_table("test_table", {"a": "int", "b": "int"}, primary_key=["a"])
+    adapter.create_state_table(
+        "test_table",
+        {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")},
+        primary_key=["a"],
+    )
 
     assert to_sql_calls(adapter) == [
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITIONED BY (`a`)",
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) PARTITIONED BY (`a`)",
     ]
 
 
@@ -542,6 +554,11 @@ def test_spark_struct_complex_to_col_to_types(type_name, spark_type):
 def test_scd_type_2(
     make_mocked_engine_adapter: t.Callable, make_temp_table_name: t.Callable, mocker: MockerFixture
 ):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        return_value=False,
+    )
+
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
     adapter.spark.catalog.currentCatalog.return_value = "spark_catalog"
     adapter.spark.catalog.currentDatabase.return_value = "default"
@@ -551,6 +568,17 @@ def test_scd_type_2(
     temp_table_id = "abcdefgh"
     temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
 
+    def check_table_exists(table_name: exp.Table) -> bool:
+        for sql in to_sql_calls(adapter):
+            if f"CREATE TABLE IF NOT EXISTS {table_name.sql(dialect=adapter.dialect)}" in sql:
+                return True
+        return False
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        side_effect=check_table_exists,
+    )
+
     adapter.scd_type_2(
         target_table="db.target",
         source_table=t.cast(
@@ -573,7 +601,6 @@ def test_scd_type_2(
 
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `db`.`target` (`id` INT, `name` STRING, `price` DOUBLE, `test_updated_at` TIMESTAMP, `test_valid_from` TIMESTAMP, `test_valid_to` TIMESTAMP)",
-        "DESCRIBE `db`.`target`",
         "CREATE SCHEMA IF NOT EXISTS `db`",
         "CREATE TABLE IF NOT EXISTS `db`.`temp_target_abcdefgh` AS SELECT `id`, `name`, `price`, `test_updated_at`, `test_valid_from`, `test_valid_to` FROM `db`.`target`",
         parse_one(
@@ -850,13 +877,13 @@ def test_create_table_with_wap(make_mocked_engine_adapter: t.Callable, mocker: M
 
     adapter.create_table(
         "catalog.schema.table.branch_wap_12345",
-        {"a": "int"},
+        {"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) USING ICEBERG",
+        "CREATE TABLE IF NOT EXISTS `catalog`.`schema`.`table` (`a` INT) USING ICEBERG",
         "INSERT INTO `catalog`.`schema`.`table` SELECT * FROM `catalog`.`schema`.`table`",
     ]