feat: databricks allow disabling spark session (#2703)

Author: eakmanrq

docs/integrations/engines/databricks.md

@@ -38,6 +38,7 @@ Note: If using Databricks Connect please note the [requirements](https://docs.da
 | `databricks_connect_cluster_id`      | Databricks Connect Only: Databricks Connect cluster ID. Uses `http_path` if not set. Cannot be a Databricks SQL Warehouse.                                                                                                                                                                      | string |    N     |
 | `force_databricks_connect`           | When running locally, force the use of Databricks Connect for all model operations (so don't use SQL Connector for SQL models)                                                                                                                                                                  |  bool  |    N     |
 | `disable_databricks_connect`         | When running locally, disable the use of Databricks Connect for all model operations (so use SQL Connector for all models)                                                                                                                                                                      |  bool  |    N     |
+| `disable_spark_session`              | Do not use SparkSession if it is available (like when running in a notebook).                                                                                                                                                                                                                   |  bool  |    N     |
 
 ## Airflow Scheduler
 **Engine Name:** `databricks` / `databricks-submit` / `databricks-sql`.

sqlmesh/core/config/connection.py

@@ -511,6 +511,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
             Defaults to deriving the cluster id from the `http_path` value.
         force_databricks_connect: Force all queries to run using Databricks Connect instead of the SQL connector.
         disable_databricks_connect: Even if databricks connect is installed, do not use it.
+        disable_spark_session: Do not use SparkSession if it is available (like when running in a notebook).
         pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
@@ -525,6 +526,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
     databricks_connect_cluster_id: t.Optional[str] = None
     force_databricks_connect: bool = False
     disable_databricks_connect: bool = False
+    disable_spark_session: bool = False
 
     concurrent_tasks: int = 1
     register_comments: bool = True
@@ -538,12 +540,11 @@ class DatabricksConnectionConfig(ConnectionConfig):
     @model_validator(mode="before")
     @model_validator_v1_args
     def _databricks_connect_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
-        from sqlmesh import RuntimeEnv
         from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 
-        runtime_env = RuntimeEnv.get()
-
-        if runtime_env.is_databricks:
+        if DatabricksEngineAdapter.can_access_spark_session(
+            bool(values.get("disable_spark_session"))
+        ):
             return values
         server_hostname, http_path, access_token = (
             values.get("server_hostname"),
@@ -554,7 +555,9 @@ def _databricks_connect_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str
             raise ValueError(
                 "`server_hostname`, `http_path`, and `access_token` are required for Databricks connections when not running in a notebook"
             )
-        if DatabricksEngineAdapter.can_access_spark_session:
+        if DatabricksEngineAdapter.can_access_databricks_connect(
+            bool(values.get("disable_databricks_connect"))
+        ):
             if not values.get("databricks_connect_server_hostname"):
                 values["databricks_connect_server_hostname"] = f"https://{server_hostname}"
             if not values.get("databricks_connect_access_token"):
@@ -585,14 +588,18 @@ def _extra_engine_config(self) -> t.Dict[str, t.Any]:
         return {
             k: v
             for k, v in self.dict().items()
-            if k.startswith("databricks_connect_") or k in ("catalog", "disable_databricks_connect")
+            if k.startswith("databricks_connect_")
+            or k in ("catalog", "disable_databricks_connect", "disable_spark_session")
         }
 
     @property
     def use_spark_session_only(self) -> bool:
-        from sqlmesh import RuntimeEnv
+        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 
-        return RuntimeEnv.get().is_databricks or self.force_databricks_connect
+        return (
+            DatabricksEngineAdapter.can_access_spark_session(self.disable_spark_session)
+            or self.force_databricks_connect
+        )
 
     @property
     def _connection_factory(self) -> t.Callable:
@@ -607,14 +614,14 @@ def _connection_factory(self) -> t.Callable:
 
     @property
     def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
-        from sqlmesh import RuntimeEnv
+        from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 
         if not self.use_spark_session_only:
             return {
                 "_user_agent_entry": "sqlmesh",
             }
 
-        if RuntimeEnv.get().is_databricks:
+        if DatabricksEngineAdapter.can_access_spark_session(self.disable_spark_session):
             from pyspark.sql import SparkSession
 
             return dict(

sqlmesh/core/engine_adapter/databricks.py

@@ -14,7 +14,6 @@
 )
 from sqlmesh.core.engine_adapter.spark import SparkEngineAdapter
 from sqlmesh.core.schema_diff import SchemaDiffer
-from sqlmesh.utils import classproperty
 from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
@@ -47,12 +46,20 @@ def __init__(self, *args: t.Any, **kwargs: t.Any):
         super().__init__(*args, **kwargs)
         self._spark: t.Optional[PySparkSession] = None
 
-    @classproperty
-    def can_access_spark_session(cls) -> bool:
+    @classmethod
+    def can_access_spark_session(cls, disable_spark_session: bool) -> bool:
         from sqlmesh import RuntimeEnv
 
-        if RuntimeEnv.get().is_databricks:
-            return True
+        if disable_spark_session:
+            return False
+
+        return RuntimeEnv.get().is_databricks
+
+    @classmethod
+    def can_access_databricks_connect(cls, disable_databricks_connect: bool) -> bool:
+        if disable_databricks_connect:
+            return False
+
         try:
             from databricks.connect import DatabricksSession  # noqa
 
@@ -62,19 +69,15 @@ def can_access_spark_session(cls) -> bool:
 
     @property
     def _use_spark_session(self) -> bool:
-        from sqlmesh import RuntimeEnv
-
-        if RuntimeEnv.get().is_databricks:
+        if self.can_access_spark_session(bool(self._extra_config.get("disable_spark_session"))):
             return True
-        return (
-            self.can_access_spark_session
-            and {
-                "databricks_connect_server_hostname",
-                "databricks_connect_access_token",
-                "databricks_connect_cluster_id",
-            }.issubset(self._extra_config)
-            and not self._extra_config.get("disable_databricks_connect")
-        )
+        return self.can_access_databricks_connect(
+            bool(self._extra_config.get("disable_databricks_connect"))
+        ) and {
+            "databricks_connect_server_hostname",
+            "databricks_connect_access_token",
+            "databricks_connect_cluster_id",
+        }.issubset(self._extra_config)
 
     @property
     def is_spark_session_cursor(self) -> bool:
@@ -97,11 +100,15 @@ def spark(self) -> PySparkSession:
         from databricks.connect import DatabricksSession
 
         if self._spark is None:
-            self._spark = DatabricksSession.builder.remote(
-                host=self._extra_config["databricks_connect_server_hostname"],
-                token=self._extra_config["databricks_connect_access_token"],
-                cluster_id=self._extra_config["databricks_connect_cluster_id"],
-            ).getOrCreate()
+            self._spark = (
+                DatabricksSession.builder.remote(
+                    host=self._extra_config["databricks_connect_server_hostname"],
+                    token=self._extra_config["databricks_connect_access_token"],
+                    cluster_id=self._extra_config["databricks_connect_cluster_id"],
+                )
+                .userAgent("sqlmesh")
+                .getOrCreate()
+            )
             catalog = self._extra_config.get("catalog")
             if catalog:
                 self.set_current_catalog(catalog)

sqlmesh/engines/spark/db_api/spark_session.py

@@ -89,8 +89,12 @@ def set_current_catalog(self, catalog_name: str) -> None:
     def cursor(self) -> SparkSessionCursor:
         try:
             self.spark.sparkContext.setLocalProperty("spark.scheduler.pool", f"pool_{get_ident()}")
+            self.spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
+            self.spark.conf.set("hive.exec.dynamic.partition", "true")
+            self.spark.conf.set("hive.exec.dynamic.partition.mode", "nonstrict")
         except NotImplementedError:
-            # Databricks Connect does not support accessing the SparkContext
+            # Databricks Connect does not support accessing the SparkContext nor does it support
+            # setting dynamic partition overwrite since it uses replace where
             pass
         if self.catalog:
             from py4j.protocol import Py4JError
@@ -101,9 +105,6 @@ def cursor(self) -> SparkSessionCursor:
             # and shared clusters so we use the Databricks Unity only SQL command instead
             except Py4JError:
                 self.spark.sql(f"USE CATALOG {self.catalog}")
-        self.spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
-        self.spark.conf.set("hive.exec.dynamic.partition", "true")
-        self.spark.conf.set("hive.exec.dynamic.partition.mode", "nonstrict")
         return SparkSessionCursor(self.spark)
 
     def commit(self) -> None: