fix: set catalog properly across connection and engine (#2428)

* fix: remove extra spark set catalog

* fix: consolidate spark catalog operations to connection level

* properly handle if pandas df results are returned

Author: eakmanrq

sqlmesh/core/engine_adapter/spark.py

@@ -22,6 +22,7 @@
     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 @@ def _get_data_objects(
             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 @@ def _get_data_objects(
             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:

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 sqlmesh.engines.spark.db_api.errors import NotSupportedError, ProgrammingError
 
+logger = logging.getLogger(__name__)
+
 
 class SparkSessionCursor:
     def __init__(self, spark: SparkSession):
@@ -65,18 +68,35 @@ def __init__(self, spark: SparkSession, catalog: t.Optional[str] = None):
         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()}")
         except NotImplementedError:
             # 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:

tests/conftest.py

@@ -344,7 +344,7 @@ def _make_function(
         )
         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

tests/core/engine_adapter/test_spark.py

@@ -787,7 +787,7 @@ def check_table_exists(table_name: exp.Table) -> bool:
 
 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,)]