Allow for spark session provider override

Author: markdlv-sf

src/datacustomcode/client.py

@@ -24,6 +24,8 @@
 from pyspark.sql import SparkSession
 
 from datacustomcode.config import SparkConfig, config
+from datacustomcode.spark.base import BaseSparkSessionProvider
+from datacustomcode.spark.default import DefaultSparkSessionProvider
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.io.reader.base import BaseDataCloudReader
 
@@ -36,18 +38,6 @@
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
 
 
-def _setup_spark(spark_config: SparkConfig) -> SparkSession:
-    """Setup Spark session from config."""
-    builder = SparkSession.builder
-    if spark_config.master is not None:
-        builder = builder.master(spark_config.master)
-
-    builder = builder.appName(spark_config.app_name)
-    for key, value in spark_config.options.items():
-        builder = builder.config(key, value)
-    return builder.getOrCreate()
-
-
 class DataCloudObjectType(Enum):
     DLO = "dlo"
     DMO = "dmo"
@@ -124,6 +114,7 @@ def __new__(
         cls,
         reader: Optional[BaseDataCloudReader] = None,
         writer: Optional[BaseDataCloudWriter] = None,
+        spark_provider: Optional[BaseSparkSessionProvider] = None,
     ) -> Client:
         if cls._instance is None:
             cls._instance = super().__new__(cls)
@@ -136,7 +127,15 @@ def __new__(
                     raise ValueError(
                         "Spark config is required when reader/writer is not provided"
                     )
-                spark = _setup_spark(config.spark_config)
+                provider: BaseSparkSessionProvider
+                if spark_provider is not None:
+                    provider = spark_provider
+                elif config.spark_provider_config is not None:
+                    provider = config.spark_provider_config.to_object()
+                else:
+                    provider = DefaultSparkSessionProvider()
+
+                spark = provider.get_session(config.spark_config)
 
             if config.reader_config is None and reader is None:
                 raise ValueError(

src/datacustomcode/config.py

@@ -36,6 +36,8 @@
 # This lets all readers and writers to be findable via config
 from datacustomcode.io import *  # noqa: F403
 from datacustomcode.io.base import BaseDataAccessLayer
+from datacustomcode.spark.base import BaseSparkSessionProvider
+from datacustomcode.spark.default import DefaultSparkSessionProvider
 from datacustomcode.io.reader.base import BaseDataCloudReader  # noqa: TCH001
 from datacustomcode.io.writer.base import BaseDataCloudWriter  # noqa: TCH001
 
@@ -89,10 +91,25 @@ class SparkConfig(ForceableConfig):
     )
 
 
+_P = TypeVar("_P", bound=BaseSparkSessionProvider)
+
+
+class SparkProviderConfig(ForceableConfig, Generic[_P]):
+    model_config = ConfigDict(validate_default=True, extra="forbid")
+    type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
+    type_config_name: str = Field(description="CONFIG_NAME of the Spark session provider.")
+    options: dict[str, Any] = Field(default_factory=dict)
+
+    def to_object(self) -> _P:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_P, type_(**self.options))
+
+
 class ClientConfig(BaseModel):
     reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
     writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
     spark_config: Union[SparkConfig, None] = None
