fix: Spark Create Table DDL WAP (#2033)

* fix: spark create table ddl wap

* feedback

Author: eakmanrq

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 @@ def _create_table(
         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 @@ def _create_table(
             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 @@ def wap_supported(self, table_name: TableName) -> bool:
         )
 
     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 @@ def _build_create_comment_column_exp(
 
         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}"

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,19 +329,19 @@ 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):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
 
     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`",
+    ]