Repository: TobikoData/sqlmesh
Pull request #2033 (2024-01-27 04:54:55)

fix: Spark Create Table DDL WAP
Fix after this PR: https://github.com/TobikoData/sqlmesh/pull/2031

The PR introduced a flow where we would create a table if not exists if the query is a self-referencing query. This doesn't work for Spark WAP since spark doesn't let you create a table with a WAP ID at the end. This makes sense because to have the branch you needed to create the table first and then altered it. Therefore this PR updates Spark's create table to issue the create table on the table itself without the branch. 

While working on that fix I noticed a different issue: If `_create_table` was called with `IF NOT EXISTS` included and was Iceberg and WAP capable then we would do the "Dummy" Insert into call: https://github.com/TobikoData/sqlmesh/blob/f583fcec955e9433f755737ecbab279888643fc0/sqlmesh/core/engine_adapter/spark.py#L466

The problem would be if the table has data in it which is possible since we could have done "CREATE IF NOT EXIST" and it did exist and have data. In that case we would duplicate all of the data. Therefore I also add an exists check to ensure that we only create the dummy snapshot if the table was actually created. 


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 73e5f26e..b66558f7 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import abc
 import base64
+import logging
 import os
 import pathlib
 import sys
@@ -33,6 +34,9 @@ else:
     from typing_extensions import Literal
 
 
+logger = logging.getLogger(__name__)
+
+
 class ConnectionConfig(abc.ABC, BaseConfig):
     concurrent_tasks: int
     register_comments: bool
@@ -122,6 +126,8 @@ class DuckDBConnectionConfig(ConnectionConfig):
 
     type_: Literal["duckdb"] = Field(alias="type", default="duckdb")
 
+    _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
+
     @model_validator(mode="before")
     @model_validator_v1_args
     def _validate_database_catalogs(
@@ -182,6 +188,29 @@ class DuckDBConnectionConfig(ConnectionConfig):
 
         return init
 
+    def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
+        """Checks if another engine adapter has already been created that shares a catalog that points to the same data
+        file. If so, it uses that same adapter instead of creating a new one. As a result, any additional configuration
+        associated with the new adapter will be ignored."""
+        data_files = set((self.catalogs or {}).values())
+        if self.database:
+            data_files.add(self.database)
+        data_files.discard(":memory:")
+        for data_file in data_files:
+            if adapter := DuckDBConnectionConfig._data_file_to_adapter.get(data_file):
+                logger.info(
+                    f"Using existing DuckDB adapter due to overlapping data file: {data_file}"
+                )
+                return adapter
+        if data_files:
+            logger.info(f"Creating new DuckDB adapter for data files: {data_files}")
+        else:
+            logger.info("Creating new DuckDB adapter for in-memory database")
+        adapter = super().create_engine_adapter(register_comments_override)
+        for data_file in data_files:
+            DuckDBConnectionConfig._data_file_to_adapter[data_file] = adapter
+        return adapter
+
     def get_catalog(self) -> t.Optional[str]:
         if self.database:
             # Remove `:` from the database name in order to handle if `:memory:` is passed in
diff --git a/sqlmesh/core/engine_adapter/spark.py b/sqlmesh/core/engine_adapter/spark.py
index f4b45227..24e2fac1 100644
--- a/sqlmesh/core/engine_adapter/spark.py
+++ b/sqlmesh/core/engine_adapter/spark.py
@@ -61,6 +61,9 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
     # currently check for storage formats we say we don't support REPLACE TABLE
     SUPPORTS_REPLACE_TABLE = False
 
+    WAP_PREFIX = "wap_"
+    BRANCH_PREFIX = "branch_"
+
     @property
     def spark(self) -> PySparkSession:
         return self._connection_pool.get().spark
@@ -428,6 +431,24 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
         column_descriptions: t.Optional[t.Dict[str, str]] = None,
         **kwargs: t.Any,
     ) -> None:
+        table_name = (
+            table_name_or_schema.this
+            if isinstance(table_name_or_schema, exp.Schema)
+            else exp.to_table(table_name_or_schema)
+        )
+        # Spark doesn't support creating a wap table DDL. Therefore we check if this is a wap table and if it is,
+        # this is not a replace, and the table already exists then we can safely just return. Otherwise we let it error.
+        if not expression and isinstance(table_name.this, exp.Dot):
+            wap_id = table_name.this.parts[-1].name
+            if wap_id.startswith(f"{self.BRANCH_PREFIX}{self.WAP_PREFIX}"):
+                table_name.set("this", table_name.this.this)
+
+        wap_supported = (
+            kwargs.get("storage_format") or ""
+        ).lower() == "iceberg" or self.wap_supported(table_name)
+        do_dummy_insert = (
+            False if not wap_supported or not exists else not self.table_exists(table_name)
+        )
         super()._create_table(
             table_name_or_schema,
             expression,
@@ -443,9 +464,7 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
             if isinstance(table_name_or_schema, exp.Schema)
             else exp.to_table(table_name_or_schema)
         )
