Fix(spark): support replacing self-referencing queries (#1642)

* Handle spark self-referencing queries

* Enable SCD type 2 test for spark and mssql

* Fix self referencing return

* Move unit test from mixins to spark

Author: treysp

sqlmesh/core/engine_adapter/spark.py

@@ -6,14 +6,18 @@
 
 import pandas as pd
 from sqlglot import exp
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.base import (
     CatalogSupport,
     InsertOverwriteStrategy,
     SourceQuery,
 )
-from sqlmesh.core.engine_adapter.mixins import HiveMetastoreTablePropertiesMixin
+from sqlmesh.core.engine_adapter.mixins import (
+    HiveMetastoreTablePropertiesMixin,
+    LogicalReplaceQueryMixin,
+)
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType, set_catalog
 from sqlmesh.utils import classproperty
 from sqlmesh.utils.errors import SQLMeshError
@@ -338,6 +342,21 @@ def replace_query(
         columns_to_types = columns_to_types or self.columns(table_name)
         if not columns_to_types:
             raise SQLMeshError("Cannot replace table without columns to types")
+
+        # Self-referential queries: cannot insert overwrite a SELECT from itself, so
+        # use LogicalReplaceQuery (which creates a temp table and SELECTs from it)
+        if len(source_queries) > 1:
+            raise SQLMeshError("Cannot replace table with a batched dataframe")
+        with source_queries[0] as query:
+            target_table = exp.to_table(table_name)
+            self_referencing = any(
+                quote_identifiers(table) == quote_identifiers(target_table)
+                for table in query.find_all(exp.Table)
+            )
+
+            if self_referencing:
+                return LogicalReplaceQueryMixin.replace_query(self, table_name, query, columns_to_types)  # type: ignore
+
         self.create_table(table_name, columns_to_types)
         return self._insert_overwrite_by_condition(
             table_name, source_queries, columns_to_types, where=exp.true()
@@ -382,3 +401,7 @@ def create_view(
         super().create_view(
             view_name, query_or_df, columns_to_types, replace, materialized, **create_kwargs
         )
+
+    def _truncate_table(self, table_name: TableName) -> str:
+        table = quote_identifiers(exp.to_table(table_name))
+        return f"TRUNCATE TABLE {table.sql(dialect=self.dialect)}"

tests/core/engine_adapter/test_integration.py

@@ -138,7 +138,7 @@ def input_data(
         return self._format_df(data, to_datetime=self.dialect != "trino")
 
     def output_data(self, data: pd.DataFrame) -> pd.DataFrame:
-        return self._format_df(data, include_tz=self.dialect in ("spark", "databricks"))
+        return self._format_df(data, include_tz=self.dialect == "databricks")
 
     def table(self, table_name: str, schema: str = TEST_SCHEMA) -> exp.Table:
         return exp.to_table(
@@ -737,17 +737,9 @@ def test_merge(ctx: TestContext):
 
 
 def test_scd_type_2(ctx: TestContext):
-    if ctx.dialect == "tsql":
-        pytest.skip(f"MSSQL scd type 2 functionality waiting on sqlglot cte in FROM fix")
-    if ctx.dialect == "spark":
-        pytest.skip(
-            "Spark SCD Type 2 does not currently work due to this error: `[UNSUPPORTED_OVERWRITE.TABLE] Can't overwrite the target that is also being read from`"
-        )
-
-    name_type = "varchar(max)" if ctx.dialect == "tsql" else "string"
     ctx.columns_to_types = {
         "id": "int",
-        "name": name_type,
+        "name": "string",
         "updated_at": "timestamp",
         "valid_from": "timestamp",
         "valid_to": "timestamp",

tests/core/engine_adapter/test_spark.py

@@ -129,6 +129,38 @@ def test_replace_query_pandas(make_mocked_engine_adapter: t.Callable, mocker: Mo
     ]
 
 
+def test_replace_query_self_ref(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    adapter.cursor.fetchone.return_value = (1,)
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "db.table"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.LogicalReplaceQueryMixin.table_exists",
+        return_value=True,
+    )
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.columns",
+        return_value={"col": exp.DataType(this=exp.DataType.Type.INT)},
+    )
+
+    adapter.replace_query(table_name, parse_one(f"SELECT col + 1 AS col FROM {table_name}"))
+
+    assert to_sql_calls(adapter) == [
+        "DESCRIBE `db`.`table`",
+        "CREATE SCHEMA IF NOT EXISTS `db`",
+        f"CREATE TABLE IF NOT EXISTS `db`.`__temp_table_{temp_table_id}` AS SELECT `col` FROM `db`.`table`",
+        "TRUNCATE TABLE `db`.`table`",
+        f"INSERT INTO `db`.`table` (`col`) SELECT `col` + 1 AS `col` FROM `db`.`__temp_table_{temp_table_id}`",
+        f"DROP TABLE IF EXISTS `db`.`__temp_table_{temp_table_id}`",
+    ]
+
+
 def test_create_table_table_options(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)