+    spark_provider_config: Union[SparkProviderConfig[BaseSparkSessionProvider], None] = None
 
     def update(self, other: ClientConfig) -> ClientConfig:
         """Merge this ClientConfig with another, respecting force flags.
@@ -117,6 +134,7 @@ def merge(
         self.reader_config = merge(self.reader_config, other.reader_config)
         self.writer_config = merge(self.writer_config, other.writer_config)
         self.spark_config = merge(self.spark_config, other.spark_config)
+        self.spark_provider_config = merge(self.spark_provider_config, other.spark_provider_config)
         return self
 
     def load(self, config_path: str) -> ClientConfig:

src/datacustomcode/spark/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+
+from datacustomcode.spark.base import BaseSparkSessionProvider
+from datacustomcode.spark.default import DefaultSparkSessionProvider
+
+__all__ = ["BaseSparkSessionProvider", "DefaultSparkSessionProvider"]
+

src/datacustomcode/spark/base.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+    from datacustomcode.config import SparkConfig
+
+
+class BaseSparkSessionProvider(UserExtendableNamedConfigMixin):
+    def get_session(self, spark_config: SparkConfig) -> "SparkSession":
+        raise NotImplementedError
+
+

src/datacustomcode/spark/default.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from datacustomcode.spark.base import BaseSparkSessionProvider
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+    from datacustomcode.config import SparkConfig
+
+
+class DefaultSparkSessionProvider(BaseSparkSessionProvider):
+    CONFIG_NAME = "DefaultSparkSessionProvider"
+
+    def get_session(self, spark_config: SparkConfig) -> "SparkSession":
+        from pyspark.sql import SparkSession
+
+        builder = SparkSession.builder
+        if spark_config.master is not None:
+            builder = builder.master(spark_config.master)
+        builder = builder.appName(spark_config.app_name)
+        for key, value in spark_config.options.items():
+            builder = builder.config(key, value)
+        return builder.getOrCreate()
+
+

tests/test_client.py

@@ -9,7 +9,6 @@
     Client,
     DataCloudAccessLayerException,
     DataCloudObjectType,
-    _setup_spark,
 )
 from datacustomcode.config import (
     AccessLayerObjectConfig,
@@ -100,37 +99,41 @@ def test_singleton_pattern(self, reset_client, mock_spark):
             Client(reader=MagicMock(spec=BaseDataCloudReader))
 
     @patch("datacustomcode.client.config")
-    @patch("datacustomcode.client._setup_spark")
     def test_initialization_with_config(
-        self, mock_setup_spark, mock_config, reset_client, mock_spark
+        self, mock_config, reset_client, mock_spark
     ):
         """Test client initialization using configuration."""
-        mock_setup_spark.return_value = mock_spark
+        from datacustomcode.spark.default import DefaultSparkSessionProvider
+        from unittest.mock import patch as mock_patch
 
-        mock_reader = MagicMock(spec=BaseDataCloudReader)
-        mock_reader_config = MagicMock()
-        mock_reader_config.to_object.return_value = mock_reader
-        mock_reader_config.force = False
+        with mock_patch.object(DefaultSparkSessionProvider, "get_session") as mock_get_session:
+            mock_get_session.return_value = mock_spark
 
-        mock_writer = MagicMock(spec=BaseDataCloudWriter)
-        mock_writer_config = MagicMock()
-        mock_writer_config.to_object.return_value = mock_writer
-        mock_writer_config.force = False
+            mock_reader = MagicMock(spec=BaseDataCloudReader)
+            mock_reader_config = MagicMock()
+            mock_reader_config.to_object.return_value = mock_reader
+            mock_reader_config.force = False
 
-        mock_spark_config = MagicMock(spec=SparkConfig)
+            mock_writer = MagicMock(spec=BaseDataCloudWriter)
+            mock_writer_config = MagicMock()
+            mock_writer_config.to_object.return_value = mock_writer
+            mock_writer_config.force = False
 
-        mock_config.reader_config = mock_reader_config
-        mock_config.writer_config = mock_writer_config
-        mock_config.spark_config = mock_spark_config
+            mock_spark_config = MagicMock(spec=SparkConfig)
+            mock_config.spark_provider_config = None
 
-        client = Client()
+            mock_config.reader_config = mock_reader_config
+            mock_config.writer_config = mock_writer_config
+            mock_config.spark_config = mock_spark_config
 
-        mock_setup_spark.assert_called_once_with(mock_spark_config)
-        mock_reader_config.to_object.assert_called_once_with(mock_spark)
-        mock_writer_config.to_object.assert_called_once_with(mock_spark)
+            client = Client()
 
-        assert client._reader is mock_reader
-        assert client._writer is mock_writer
+            mock_get_session.assert_called_once_with(mock_spark_config)
+            mock_reader_config.to_object.assert_called_once_with(mock_spark)
+            mock_writer_config.to_object.assert_called_once_with(mock_spark)
+
+            assert client._reader is mock_reader
+            assert client._writer is mock_writer
 
     def test_read_dlo(self, reset_client, mock_spark):
         reader = MagicMock(spec=BaseDataCloudReader)
@@ -249,12 +252,12 @@ def test_read_pattern_flow(self, reset_client, mock_spark):
         assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO]
 
 
-# Add tests for _setup_spark function
-class TestSetupSpark:
+# Add tests for DefaultSparkSessionProvider
+class TestDefaultSparkSessionProvider:
 
-    @patch("datacustomcode.client.SparkSession")
-    def test_setup_spark_with_master(self, mock_spark_session):
-        """Test _setup_spark with master specified"""
+    @patch("pyspark.sql.SparkSession")
+    def test_get_session_with_master(self, mock_spark_session):
+        """Test DefaultSparkSessionProvider with master specified"""
         mock_builder = MagicMock()
         mock_master_builder = MagicMock()
         mock_app_name_builder = MagicMock()
@@ -273,7 +276,9 @@ def test_setup_spark_with_master(self, mock_spark_session):
             options={"spark.executor.memory": "1g"},
         )
 
-        result = _setup_spark(spark_config)
+        from datacustomcode.spark.default import DefaultSparkSessionProvider
+        provider = DefaultSparkSessionProvider()
+        result = provider.get_session(spark_config)
 
         mock_builder.master.assert_called_once_with("local[1]")
         mock_master_builder.appName.assert_called_once_with("test-app")
@@ -283,9 +288,9 @@ def test_setup_spark_with_master(self, mock_spark_session):
         mock_config_builder.getOrCreate.assert_called_once()
         assert result is mock_session
 
-    @patch("datacustomcode.client.SparkSession")
-    def test_setup_spark_with_multiple_options(self, mock_spark_session):
-        """Test _setup_spark with multiple config options"""
+    @patch("pyspark.sql.SparkSession")
+    def test_get_session_with_multiple_options(self, mock_spark_session):
+        """Test DefaultSparkSessionProvider with multiple config options"""
         mock_builder = MagicMock()
         mock_app_name_builder = MagicMock()
         mock_config_builder1 = MagicMock()
@@ -310,7 +315,9 @@ def test_setup_spark_with_multiple_options(self, mock_spark_session):
             },
         )
 
-        result = _setup_spark(spark_config)
+        from datacustomcode.spark.default import DefaultSparkSessionProvider
+        provider = DefaultSparkSessionProvider()
+        result = provider.get_session(spark_config)
 
         mock_builder.appName.assert_called_once_with("test-app")