Repository: TobikoData/sqlmesh
Pull request #2428 (2024-04-12 01:14:50)

fix: set catalog properly across connection and engine
Prior to this change we were doing catalog setting in two different places (at connection level and at the adapter level) which made it easy to forget one of them when making changes like in https://github.com/TobikoData/sqlmesh/pull/2409. Therefore this consolidates to the connection level. Passes integration tests for both Spark and Databricks. 

An alternative change would be to remove setting catalog at the connection level. This is done to protect against ambiguous references to make sure the behavior is consistent. We shouldn't expect this internally within SQLMesh but users could use public facing methods like `fetchdf` and they would expect the catalog they set at the connection level to be applied when running those queries. 


Patch:
diff --git a/sqlmesh/core/engine_adapter/spark.py b/sqlmesh/core/engine_adapter/spark.py
index a6b68527..ac0806d5 100644
--- a/sqlmesh/core/engine_adapter/spark.py
+++ b/sqlmesh/core/engine_adapter/spark.py
@@ -22,6 +22,7 @@ from sqlmesh.core.engine_adapter.shared import (
     SourceQuery,
     set_catalog,
 )
+from sqlmesh.engines.spark.db_api.spark_session import SparkSessionConnection
 from sqlmesh.utils import classproperty
 from sqlmesh.utils.errors import SQLMeshError
 
@@ -58,9 +59,13 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
     WAP_PREFIX = "wap_"
     BRANCH_PREFIX = "branch_"
 
+    @property
+    def connection(self) -> SparkSessionConnection:
+        return self._connection_pool.get()
+
     @property
     def spark(self) -> PySparkSession:
-        return self._connection_pool.get().spark
+        return self.connection.spark
 
     @property
     def _use_spark_session(self) -> bool:
@@ -319,7 +324,8 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
             DataObject(
                 catalog=self.get_current_catalog(),
                 # This varies between Spark and Databricks
-                schema=row.asDict().get("namespace") or row["database"],
+                schema=(row.asDict() if not isinstance(row, dict) else row).get("namespace")
+                or row["database"],
                 name=row["tableName"],
                 type=(
                     DataObjectType.VIEW
@@ -330,26 +336,13 @@ class SparkEngineAdapter(GetCurrentCatalogFromFunctionMixin, HiveMetastoreTableP
             for row in results  # type: ignore
         ]
 
-    @property
-    def _spark_major_minor(self) -> t.Tuple[int, int]:
-        return tuple(int(x) for x in self.spark.version.split(".")[:2])  # type: ignore
-
     def get_current_catalog(self) -> t.Optional[str]:
         if self._use_spark_session:
-            if self._spark_major_minor >= (3, 4):
-                return self.spark.catalog.currentCatalog()
-            else:
-                return self._default_catalog or "spark_catalog"
+            return self.connection.get_current_catalog()
         return super().get_current_catalog()
 
     def set_current_catalog(self, catalog_name: str) -> None:
-        if self._spark_major_minor >= (3, 4):
-            return self.spark.catalog.setCurrentCatalog(catalog_name)
-        current_catalog = self.get_current_catalog()
-        if current_catalog != catalog_name:
-            logger.warning(
-                "Spark <3.4 does not support certain cross catalog queries since the default catalog cannot be set <3.4"
-            )
+        self.connection.set_current_catalog(catalog_name)
 
     def get_current_database(self) -> str:
         if self._use_spark_session:
diff --git a/sqlmesh/engines/spark/db_api/spark_session.py b/sqlmesh/engines/spark/db_api/spark_session.py
index d10ccd81..a46d1f35 100644
--- a/sqlmesh/engines/spark/db_api/spark_session.py
+++ b/sqlmesh/engines/spark/db_api/spark_session.py
@@ -1,3 +1,4 @@
+import logging
 import typing as t
 from threading import get_ident
 
@@ -6,6 +7,8 @@ from pyspark.sql.types import Row
 
 from sqlmesh.engines.spark.db_api.errors import NotSupportedError, ProgrammingError
 
+logger = logging.getLogger(__name__)
+
 
 class SparkSessionCursor:
     def __init__(self, spark: SparkSession):
@@ -65,6 +68,24 @@ class SparkSessionConnection:
         self.spark = spark
         self.catalog = catalog
 
+    @property
+    def _spark_major_minor(self) -> t.Tuple[int, int]:
+        return tuple(int(x) for x in self.spark.version.split(".")[:2])  # type: ignore
+
+    def get_current_catalog(self) -> t.Optional[str]:
+        if self._spark_major_minor >= (3, 4):
+            return self.spark.catalog.currentCatalog()
+        return self.catalog or "spark_catalog"
+
+    def set_current_catalog(self, catalog_name: str) -> None:
+        if self._spark_major_minor >= (3, 4):
+            return self.spark.catalog.setCurrentCatalog(catalog_name)
+        current_catalog = self.get_current_catalog()
+        if current_catalog != catalog_name:
+            logger.warning(
+                "Spark <3.4 does not support certain cross catalog queries since the default catalog cannot be set <3.4"
+            )
+
     def cursor(self) -> SparkSessionCursor:
         try:
             self.spark.sparkContext.setLocalProperty("spark.scheduler.pool", f"pool_{get_ident()}")
@@ -72,11 +93,10 @@ class SparkSessionConnection:
             # Databricks Connect does not support accessing the SparkContext
             pass
         if self.catalog:
-            # Note: Spark 3.4+ Only API
             from py4j.protocol import Py4JError
 
             try:
-                self.spark.catalog.setCurrentCatalog(self.catalog)
+                self.set_current_catalog(self.catalog)
             # Databricks does not support `setCurrentCatalog` with Unity catalog
             # and shared clusters so we use the Databricks Unity only SQL command instead
             except Py4JError:


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index dd86bb6d..2f43d758 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -344,7 +344,7 @@ def make_mocked_engine_adapter(mocker: MockerFixture) -> t.Callable:
         )
         if isinstance(adapter, SparkEngineAdapter):
             mocker.patch(
-                "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter._spark_major_minor",
+                "sqlmesh.engines.spark.db_api.spark_session.SparkSessionConnection._spark_major_minor",
                 new_callable=PropertyMock(return_value=(3, 5)),
             )
         return adapter
diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py
index cb67764d..9d80c12d 100644
--- a/tests/core/engine_adapter/test_spark.py
+++ b/tests/core/engine_adapter/test_spark.py
@@ -787,7 +787,7 @@ FROM `inserted_rows`
 
 def test_wap_prepare(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
-    adapter.spark.catalog.currentCatalog.return_value = "spark_catalog"
+    adapter.connection.get_current_catalog.return_value = "spark_catalog"
     adapter.spark.catalog.currentDatabase.return_value = "default"
 
     table_name = "test_db.test_table"
@@ -805,7 +805,7 @@ def test_wap_publish(make_mocked_engine_adapter: t.Callable, mocker: MockerFixtu
     iceberg_snapshot_id = 123
 
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
-    adapter.spark.catalog.currentCatalog.return_value = "spark_catalog"
+    adapter.connection.get_current_catalog.return_value = "spark_catalog"
     adapter.spark.catalog.currentDatabase.return_value = "default"
     adapter.cursor.fetchall.return_value = [(iceberg_snapshot_id,)]