diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index 8d47cdfd61..7ee351950b 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -258,13 +258,17 @@ def _df_to_source_queries( df, target_columns_to_types, source_columns=source_columns ) + temp_table = self._get_temp_table(target_table or "spark", table_only=True) + def query_factory() -> Query: - temp_table = self._get_temp_table(target_table or "spark", table_only=True) pyspark_df.createOrReplaceTempView(temp_table.sql(dialect=self.dialect)) self._connection_pool.set_attribute("use_spark_engine_adapter", True) return exp.select(*self._select_columns(target_columns_to_types)).from_(temp_table) - return [SourceQuery(query_factory=query_factory)] + def cleanup() -> None: + self.spark.catalog.dropTempView(temp_table.name) + + return [SourceQuery(query_factory=query_factory, cleanup_func=cleanup)] def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False diff --git a/sqlmesh/core/engine_adapter/spark.py b/sqlmesh/core/engine_adapter/spark.py index 9199aa3bcd..72672bafd7 100644 --- a/sqlmesh/core/engine_adapter/spark.py +++ b/sqlmesh/core/engine_adapter/spark.py @@ -282,14 +282,18 @@ def _df_to_source_queries( source_columns: t.Optional[t.List[str]] = None, ) -> t.List[SourceQuery]: df = self._ensure_pyspark_df(df, target_columns_to_types, source_columns=source_columns) + temp_table = self._get_temp_table(target_table or "spark", table_only=True) + temp_view_name = temp_table.sql(dialect=self.dialect) + temp_table.set("db", "global_temp") def query_factory() -> Query: - temp_table = self._get_temp_table(target_table or "spark", table_only=True) - df.createOrReplaceGlobalTempView(temp_table.sql(dialect=self.dialect)) # type: ignore - temp_table.set("db", "global_temp") + df.createOrReplaceGlobalTempView(temp_view_name) # type: ignore return exp.select(*self._select_columns(target_columns_to_types)).from_(temp_table) - return [SourceQuery(query_factory=query_factory)] + def cleanup() -> None: + self.spark.catalog.dropGlobalTempView(temp_table.name) + + return [SourceQuery(query_factory=query_factory, cleanup_func=cleanup)] def _ensure_pyspark_df( self, diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py index 65d8e0388f..7e03a1f11e 100644 --- a/tests/core/engine_adapter/test_databricks.py +++ b/tests/core/engine_adapter/test_databricks.py @@ -837,3 +837,41 @@ def test_databricks_connect_cluster_id_not_overridden_by_falsy_serverless( token="mytoken", cluster_id="0123-456789-mycluster", ) + + +def test_df_source_query_drops_temp_view( + mocker: MockFixture, make_mocked_engine_adapter: t.Callable +): + mocker.patch( + "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog" + ) + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + mocker.patch.object( + DatabricksEngineAdapter, + "_use_spark_session", + new_callable=mocker.PropertyMock, + return_value=True, + ) + spark = mocker.Mock() + mocker.patch.object( + DatabricksEngineAdapter, "spark", new_callable=mocker.PropertyMock, return_value=spark + ) + pyspark_df = mocker.Mock() + mocker.patch.object(adapter, "_ensure_pyspark_df", return_value=pyspark_df) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.to_table("`temp_test_abcdefgh`", dialect="databricks"), + ) + + source_queries = adapter._df_to_source_queries( + pyspark_df, {"a": exp.DataType.build("int")}, batch_size=0, target_table="db.test" + ) + assert len(source_queries) == 1 + + with source_queries[0] as query: + pyspark_df.createOrReplaceTempView.assert_called_once_with("`temp_test_abcdefgh`") + assert query.sql(dialect="databricks") == "SELECT `a` FROM `temp_test_abcdefgh`" + spark.catalog.dropTempView.assert_not_called() + + spark.catalog.dropTempView.assert_called_once_with("temp_test_abcdefgh") diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py index d7c3127f05..fa257161d3 100644 --- a/tests/core/engine_adapter/test_spark.py +++ b/tests/core/engine_adapter/test_spark.py @@ -1106,3 +1106,28 @@ def test_get_data_object_wap_branch(make_mocked_engine_adapter: t.Callable, mock d.schema_("sqlmesh__test", "catalog"), {"test__test_view__630979748"}, ) + + +def test_df_source_query_drops_global_temp_view( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(SparkEngineAdapter) + pyspark_df = mocker.Mock() + mocker.patch.object(adapter, "_ensure_pyspark_df", return_value=pyspark_df) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.to_table("`temp_test_abcdefgh`", dialect="spark"), + ) + + source_queries = adapter._df_to_source_queries( + pyspark_df, {"a": exp.DataType.build("int")}, batch_size=0, target_table="db.test" + ) + assert len(source_queries) == 1 + + with source_queries[0] as query: + pyspark_df.createOrReplaceGlobalTempView.assert_called_once_with("`temp_test_abcdefgh`") + assert query.sql(dialect="spark") == "SELECT `a` FROM global_temp.`temp_test_abcdefgh`" + adapter.spark.catalog.dropGlobalTempView.assert_not_called() + + adapter.spark.catalog.dropGlobalTempView.assert_called_once_with("temp_test_abcdefgh")