-        if (kwargs.get("storage_format") or "").lower() == "iceberg" or self.wap_supported(
-            table_name
-        ):
+        if do_dummy_insert:
             # Performing a dummy insert to create a dummy snapshot for Iceberg tables
             # to workaround https://github.com/apache/iceberg/issues/8849.
             dummy_insert = exp.insert(exp.select("*").from_(table_name), table_name)
@@ -459,20 +478,20 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
         )
 
     def wap_table_name(self, table_name: TableName, wap_id: str) -> str:
-        branch_name = _wap_branch_name(wap_id)
+        branch_name = self._wap_branch_name(wap_id)
         fqn = self._ensure_fqn(table_name)
-        return exp.Dot.build([fqn, exp.to_identifier(f"branch_{branch_name}")]).sql(
+        return exp.Dot.build([fqn, exp.to_identifier(f"{self.BRANCH_PREFIX}{branch_name}")]).sql(
             dialect=self.dialect
         )
 
     def wap_prepare(self, table_name: TableName, wap_id: str) -> str:
-        branch_name = _wap_branch_name(wap_id)
+        branch_name = self._wap_branch_name(wap_id)
         fqn = self._ensure_fqn(table_name)
         self.execute(f"ALTER TABLE {fqn.sql(dialect=self.dialect)} CREATE BRANCH {branch_name}")
         return self.wap_table_name(table_name, wap_id)
 
     def wap_publish(self, table_name: TableName, wap_id: str) -> None:
-        branch_name = _wap_branch_name(wap_id)
+        branch_name = self._wap_branch_name(wap_id)
         fqn = self._ensure_fqn(table_name)
 
         get_snapshot_id_query = (
@@ -512,6 +531,6 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
 
         return f"ALTER TABLE {table_sql} ALTER COLUMN {column_sql} COMMENT '{column_comment}'"
 
-
-def _wap_branch_name(wap_id: str) -> str:
-    return f"wap_{wap_id}"
+    @classmethod
+    def _wap_branch_name(cls, wap_id: str) -> str:
+        return f"{cls.WAP_PREFIX}{wap_id}"


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 499fadb0..99dfa9e5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,6 +16,7 @@ from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, maybe_parse, parse_one
 from sqlglot.helper import ensure_list
 
+from sqlmesh.core.config import DuckDBConnectionConfig
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.macros import macro
@@ -176,6 +177,12 @@ def rescope_global_models(request):
     model.set_registry(existing_registry)
 
 
+@pytest.fixture(scope="function", autouse=True)
+def rescope_duckdb_classvar(request):
+    DuckDBConnectionConfig._data_file_to_adapter = {}
+    yield
+
+
 @pytest.fixture
 def duck_conn() -> duckdb.DuckDBPyConnection:
     return duckdb.connect()
diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py
index b78decfc..86e5a599 100644
--- a/tests/core/engine_adapter/test_spark.py
+++ b/tests/core/engine_adapter/test_spark.py
@@ -33,9 +33,9 @@ def test_create_table_properties(make_mocked_engine_adapter: t.Callable):
         storage_format="parquet",
     )
 
-    adapter.cursor.execute.assert_called_once_with(
-        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING PARQUET PARTITIONED BY (`colb`)"
-    )
+    assert to_sql_calls(adapter) == [
+        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING PARQUET PARTITIONED BY (`colb`)",
+    ]
 
     adapter.cursor.reset_mock()
     adapter.create_table(
@@ -45,9 +45,9 @@ def test_create_table_properties(make_mocked_engine_adapter: t.Callable):
         storage_format="parquet",
     )
 
-    adapter.cursor.execute.assert_called_once_with(
-        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING PARQUET PARTITIONED BY (`cola`, `colb`)"
-    )
+    assert to_sql_calls(adapter) == [
+        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING PARQUET PARTITIONED BY (`cola`, `colb`)",
+    ]
 
     with pytest.raises(SQLMeshError):
         adapter.create_table(
@@ -256,7 +256,7 @@ def test_replace_query_self_ref_not_exists(
     temp_table_id = "abcdefgh"
     mocker.patch(
         "sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table",
-        side_effect=lambda *args, **kwargs: make_temp_table_name(table_name, temp_table_id),
+        return_value=make_temp_table_name(table_name, temp_table_id),
     )
 
     mocker.patch(
@@ -329,9 +329,9 @@ def test_create_table_table_options(make_mocked_engine_adapter: t.Callable):
         },
     )
 
-    adapter.cursor.execute.assert_called_once_with(
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) TBLPROPERTIES ('test.conf.key'='value')"
-    )
+    assert to_sql_calls(adapter) == [
+        "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):
@@ -339,9 +339,9 @@ def test_create_state_table(make_mocked_engine_adapter: t.Callable):
 
     adapter.create_state_table("test_table", {"a": "int", "b": "int"}, primary_key=["a"])
 
-    adapter.cursor.execute.assert_called_once_with(
-        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITIONED BY (`a`)"
-    )
+    assert to_sql_calls(adapter) == [
+        "CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITIONED BY (`a`)",
+    ]
 
 
 test_primitive_params = [
@@ -783,7 +783,12 @@ def test_wap_publish(make_mocked_engine_adapter: t.Callable, mocker: MockerFixtu
     )
 
 
-def test_create_table_iceberg(make_mocked_engine_adapter: t.Callable):
+def test_create_table_iceberg(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        return_value=False,
+    )
+
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
 
     columns_to_types = {
@@ -799,17 +804,13 @@ def test_create_table_iceberg(make_mocked_engine_adapter: t.Callable):
         storage_format="ICEBERG",
     )
 
-    adapter.cursor.execute.assert_has_calls(
-        [
-            call(
-                "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING ICEBERG PARTITIONED BY (`colb`)"
-            ),
-            call("INSERT INTO `test_table` SELECT * FROM `test_table`"),
-        ]
-    )
+    assert to_sql_calls(adapter) == [
+        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING, `colc` STRING) USING ICEBERG PARTITIONED BY (`colb`)",
+        "INSERT INTO `test_table` SELECT * FROM `test_table`",
+    ]
 
 
-def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+def test_comments(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
 
     adapter._create_table_comment(
@@ -827,3 +828,54 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         "COMMENT ON TABLE `test_table` IS 'test description'",
         "ALTER TABLE `test_table` ALTER COLUMN `a` COMMENT 'a description'",
     ]
+
+
+def test_create_table_with_wap(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        return_value=False,
+    )
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+
+    adapter.create_table(
+        "catalog.schema.table.branch_wap_12345",
+        {"a": "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",
+        "INSERT INTO `catalog`.`schema`.`table` SELECT * FROM `catalog`.`schema`.`table`",
+    ]
+
+
+def test_replace_query_with_wap_self_reference(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
+        return_value=True,
+    )
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.base.random_id",
+        return_value="abcdefgh",
+    )
+
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+
+    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"},
+        storage_format="ICEBERG",
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        "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`",
+        "DROP TABLE IF EXISTS `catalog`.`schema`.`temp_branch_wap_12345_abcdefgh`",
+    ]
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 64937fb1..702be84e 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -366,3 +366,145 @@ def test_duckdb(make_config):
     )
     assert isinstance(config, DuckDBConnectionConfig)
     assert config._static_connection_kwargs == {"config": {"foo": "bar"}}
+
+
+@pytest.mark.parametrize(
+    "kwargs1, kwargs2, shared_adapter",
+    [
+        (
+            {
+                "database": "test.duckdb",
+            },
+            {
+                "database": "test.duckdb",
+            },
+            True,
+        ),
+        (
+            {},
+            {},
+            False,
+        ),
+        (
+            {
+                "database": "test1.duckdb",
+            },
+            {
+                "database": "test2.duckdb",
+            },
+            False,
+        ),
+        (
+            {
+                "database": ":memory:",
+            },
+            {
+                "database": ":memory:",
+            },
+            False,
+        ),
+        (
+            {
+                "database": ":memory:",
+            },
+            {
+                "database": "test.duckdb",
+            },
+            False,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test": "test.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test": "test.duckdb",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test": ":memory:",
+                }
+            },
+            {
+                "catalogs": {
+                    "test": ":memory:",
+                }
+            },
+            False,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": "test2.duckdb",
+                    "test2": ":memory:",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": "test1.duckdb",
+                    "test2": "test2.duckdb",
+                    "test3": "test3.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": "test1_miss.duckdb",
+                    "test2": "test2_miss.duckdb",
+                    "test3": "test3.duckdb",
+                }
+            },
+            True,
+        ),
+    ],
+)
+def test_duckdb_shared(make_config, caplog, kwargs1, kwargs2, shared_adapter):
+    config1 = make_config(
+        type="duckdb",
+        **kwargs1,
+    )
+    config2 = make_config(
+        type="duckdb",
+        **kwargs2,
+    )
+    assert isinstance(config1, DuckDBConnectionConfig)
+    assert isinstance(config2, DuckDBConnectionConfig)
+    adapter1 = config1.create_engine_adapter()
+    adapter2 = config2.create_engine_adapter()
+    if shared_adapter:
+        assert id(adapter1) == id(adapter2)
+        assert "Creating new DuckDB adapter" in caplog.messages[0]
+        assert "Using existing DuckDB adapter" in caplog.messages[1]
+    else:
+        assert id(adapter1) != id(adapter2)
+        assert "Creating new DuckDB adapter" in caplog.messages[0]
+        assert "Creating new DuckDB adapter" in caplog.messages[1]