From 3da41b60d2e64a24d724c8918bf52580ecbcb2c7 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Thu, 2 Jul 2026 15:11:25 -0400 Subject: [PATCH 1/5] rely on config for name source --- CHANGELOG.md | 4 +- README.md | 7 +-- src/datacustomcode/client.py | 31 +++++++------ src/datacustomcode/io/reader/base.py | 22 +--------- .../examples/streaming_deltas/entrypoint.py | 4 +- tests/io/reader/test_query_api.py | 4 +- tests/test_client.py | 44 +++++++++++++++---- 7 files changed, 66 insertions(+), 50 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 93b82ae..80699ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,11 +8,11 @@ New methods let an entry point process a Data Lake Object's Change Data Feed continuously instead of reading a bounded snapshot: - - `read_dlo_deltas(name)` / `read_dmo_deltas(name)` – return a streaming DataFrame over the object's change feed. + - `read_dlo_deltas()` / `read_dmo_deltas()` – return a streaming DataFrame over the object's change feed. - `write_dlo_deltas(name, dataframe)` – start a streaming query that writes each micro-batch to the target DLO and return the `StreamingQuery` handle. ```python - deltas = client.read_dlo_deltas("Input__dll") + deltas = client.read_dlo_deltas() transformed = deltas.withColumn("description__c", upper(col("description__c"))) query = client.write_dlo_deltas("Output__dll", transformed) query.awaitTermination() diff --git a/README.md b/README.md index 789eef7..16c8394 100644 --- a/README.md +++ b/README.md @@ -155,8 +155,8 @@ You should only need the following methods: * `write_to_dmo(name, spark_dataframe, write_mode)` – Write to a Data Lake Object by name with a Spark dataframe For streaming (delta) transforms, the streaming counterparts are: -* `read_dlo_deltas(name)` – Read the streaming change feed (deltas) of a Data Lake Object as a streaming DataFrame -* `read_dmo_deltas(name)` – Read the streaming change feed (deltas) of a Data Model Object as a streaming DataFrame +* `read_dlo_deltas()` – Read the streaming change feed (deltas) of a Data Lake Object as a streaming DataFrame. +* `read_dmo_deltas()` – Read the streaming change feed (deltas) of a Data Model Object as a streaming DataFrame. * `write_dlo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` For example: @@ -186,7 +186,8 @@ from datacustomcode import Client client = Client() # read_dlo_deltas returns a *streaming* DataFrame over the change feed. -deltas = client.read_dlo_deltas("Input__dll") +# The runtime resolves the single streaming source, so no name is passed. +deltas = client.read_dlo_deltas() # Ordinary PySpark transform. transformed = deltas.withColumn("description__c", upper(col("description__c"))) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 143796d..f0f4529 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -15,6 +15,7 @@ from __future__ import annotations from enum import Enum +import os from typing import ( TYPE_CHECKING, Any, @@ -45,6 +46,15 @@ from datacustomcode.spark.base import BaseSparkSessionProvider +_STREAMING_SOURCE_ENV = "BYOC_STREAMING_SOURCE_NAME" +_STREAMING_SOURCE_FALLBACK = "" + + +def _streaming_source_name() -> str: + """Return the runtime streaming source name, or a readable fallback.""" + return os.environ.get(_STREAMING_SOURCE_ENV, _STREAMING_SOURCE_FALLBACK) + + def _build_spark_llm_gateway() -> "SparkLLMGateway": """Instantiate the SDK-configured :class:`SparkLLMGateway`. @@ -336,7 +346,7 @@ def read_dmo(self, name: str) -> PySparkDataFrame: self._record_dmo_access(name) return self._reader.read_dmo(name) # type: ignore[no-any-return] - def read_dlo_deltas(self, name: str) -> PySparkDataFrame: + def read_dlo_deltas(self) -> PySparkDataFrame: """Read the streaming change feed (deltas) for a DLO from Data Cloud. Streaming counterpart to :meth:`read_dlo`, for use in a streaming @@ -345,29 +355,24 @@ def read_dlo_deltas(self, name: str) -> PySparkDataFrame: ``_commit_*``) alongside the source columns. Pair with :meth:`write_dlo_deltas` to write the transformed stream back to a DLO. - Args: - name: The name of the DLO to read deltas from. - Returns: A streaming PySpark DataFrame over the DLO change feed. """ - self._record_dlo_access(name) - return self._reader.read_dlo_deltas(name) # type: ignore[no-any-return] + self._record_dlo_access(_streaming_source_name()) + return self._reader.read_dlo_deltas() # type: ignore[no-any-return] - def read_dmo_deltas(self, name: str) -> PySparkDataFrame: + def read_dmo_deltas(self) -> PySparkDataFrame: """Read the streaming change feed (deltas) for a DMO from Data Cloud. Streaming counterpart to :meth:`read_dmo`. See :meth:`read_dlo_deltas` - for the shape of the returned change feed. - - Args: - name: The name of the DMO to read deltas from. + for the shape of the returned change feed and why no source name is + passed. Returns: A streaming PySpark DataFrame over the DMO change feed. """ - self._record_dmo_access(name) - return self._reader.read_dmo_deltas(name) # type: ignore[no-any-return] + self._record_dmo_access(_streaming_source_name()) + return self._reader.read_dmo_deltas() # type: ignore[no-any-return] def write_to_dlo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs diff --git a/src/datacustomcode/io/reader/base.py b/src/datacustomcode/io/reader/base.py index 2b69ed4..c19769b 100644 --- a/src/datacustomcode/io/reader/base.py +++ b/src/datacustomcode/io/reader/base.py @@ -42,11 +42,7 @@ def read_dmo( schema: Union[AtomicType, StructType, str, None] = None, ) -> PySparkDataFrame: ... - def read_dlo_deltas( - self, - name: str, - schema: Union[AtomicType, StructType, str, None] = None, - ) -> PySparkDataFrame: + def read_dlo_deltas(self) -> PySparkDataFrame: """Read the streaming change feed (deltas) for a Data Lake Object. This is the streaming counterpart to :meth:`read_dlo`. It returns a @@ -56,11 +52,6 @@ def read_dlo_deltas( base implementation raises :class:`NotImplementedError` so local readers that do not support streaming fail clearly. - Args: - name: Data Lake Object name. - schema: Accepted for parity with :meth:`read_dlo`; implementations - may ignore it. - Returns: A streaming PySpark DataFrame over the DLO change feed. @@ -74,21 +65,12 @@ def read_dlo_deltas( "deltas." ) - def read_dmo_deltas( - self, - name: str, - schema: Union[AtomicType, StructType, str, None] = None, - ) -> PySparkDataFrame: + def read_dmo_deltas(self) -> PySparkDataFrame: """Read the streaming change feed (deltas) for a Data Model Object. Streaming counterpart to :meth:`read_dmo`. See :meth:`read_dlo_deltas` for behavior and the local-development caveat. - Args: - name: Data Model Object name. - schema: Accepted for parity with :meth:`read_dmo`; implementations - may ignore it. - Returns: A streaming PySpark DataFrame over the DMO change feed. diff --git a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py index 8919cf7..f8b078c 100644 --- a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -4,7 +4,7 @@ of ``read_dlo`` / ``write_to_dlo`` (which read and write a bounded snapshot), it uses the streaming delta methods: -* ``client.read_dlo_deltas(name)`` returns a *streaming* DataFrame over the +* ``client.read_dlo_deltas()`` returns a *streaming* DataFrame over the Change Data Feed of the source DLO. Each row carries the source columns plus change-feed metadata columns (``_record_type``, ``_commit_*``). * ``client.write_dlo_deltas(name, df)`` starts a streaming query that writes @@ -30,7 +30,7 @@ def main(): client = Client() # Streaming DataFrame over the source DLO's change feed. - deltas = client.read_dlo_deltas("Account_std__dll") + deltas = client.read_dlo_deltas() # Ordinary PySpark transform. transformed = deltas.withColumn("description__c", upper(col("description__c"))) diff --git a/tests/io/reader/test_query_api.py b/tests/io/reader/test_query_api.py index eaa8848..0b19081 100644 --- a/tests/io/reader/test_query_api.py +++ b/tests/io/reader/test_query_api.py @@ -227,13 +227,13 @@ def test_read_dlo( def test_read_dlo_deltas_not_supported_locally(self, reader_without_init): """Streaming delta reads are not supported by the local reader.""" with pytest.raises(NotImplementedError) as exc_info: - reader_without_init.read_dlo_deltas("test_dlo") + reader_without_init.read_dlo_deltas() assert "read_dlo_deltas" in str(exc_info.value) def test_read_dmo_deltas_not_supported_locally(self, reader_without_init): """Streaming delta reads are not supported by the local reader.""" with pytest.raises(NotImplementedError) as exc_info: - reader_without_init.read_dmo_deltas("test_dmo") + reader_without_init.read_dmo_deltas() assert "read_dmo_deltas" in str(exc_info.value) def test_read_dlo_with_schema( diff --git a/tests/test_client.py b/tests/test_client.py index e661cbc..a0cf4cd 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from unittest.mock import MagicMock, patch from pyspark.sql import DataFrame, SparkSession @@ -200,11 +201,32 @@ def test_read_dlo_deltas(self, reset_client, mock_spark): reader.read_dlo_deltas.return_value = mock_df client = Client(reader=reader, writer=writer) - result = client.read_dlo_deltas("test_dlo") + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("BYOC_STREAMING_SOURCE_NAME", None) + result = client.read_dlo_deltas() - reader.read_dlo_deltas.assert_called_once_with("test_dlo") + reader.read_dlo_deltas.assert_called_once_with() assert result is mock_df - assert "test_dlo" in client._data_layer_history[DataCloudObjectType.DLO] + assert ( + "" + in client._data_layer_history[DataCloudObjectType.DLO] + ) + + def test_read_dlo_deltas_records_runtime_source_name( + self, reset_client, mock_spark + ): + """The runtime source env var populates the access-history entry.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + reader.read_dlo_deltas.return_value = MagicMock(spec=DataFrame) + + client = Client(reader=reader, writer=writer) + with patch.dict( + "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_std__dll"} + ): + client.read_dlo_deltas() + + assert "Account_std__dll" in client._data_layer_history[DataCloudObjectType.DLO] def test_read_dmo_deltas(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) @@ -213,11 +235,16 @@ def test_read_dmo_deltas(self, reset_client, mock_spark): reader.read_dmo_deltas.return_value = mock_df client = Client(reader=reader, writer=writer) - result = client.read_dmo_deltas("test_dmo") + with patch.dict( + "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_model__dlm"} + ): + result = client.read_dmo_deltas() - reader.read_dmo_deltas.assert_called_once_with("test_dmo") + reader.read_dmo_deltas.assert_called_once_with() assert result is mock_df - assert "test_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + assert ( + "Account_model__dlm" in client._data_layer_history[DataCloudObjectType.DMO] + ) def test_write_dlo_deltas(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) @@ -262,10 +289,11 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): client = Client(reader=reader, writer=writer) - df = client.read_dlo_deltas("source_dll") + with patch.dict("os.environ", {"BYOC_STREAMING_SOURCE_NAME": "source_dll"}): + df = client.read_dlo_deltas() client.write_dlo_deltas("target_dll", df) - reader.read_dlo_deltas.assert_called_once_with("source_dll") + reader.read_dlo_deltas.assert_called_once_with() writer.write_dlo_deltas.assert_called_once_with("target_dll", stream_df) assert "source_dll" in client._data_layer_history[DataCloudObjectType.DLO] From 2c3d2858d5bcd25ea9004b030bca5f0da6a9ff82 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Tue, 7 Jul 2026 13:03:57 -0400 Subject: [PATCH 2/5] separate streaming client --- CHANGELOG.md | 9 +- README.md | 14 +- src/datacustomcode/__init__.py | 5 + src/datacustomcode/client.py | 352 ++++++++++-------- .../examples/streaming_deltas/entrypoint.py | 9 +- tests/test_client.py | 260 +++++++++---- 6 files changed, 419 insertions(+), 230 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80699ea..fb41eee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,14 +4,19 @@ ### Added -- **Streaming (delta) read/write methods on `Client` for BYOC streaming transforms.** +- **`StreamingClient` for BYOC streaming (delta) transforms.** - New methods let an entry point process a Data Lake Object's Change Data Feed continuously instead of reading a bounded snapshot: + A dedicated `StreamingClient` (alongside the batch `Client`) lets an entry point process a Data Lake Object's Change Data Feed continuously instead of reading a bounded snapshot. - `read_dlo_deltas()` / `read_dmo_deltas()` – return a streaming DataFrame over the object's change feed. - `write_dlo_deltas(name, dataframe)` – start a streaming query that writes each micro-batch to the target DLO and return the `StreamingQuery` handle. + The shared functions (`find_file_path`, `llm_gateway_generate_text`, `einstein_predict`) are available on both `Client` and `StreamingClient`. + ```python + from datacustomcode import StreamingClient + + client = StreamingClient() deltas = client.read_dlo_deltas() transformed = deltas.withColumn("description__c", upper(col("description__c"))) query = client.write_dlo_deltas("Output__dll", transformed) diff --git a/README.md b/README.md index 16c8394..a5d77fc 100644 --- a/README.md +++ b/README.md @@ -145,20 +145,22 @@ Your Python dependencies can be packaged as .py files, .zip archives (containing ## API -Your entry point script will define logic using the `Client` object which wraps data access layers. +Your entry point script will define logic using the `Client` object (for batch transforms) or the `StreamingClient` object (for streaming delta transforms), which wrap the data access layers. Both are singletons; a single transform should use one or the other, not both. -You should only need the following methods: +For a batch transform, use `Client`. You should only need the following methods: * `find_file_path(file_name)` – Resolve a bundled file (placed under `payload/files/`) to a `pathlib.Path` that exists. Works the same locally and inside Data Cloud — see [Bundled file resolution](#bundled-file-resolution) below for the full lookup order. Raises `FileNotFoundError` if the file isn't found. * `read_dlo(name)` – Read from a Data Lake Object by name * `read_dmo(name)` – Read from a Data Model Object by name * `write_to_dlo(name, spark_dataframe, write_mode)` – Write to a Data Model Object by name with a Spark dataframe * `write_to_dmo(name, spark_dataframe, write_mode)` – Write to a Data Lake Object by name with a Spark dataframe -For streaming (delta) transforms, the streaming counterparts are: +For a streaming (delta) transform, use `StreamingClient`, which exposes the streaming counterparts: * `read_dlo_deltas()` – Read the streaming change feed (deltas) of a Data Lake Object as a streaming DataFrame. * `read_dmo_deltas()` – Read the streaming change feed (deltas) of a Data Model Object as a streaming DataFrame. * `write_dlo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` +`find_file_path`, `llm_gateway_generate_text`, and `einstein_predict` are available on both clients. + For example: ```python from datacustomcode import Client @@ -176,14 +178,14 @@ client.write_to_dlo('output_DLO') ### Streaming (delta) transforms -Streaming BYOC transforms process a Data Lake Object's Change Data Feed continuously instead of reading a bounded snapshot. Use the `*_deltas` methods in place of the batch read/write methods: +Streaming BYOC transforms process a Data Lake Object's Change Data Feed continuously instead of reading a bounded snapshot. Use a `StreamingClient` and its `*_deltas` methods in place of the batch `Client` read/write methods: ```python from pyspark.sql.functions import col, upper -from datacustomcode import Client +from datacustomcode import StreamingClient -client = Client() +client = StreamingClient() # read_dlo_deltas returns a *streaming* DataFrame over the change feed. # The runtime resolves the single streaming source, so no name is passed. diff --git a/src/datacustomcode/__init__.py b/src/datacustomcode/__init__.py index be123ff..4cd56f5 100644 --- a/src/datacustomcode/__init__.py +++ b/src/datacustomcode/__init__.py @@ -23,6 +23,7 @@ "QueryAPIDataCloudReader", "SparkEinsteinPredictions", "SparkLLMGateway", + "StreamingClient", "einstein_predict_col", "llm_gateway_generate_text_col", ] @@ -34,6 +35,10 @@ def __getattr__(name: str): from datacustomcode.client import Client return Client + elif name == "StreamingClient": + from datacustomcode.client import StreamingClient + + return StreamingClient elif name == "AuthType": from datacustomcode.credentials import AuthType diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index f0f4529..8392ef9 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -22,7 +22,9 @@ ClassVar, Dict, Optional, + TypeVar, Union, + cast, ) from datacustomcode.config import config @@ -35,7 +37,11 @@ if TYPE_CHECKING: from pathlib import Path - from pyspark.sql import Column, DataFrame as PySparkDataFrame + from pyspark.sql import ( + Column, + DataFrame as PySparkDataFrame, + SparkSession, + ) from pyspark.sql.streaming import StreamingQuery from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions @@ -47,12 +53,33 @@ _STREAMING_SOURCE_ENV = "BYOC_STREAMING_SOURCE_NAME" -_STREAMING_SOURCE_FALLBACK = "" def _streaming_source_name() -> str: - """Return the runtime streaming source name, or a readable fallback.""" - return os.environ.get(_STREAMING_SOURCE_ENV, _STREAMING_SOURCE_FALLBACK) + """Return the runtime streaming source name. + + Raises: + RuntimeError: If ``BYOC_STREAMING_SOURCE_NAME`` is not set + """ + source = os.environ.get(_STREAMING_SOURCE_ENV) + if not source: + raise RuntimeError(f"{_STREAMING_SOURCE_ENV} is not set.") + return source + + +def _active_client() -> "_BaseClient": + """Return the client backing the module-level Spark column helpers. + + Prefers an already-initialized singleton so a streaming job reuses its + :class:`StreamingClient` (and a batch job its :class:`Client`) rather than + forcing an unrelated client into existence. Falls back to building the + batch :class:`Client` when neither has been created yet. + """ + if Client._instance is not None: + return Client._instance + if StreamingClient._instance is not None: + return StreamingClient._instance + return Client() def _build_spark_llm_gateway() -> "SparkLLMGateway": @@ -110,7 +137,7 @@ def llm_gateway_generate_text_col( the generated text; on failure, ``status == "ERROR"`` and the ``error_*`` fields carry diagnostic detail. """ - gateway = Client()._get_spark_llm_gateway() + gateway = _active_client()._get_spark_llm_gateway() return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id) @@ -172,7 +199,7 @@ def einstein_predict_col( the JSON-serialized prediction payload; on failure, ``status == "ERROR"`` and the ``error_*`` fields carry diagnostic detail. """ - predictions = Client()._get_spark_einstein_predictions() + predictions = _active_client()._get_spark_einstein_predictions() return predictions.einstein_predict_col( model_api_name, prediction_type, features, settings=settings ) @@ -216,39 +243,39 @@ def __str__(self) -> str: return msg -class Client: - """Entrypoint for accessing DataCloud objects. +_ClientT = TypeVar("_ClientT", bound="_BaseClient") + + +class _BaseClient: + """Shared machinery for the Data Cloud client singletons. - This is the object used to access Data Cloud DLOs and DMOs. Accessing DLOs/DMOs - are tracked and will throw an exception if they are mixed. In other words, you - can read from DLOs and write to DLOs, read from DMOs and write to DMOs, but you - cannot read from DLOs and write to DMOs or read from DMOs and write to DLOs. - Furthermore you cannot mix during merging tables. This class is a singleton to - prevent accidental mixing of DLOs and DMOs. + Holds the wiring common to :class:`Client` (batch) and + :class:`StreamingClient` - You can provide custom readers and writers to the client for advanced use - cases, but this is not recommended for testing as they may result in unexpected - behavior once deployed to Data Cloud. By default, the client intercepts all - read/write operations and mocks access to Data Cloud. For example, during - writing, we print to the console instead of writing to Data Cloud. + This base class is not meant to be instantiated directly; use + :class:`Client` or :class:`StreamingClient`. Args: - finder: Find a file path reader: A custom reader to use for reading Data Cloud objects. writer: A custom writer to use for writing Data Cloud objects. + spark_provider: Optional custom :class:`BaseSparkSessionProvider`. spark_llm_gateway: Optional custom :class:`SparkLLMGateway`. spark_einstein_predictions: Optional custom :class:`SparkEinsteinPredictions`. - - Example: - >>> client = Client() - >>> file_path = client.find_file_path("data.csv") - >>> dlo = client.read_dlo("my_dlo") - >>> client.write_to_dmo("my_dmo", dlo) - >>> answer = client.llm_gateway_generate_text("Generate a greeting message") """ - _instance: ClassVar[Optional[Client]] = None + # Each concrete subclass gets its own ``_instance`` slot: reads fall through + # to this base default of ``None``, but ``cls._instance = ...`` in __new__ + # always writes to the subclass, so ``Client`` and ``StreamingClient`` never + # share an instance. + _instance: ClassVar[Optional[_BaseClient]] = None + # Process-wide Spark session shared across BOTH client types. Unlike + # ``_instance``, this is written via ``_BaseClient._shared_spark`` (never + # ``cls._shared_spark``), so the slot lives on the base class and a + # ``Client`` and a ``StreamingClient`` in the same process reuse one session + # — and therefore one underlying connection — instead of opening two + # containing differing state + _shared_spark: ClassVar[Optional[SparkSession]] = None _reader: BaseDataCloudReader _writer: BaseDataCloudWriter _file: DefaultFindFilePath @@ -258,37 +285,44 @@ class Client: _code_type: str def __new__( - cls, + cls: type[_ClientT], reader: Optional[BaseDataCloudReader] = None, writer: Optional[BaseDataCloudWriter] = None, spark_provider: Optional[BaseSparkSessionProvider] = None, spark_llm_gateway: Optional[SparkLLMGateway] = None, spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None, code_type: str = "script", - ) -> Client: + ) -> _ClientT: if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._spark_llm_gateway = spark_llm_gateway - cls._instance._spark_einstein_predictions = spark_einstein_predictions + instance = super().__new__(cls) + instance._spark_llm_gateway = spark_llm_gateway + instance._spark_einstein_predictions = spark_einstein_predictions # Initialize Readers and Writers from config # and/or provided reader and writer if reader is None or writer is None: - # We need a spark because we will initialize readers and writers - if config.spark_config is None: - raise ValueError( - "Spark config is required when reader/writer is not provided" - ) - - 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() + # We need a spark because we will initialize readers and writers. + # Reuse the process-wide session if one client already built it, + # so a Client and a StreamingClient share a single connection. + if _BaseClient._shared_spark is not None: + spark = _BaseClient._shared_spark else: - provider = DefaultSparkSessionProvider() - - spark = provider.get_session(config.spark_config) + if config.spark_config is None: + raise ValueError( + "Spark config is required when reader/writer is not " + "provided" + ) + + 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) + _BaseClient._shared_spark = spark if config.reader_config is None and reader is None: raise ValueError( @@ -311,115 +345,17 @@ def __new__( else: writer_init = writer - cls._instance._reader = reader_init - cls._instance._writer = writer_init - cls._instance._file = DefaultFindFilePath() - cls._instance._data_layer_history = { + instance._reader = reader_init + instance._writer = writer_init + instance._file = DefaultFindFilePath() + instance._data_layer_history = { DataCloudObjectType.DLO: set(), DataCloudObjectType.DMO: set(), } - elif (reader is not None or writer is not None) and cls._instance is not None: + cls._instance = instance + elif reader is not None or writer is not None: raise ValueError("Cannot set reader or writer after client is initialized") - return cls._instance - - def read_dlo(self, name: str) -> PySparkDataFrame: - """Read a DLO from Data Cloud. - - Args: - name: The name of the DLO to read. - - Returns: - A PySpark DataFrame containing the DLO data. - """ - self._record_dlo_access(name) - return self._reader.read_dlo(name) # type: ignore[no-any-return] - - def read_dmo(self, name: str) -> PySparkDataFrame: - """Read a DMO from Data Cloud. - - Args: - name: The name of the DMO to read. - - Returns: - A PySpark DataFrame containing the DMO data. - """ - self._record_dmo_access(name) - return self._reader.read_dmo(name) # type: ignore[no-any-return] - - def read_dlo_deltas(self) -> PySparkDataFrame: - """Read the streaming change feed (deltas) for a DLO from Data Cloud. - - Streaming counterpart to :meth:`read_dlo`, for use in a streaming - (``DELTA_SYNC``) BYOC transform. Returns a streaming DataFrame whose - rows carry the change-feed metadata columns (``_record_type``, - ``_commit_*``) alongside the source columns. Pair with - :meth:`write_dlo_deltas` to write the transformed stream back to a DLO. - - Returns: - A streaming PySpark DataFrame over the DLO change feed. - """ - self._record_dlo_access(_streaming_source_name()) - return self._reader.read_dlo_deltas() # type: ignore[no-any-return] - - def read_dmo_deltas(self) -> PySparkDataFrame: - """Read the streaming change feed (deltas) for a DMO from Data Cloud. - - Streaming counterpart to :meth:`read_dmo`. See :meth:`read_dlo_deltas` - for the shape of the returned change feed and why no source name is - passed. - - Returns: - A streaming PySpark DataFrame over the DMO change feed. - """ - self._record_dmo_access(_streaming_source_name()) - return self._reader.read_dmo_deltas() # type: ignore[no-any-return] - - def write_to_dlo( - self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs - ) -> None: - """Write a PySpark DataFrame to a DLO in Data Cloud. - - Args: - name: The name of the DLO to write to. - dataframe: The PySpark DataFrame to write. - write_mode: The write mode to use for writing to the DLO. - """ - self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) - return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] - - def write_to_dmo( - self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs - ) -> None: - """Write a PySpark DataFrame to a DMO in Data Cloud. - - Args: - name: The name of the DMO to write to. - dataframe: The PySpark DataFrame to write. - write_mode: The write mode to use for writing to the DMO. - """ - self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) - return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] - - def write_dlo_deltas( - self, name: str, dataframe: PySparkDataFrame, **kwargs - ) -> StreamingQuery: - """Write a streaming DataFrame of deltas to a DLO in Data Cloud. - - Streaming counterpart to :meth:`write_to_dlo`. Starts a streaming query - that writes each micro-batch to the target DLO and returns the - ``StreamingQuery`` handle; the caller typically calls - ``query.awaitTermination()``. The runtime owns the trigger and - checkpoint location. - - Args: - name: The name of the DLO to write to. - dataframe: The streaming PySpark DataFrame to write. - - Returns: - The started ``StreamingQuery``. - """ - self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) - return self._writer.write_dlo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return] + return cast(_ClientT, cls._instance) def find_file_path(self, file_name: str) -> Path: """Resolve a bundled file shipped in the package to an absolute path. @@ -547,3 +483,115 @@ def _record_dlo_access(self, name: str) -> None: def _record_dmo_access(self, name: str) -> None: self._data_layer_history[DataCloudObjectType.DMO].add(name) + + +class Client(_BaseClient): + """Entrypoint for batch access to Data Cloud objects. + + This is the object used to read and write bounded snapshots of Data Cloud + DLOs and DMOs. + """ + + _instance: ClassVar[Optional[Client]] = None + + def read_dlo(self, name: str) -> PySparkDataFrame: + """Read a DLO from Data Cloud. + + Args: + name: The name of the DLO to read. + + Returns: + A PySpark DataFrame containing the DLO data. + """ + self._record_dlo_access(name) + return self._reader.read_dlo(name) # type: ignore[no-any-return] + + def read_dmo(self, name: str) -> PySparkDataFrame: + """Read a DMO from Data Cloud. + + Args: + name: The name of the DMO to read. + + Returns: + A PySpark DataFrame containing the DMO data. + """ + self._record_dmo_access(name) + return self._reader.read_dmo(name) # type: ignore[no-any-return] + + def write_to_dlo( + self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs + ) -> None: + """Write a PySpark DataFrame to a DLO in Data Cloud. + + Args: + name: The name of the DLO to write to. + dataframe: The PySpark DataFrame to write. + write_mode: The write mode to use for writing to the DLO. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) + return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] + + def write_to_dmo( + self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs + ) -> None: + """Write a PySpark DataFrame to a DMO in Data Cloud. + + Args: + name: The name of the DMO to write to. + dataframe: The PySpark DataFrame to write. + write_mode: The write mode to use for writing to the DMO. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) + return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] + + +class StreamingClient(_BaseClient): + """Entrypoint for streaming (``DELTA_SYNC``) access to Data Cloud objects. + + This is the streaming counterpart to :class:`Client`. Instead of reading and + writing bounded snapshots, it reads a DLO/DMO change feed as a streaming + DataFrame and writes the transformed stream back via a ``StreamingQuery``. + """ + + _instance: ClassVar[Optional[StreamingClient]] = None + + def read_dlo_deltas(self) -> PySparkDataFrame: + """Read the streaming change feed (deltas) for a DLO from Data Cloud. + + For use in a streaming (``DELTA_SYNC``) BYOC transform. Returns a + streaming DataFrame whose rows carry the change-feed metadata columns + (``_record_type``, ``_commit_*``) alongside the source columns. + + Returns: + A streaming PySpark DataFrame over the DLO change feed. + """ + self._record_dlo_access(_streaming_source_name()) + return self._reader.read_dlo_deltas() # type: ignore[no-any-return] + + def read_dmo_deltas(self) -> PySparkDataFrame: + """Read the streaming change feed (deltas) for a DMO from Data Cloud. + + Returns: + A streaming PySpark DataFrame over the DMO change feed. + """ + self._record_dmo_access(_streaming_source_name()) + return self._reader.read_dmo_deltas() # type: ignore[no-any-return] + + def write_dlo_deltas( + self, name: str, dataframe: PySparkDataFrame, **kwargs + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a DLO in Data Cloud. + + Starts a streaming query that writes each micro-batch to the + target DLO and returns the ``StreamingQuery`` handle; the caller + typically calls ``query.awaitTermination()``. + + Args: + name: The name of the DLO to write to. + dataframe: The streaming PySpark DataFrame to write. + + Returns: + The started ``StreamingQuery``. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) + return self._writer.write_dlo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return] diff --git a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py index f8b078c..97dea40 100644 --- a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -1,8 +1,9 @@ """Streaming BYOC transform: read a DLO change feed and write the deltas back. This example is the streaming counterpart to a normal batch entrypoint. Instead -of ``read_dlo`` / ``write_to_dlo`` (which read and write a bounded snapshot), it -uses the streaming delta methods: +of a batch ``Client`` with ``read_dlo`` / ``write_to_dlo`` (which read and write +a bounded snapshot), it uses a :class:`StreamingClient` and its streaming delta +methods: * ``client.read_dlo_deltas()`` returns a *streaming* DataFrame over the Change Data Feed of the source DLO. Each row carries the source columns plus @@ -23,11 +24,11 @@ from pyspark.sql.functions import col, upper -from datacustomcode.client import Client +from datacustomcode.client import StreamingClient def main(): - client = Client() + client = StreamingClient() # Streaming DataFrame over the source DLO's change feed. deltas = client.read_dlo_deltas() diff --git a/tests/test_client.py b/tests/test_client.py index a0cf4cd..d020ebe 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -10,6 +10,8 @@ Client, DataCloudAccessLayerException, DataCloudObjectType, + StreamingClient, + _BaseClient, einstein_predict_col, llm_gateway_generate_text_col, ) @@ -81,10 +83,14 @@ def mock_config(mock_spark): @pytest.fixture def reset_client(): - """Reset the Client singleton between tests.""" + """Reset the client singletons (and the shared Spark session) between tests.""" Client._instance = None + StreamingClient._instance = None + _BaseClient._shared_spark = None yield Client._instance = None + StreamingClient._instance = None + _BaseClient._shared_spark = None class TestClient: @@ -194,47 +200,136 @@ def test_write_to_dmo(self, reset_client, mock_spark): "test_dmo", mock_df, WriteMode.OVERWRITE, extra_param=True ) - def test_read_dlo_deltas(self, reset_client, mock_spark): + def test_mixed_dlo_dmo_raises_exception(self, reset_client, mock_spark): + """Test that mixing DLOs and DMOs raises an exception.""" reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) mock_df = MagicMock(spec=DataFrame) - reader.read_dlo_deltas.return_value = mock_df client = Client(reader=reader, writer=writer) - with patch.dict("os.environ", {}, clear=False): - os.environ.pop("BYOC_STREAMING_SOURCE_NAME", None) - result = client.read_dlo_deltas() + client._record_dlo_access("test_dlo") - reader.read_dlo_deltas.assert_called_once_with() - assert result is mock_df - assert ( - "" - in client._data_layer_history[DataCloudObjectType.DLO] + with pytest.raises(DataCloudAccessLayerException) as exc_info: + client.write_to_dmo("test_dmo", mock_df, WriteMode.APPEND) + + assert "test_dlo" in str(exc_info.value) + + def test_mixed_dmo_dlo_raises_exception(self, reset_client, mock_spark): + """Test that mixing DMOs and DLOs raises an exception (converse case).""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = Client(reader=reader, writer=writer) + client._record_dmo_access("test_dmo") + + with pytest.raises(DataCloudAccessLayerException) as exc_info: + client.write_to_dlo("test_dlo", mock_df, WriteMode.APPEND) + + assert "test_dmo" in str(exc_info.value) + + def test_read_pattern_flow(self, reset_client, mock_spark): + """Test a complete flow of reading and writing within the same object type.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + reader.read_dlo.return_value = mock_df + + client = Client(reader=reader, writer=writer) + + df = client.read_dlo("source_dlo") + client.write_to_dlo("target_dlo", df, WriteMode.APPEND) + + reader.read_dlo.assert_called_once_with("source_dlo") + writer.write_to_dlo.assert_called_once_with( + "target_dlo", mock_df, WriteMode.APPEND ) - def test_read_dlo_deltas_records_runtime_source_name( + assert "source_dlo" in client._data_layer_history[DataCloudObjectType.DLO] + + # Reset for DMO test + Client._instance = None + client = Client(reader=reader, writer=writer) + reader.read_dmo.return_value = mock_df + + df = client.read_dmo("source_dmo") + client.write_to_dmo("target_dmo", df, WriteMode.MERGE) + + reader.read_dmo.assert_called_once_with("source_dmo") + writer.write_to_dmo.assert_called_once_with( + "target_dmo", mock_df, WriteMode.MERGE + ) + + assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + + +class TestStreamingClient: + + def test_singleton_pattern(self, reset_client, mock_spark): + """StreamingClient is a singleton, independent of Client.""" + reader = MockDataCloudReader(mock_spark) + writer = MockDataCloudWriter(mock_spark) + + client1 = StreamingClient(reader=reader, writer=writer) + client2 = StreamingClient() + + assert client1 is client2 + + with pytest.raises(ValueError): + StreamingClient(reader=MagicMock(spec=BaseDataCloudReader)) + + def test_streaming_client_is_distinct_from_batch_client( self, reset_client, mock_spark ): - """The runtime source env var populates the access-history entry.""" + """The two clients keep separate singleton instances and histories.""" + reader = MockDataCloudReader(mock_spark) + writer = MockDataCloudWriter(mock_spark) + + batch = Client(reader=reader, writer=writer) + streaming = StreamingClient(reader=reader, writer=writer) + + assert batch is not streaming + assert batch._data_layer_history is not streaming._data_layer_history + + def test_read_dlo_deltas(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) - reader.read_dlo_deltas.return_value = MagicMock(spec=DataFrame) + mock_df = MagicMock(spec=DataFrame) + reader.read_dlo_deltas.return_value = mock_df - client = Client(reader=reader, writer=writer) + client = StreamingClient(reader=reader, writer=writer) + # The streaming source is resolved by the runtime and recorded from the + # env var it sets; the caller passes no name. with patch.dict( "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_std__dll"} ): - client.read_dlo_deltas() + result = client.read_dlo_deltas() + reader.read_dlo_deltas.assert_called_once_with() + assert result is mock_df assert "Account_std__dll" in client._data_layer_history[DataCloudObjectType.DLO] + def test_read_dlo_deltas_without_source_env_raises(self, reset_client, mock_spark): + """Delta reads require the runtime source env var; absence fails fast.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + + client = StreamingClient(reader=reader, writer=writer) + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("BYOC_STREAMING_SOURCE_NAME", None) + with pytest.raises(RuntimeError) as exc_info: + client.read_dlo_deltas() + + assert "BYOC_STREAMING_SOURCE_NAME" in str(exc_info.value) + reader.read_dlo_deltas.assert_not_called() + def test_read_dmo_deltas(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) mock_df = MagicMock(spec=DataFrame) reader.read_dmo_deltas.return_value = mock_df - client = Client(reader=reader, writer=writer) + client = StreamingClient(reader=reader, writer=writer) with patch.dict( "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_model__dlm"} ): @@ -253,7 +348,7 @@ def test_write_dlo_deltas(self, reset_client, mock_spark): mock_query = MagicMock() writer.write_dlo_deltas.return_value = mock_query - client = Client(reader=reader, writer=writer) + client = StreamingClient(reader=reader, writer=writer) client._record_dlo_access("some_dlo") result = client.write_dlo_deltas("test_dlo", mock_df, extra_param=True) @@ -271,7 +366,7 @@ def test_write_dlo_deltas_after_dmo_read_raises_exception( writer = MagicMock(spec=BaseDataCloudWriter) mock_df = MagicMock(spec=DataFrame) - client = Client(reader=reader, writer=writer) + client = StreamingClient(reader=reader, writer=writer) client._record_dmo_access("test_dmo") with pytest.raises(DataCloudAccessLayerException) as exc_info: @@ -287,7 +382,7 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): stream_df = MagicMock(spec=DataFrame) reader.read_dlo_deltas.return_value = stream_df - client = Client(reader=reader, writer=writer) + client = StreamingClient(reader=reader, writer=writer) with patch.dict("os.environ", {"BYOC_STREAMING_SOURCE_NAME": "source_dll"}): df = client.read_dlo_deltas() @@ -297,70 +392,103 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): writer.write_dlo_deltas.assert_called_once_with("target_dll", stream_df) assert "source_dll" in client._data_layer_history[DataCloudObjectType.DLO] - def test_mixed_dlo_dmo_raises_exception(self, reset_client, mock_spark): - """Test that mixing DLOs and DMOs raises an exception.""" - reader = MagicMock(spec=BaseDataCloudReader) - writer = MagicMock(spec=BaseDataCloudWriter) - mock_df = MagicMock(spec=DataFrame) - client = Client(reader=reader, writer=writer) - client._record_dlo_access("test_dlo") +class TestSharedSparkSession: + """Both client types must share a single Spark session (one connection).""" - with pytest.raises(DataCloudAccessLayerException) as exc_info: - client.write_to_dmo("test_dmo", mock_df, WriteMode.APPEND) + def _make_config_client(self, client_cls, mock_spark): + """Build ``client_cls`` through the config path so it resolves a session + via the provider (rather than skipping it with an injected reader/writer). + Returns the provider's patched ``get_session`` mock.""" + from datacustomcode.spark.default import DefaultSparkSessionProvider - assert "test_dlo" in str(exc_info.value) + with ( + patch("datacustomcode.client.config") as mock_config, + patch.object( + DefaultSparkSessionProvider, "get_session" + ) as mock_get_session, + ): + mock_get_session.return_value = mock_spark - def test_mixed_dmo_dlo_raises_exception(self, reset_client, mock_spark): - """Test that mixing DMOs and DLOs raises an exception (converse case).""" - reader = MagicMock(spec=BaseDataCloudReader) - writer = MagicMock(spec=BaseDataCloudWriter) - mock_df = MagicMock(spec=DataFrame) + mock_reader_config = MagicMock() + mock_reader_config.to_object.return_value = MagicMock( + spec=BaseDataCloudReader + ) + mock_reader_config.force = False - client = Client(reader=reader, writer=writer) - client._record_dmo_access("test_dmo") + mock_writer_config = MagicMock() + mock_writer_config.to_object.return_value = MagicMock( + spec=BaseDataCloudWriter + ) + mock_writer_config.force = False - with pytest.raises(DataCloudAccessLayerException) as exc_info: - client.write_to_dlo("test_dlo", mock_df, WriteMode.APPEND) + mock_config.spark_provider_config = None + mock_config.reader_config = mock_reader_config + mock_config.writer_config = mock_writer_config + mock_config.spark_config = MagicMock(spec=SparkConfig) - assert "test_dmo" in str(exc_info.value) + client_cls() + return mock_get_session - def test_read_pattern_flow(self, reset_client, mock_spark): - """Test a complete flow of reading and writing within the same object type.""" - reader = MagicMock(spec=BaseDataCloudReader) - writer = MagicMock(spec=BaseDataCloudWriter) - mock_df = MagicMock(spec=DataFrame) - reader.read_dlo.return_value = mock_df + def test_two_client_types_reuse_one_session(self, reset_client, mock_spark): + """A StreamingClient created after a Client reuses the same session and + does not open a second connection.""" + batch_get_session = self._make_config_client(Client, mock_spark) + streaming_get_session = self._make_config_client(StreamingClient, mock_spark) - client = Client(reader=reader, writer=writer) + # The first client builds the session; the second reuses the cached one + # instead of asking the provider for another. + batch_get_session.assert_called_once() + streaming_get_session.assert_not_called() - df = client.read_dlo("source_dlo") - client.write_to_dlo("target_dlo", df, WriteMode.APPEND) + assert _BaseClient._shared_spark is mock_spark + assert Client._instance is not StreamingClient._instance - reader.read_dlo.assert_called_once_with("source_dlo") - writer.write_to_dlo.assert_called_once_with( - "target_dlo", mock_df, WriteMode.APPEND - ) + def test_reader_and_writer_built_against_shared_session( + self, reset_client, mock_spark + ): + """The reused session is the one handed to the second client's + reader/writer factories.""" + self._make_config_client(Client, mock_spark) - assert "source_dlo" in client._data_layer_history[DataCloudObjectType.DLO] + from datacustomcode.spark.default import DefaultSparkSessionProvider - # Reset for DMO test - Client._instance = None - client = Client(reader=reader, writer=writer) - reader.read_dmo.return_value = mock_df + with ( + patch("datacustomcode.client.config") as mock_config, + patch.object( + DefaultSparkSessionProvider, "get_session" + ) as mock_get_session, + ): + # Give the second client a *different* session if it were to build one, + # so a stale/duplicate build would be detectable. + mock_get_session.return_value = MagicMock(spec=SparkSession) - df = client.read_dmo("source_dmo") - client.write_to_dmo("target_dmo", df, WriteMode.MERGE) + mock_reader_config = MagicMock() + mock_reader_config.force = False + mock_writer_config = MagicMock() + mock_writer_config.force = False + mock_config.spark_provider_config = None + mock_config.reader_config = mock_reader_config + mock_config.writer_config = mock_writer_config + mock_config.spark_config = MagicMock(spec=SparkConfig) - reader.read_dmo.assert_called_once_with("source_dmo") - writer.write_to_dmo.assert_called_once_with( - "target_dmo", mock_df, WriteMode.MERGE - ) + StreamingClient() - assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + mock_get_session.assert_not_called() + mock_reader_config.to_object.assert_called_once_with(mock_spark) + mock_writer_config.to_object.assert_called_once_with(mock_spark) + + def test_injected_reader_writer_does_not_build_session( + self, reset_client, mock_spark + ): + """Injecting reader+writer skips session creation entirely, leaving the + shared session untouched for a later config-based client to populate.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + Client(reader=reader, writer=writer) -class TestClientLlmGatewayGenerateText: + assert _BaseClient._shared_spark is None @patch("datacustomcode.client._build_spark_llm_gateway") def test_forwards_args_to_spark_llm_gateway(self, mock_build_gateway, reset_client): From 5baa65afd1ebd7c65bb72e14869b60403dc68532 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 8 Jul 2026 08:47:52 -0400 Subject: [PATCH 3/5] streaming source from config --- src/datacustomcode/client.py | 21 ++++---- src/datacustomcode/config.py | 5 ++ src/datacustomcode/run.py | 18 +++++++ tests/test_client.py | 33 ++++++------ tests/test_run.py | 99 ++++++++++++++++++++++++++++++++++++ 5 files changed, 152 insertions(+), 24 deletions(-) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 8392ef9..c1cbff6 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -15,7 +15,6 @@ from __future__ import annotations from enum import Enum -import os from typing import ( TYPE_CHECKING, Any, @@ -52,18 +51,22 @@ from datacustomcode.spark.base import BaseSparkSessionProvider -_STREAMING_SOURCE_ENV = "BYOC_STREAMING_SOURCE_NAME" - - def _streaming_source_name() -> str: - """Return the runtime streaming source name. + """Return the streaming transform's read-source name. + + Resolved from ``config.streaming_source``, which ``run_entrypoint`` + populates from config.json's ``permissions.read`` entry. Raises: - RuntimeError: If ``BYOC_STREAMING_SOURCE_NAME`` is not set + RuntimeError: If no ``streaming_source`` has been configured (e.g. the + transform's config.json has no ``permissions.read`` entry). """ - source = os.environ.get(_STREAMING_SOURCE_ENV) + source = config.streaming_source if not source: - raise RuntimeError(f"{_STREAMING_SOURCE_ENV} is not set.") + raise RuntimeError( + "No streaming source configured. A streaming transform must declare " + "its read source in config.json under 'permissions.read'." + ) return source @@ -583,7 +586,7 @@ def write_dlo_deltas( """Write a streaming DataFrame of deltas to a DLO in Data Cloud. Starts a streaming query that writes each micro-batch to the - target DLO and returns the ``StreamingQuery`` handle; the caller + target DLO and returns the ``StreamingQuery`` handle; the caller typically calls ``query.awaitTermination()``. Args: diff --git a/src/datacustomcode/config.py b/src/datacustomcode/config.py index 901b295..779b5c9 100644 --- a/src/datacustomcode/config.py +++ b/src/datacustomcode/config.py @@ -89,6 +89,9 @@ class ClientConfig(BaseConfig): spark_provider_config: Union[ SparkProviderConfig[BaseSparkSessionProvider], None ] = None + # Source object name for a streaming (DELTA_SYNC) transform, populated by + # ``run_entrypoint`` from config.json's ``permissions.read`` + streaming_source: Union[str, None] = None def update(self, other: ClientConfig) -> ClientConfig: """Merge this ClientConfig with another, respecting force flags. @@ -116,6 +119,8 @@ def merge( self.spark_provider_config = merge( self.spark_provider_config, other.spark_provider_config ) + if other.streaming_source is not None: + self.streaming_source = other.streaming_source return self diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index 006055c..6167b0b 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -42,6 +42,22 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None: config_obj.options[key] = value +def _read_source_from_permissions(config_json: dict) -> Optional[str]: + """Return the read-source name from config.json ``permissions.read``. + """ + permissions = config_json.get("permissions") + if not isinstance(permissions, dict): + return None + read = permissions.get("read") + if not isinstance(read, dict): + return None + for layer in ("dlo", "dmo"): + names = read.get(layer) + if names: + return names[0] + return None + + def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]): if sf_cli_org: config_key = "sf_cli_org" @@ -125,6 +141,8 @@ def run_entrypoint( _set_config_option(config.reader_config, "dataspace", dataspace) _set_config_option(config.writer_config, "dataspace", dataspace) + config.streaming_source = _read_source_from_permissions(config_json) + _update_config_options(profile, sf_cli_org) for dependency in dependencies: diff --git a/tests/test_client.py b/tests/test_client.py index d020ebe..bc6f571 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from unittest.mock import MagicMock, patch from pyspark.sql import DataFrame, SparkSession @@ -84,13 +83,17 @@ def mock_config(mock_spark): @pytest.fixture def reset_client(): """Reset the client singletons (and the shared Spark session) between tests.""" + from datacustomcode.client import config as client_config + Client._instance = None StreamingClient._instance = None _BaseClient._shared_spark = None + client_config.streaming_source = None yield Client._instance = None StreamingClient._instance = None _BaseClient._shared_spark = None + client_config.streaming_source = None class TestClient: @@ -298,29 +301,29 @@ def test_read_dlo_deltas(self, reset_client, mock_spark): reader.read_dlo_deltas.return_value = mock_df client = StreamingClient(reader=reader, writer=writer) - # The streaming source is resolved by the runtime and recorded from the - # env var it sets; the caller passes no name. - with patch.dict( - "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_std__dll"} - ): + + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = "Account_std__dll" result = client.read_dlo_deltas() reader.read_dlo_deltas.assert_called_once_with() assert result is mock_df assert "Account_std__dll" in client._data_layer_history[DataCloudObjectType.DLO] - def test_read_dlo_deltas_without_source_env_raises(self, reset_client, mock_spark): - """Delta reads require the runtime source env var; absence fails fast.""" + def test_read_dlo_deltas_without_configured_source_raises( + self, reset_client, mock_spark + ): + """Delta reads require a configured streaming source; absence fails fast.""" reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) client = StreamingClient(reader=reader, writer=writer) - with patch.dict("os.environ", {}, clear=False): - os.environ.pop("BYOC_STREAMING_SOURCE_NAME", None) + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = None with pytest.raises(RuntimeError) as exc_info: client.read_dlo_deltas() - assert "BYOC_STREAMING_SOURCE_NAME" in str(exc_info.value) + assert "permissions.read" in str(exc_info.value) reader.read_dlo_deltas.assert_not_called() def test_read_dmo_deltas(self, reset_client, mock_spark): @@ -330,9 +333,8 @@ def test_read_dmo_deltas(self, reset_client, mock_spark): reader.read_dmo_deltas.return_value = mock_df client = StreamingClient(reader=reader, writer=writer) - with patch.dict( - "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_model__dlm"} - ): + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = "Account_model__dlm" result = client.read_dmo_deltas() reader.read_dmo_deltas.assert_called_once_with() @@ -384,7 +386,8 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): client = StreamingClient(reader=reader, writer=writer) - with patch.dict("os.environ", {"BYOC_STREAMING_SOURCE_NAME": "source_dll"}): + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = "source_dll" df = client.read_dlo_deltas() client.write_dlo_deltas("target_dll", df) diff --git a/tests/test_run.py b/tests/test_run.py index 1eace88..86c0d8f 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -488,3 +488,102 @@ def test_run_entrypoint_empty_dataspace_value(self): os.unlink(entrypoint_file) if os.path.exists(config_json_path): os.unlink(config_json_path) + + +class TestReadSourceFromPermissions: + """`_read_source_from_permissions` extracts the streaming read source from + config.json's `permissions.read`.""" + + def test_returns_single_dlo(self): + from datacustomcode.run import _read_source_from_permissions + + config_json = {"permissions": {"read": {"dlo": ["Account_std__dll"]}}} + assert _read_source_from_permissions(config_json) == "Account_std__dll" + + def test_returns_single_dmo(self): + from datacustomcode.run import _read_source_from_permissions + + config_json = {"permissions": {"read": {"dmo": ["Account_model__dlm"]}}} + assert _read_source_from_permissions(config_json) == "Account_model__dlm" + + def test_dlo_preferred_when_both_present(self): + from datacustomcode.run import _read_source_from_permissions + + config_json = { + "permissions": {"read": {"dlo": ["the_dll"], "dmo": ["the_dlm"]}} + } + assert _read_source_from_permissions(config_json) == "the_dll" + + def test_returns_first_of_multiple(self): + from datacustomcode.run import _read_source_from_permissions + + config_json = {"permissions": {"read": {"dlo": ["first__dll", "second__dll"]}}} + assert _read_source_from_permissions(config_json) == "first__dll" + + @pytest.mark.parametrize( + "config_json", + [ + {}, + {"permissions": None}, + {"permissions": {}}, + {"permissions": {"read": None}}, + {"permissions": {"read": {}}}, + {"permissions": {"read": {"dlo": []}}}, + ], + ) + def test_returns_none_when_absent_or_empty(self, config_json): + from datacustomcode.run import _read_source_from_permissions + + assert _read_source_from_permissions(config_json) is None + + +class TestStreamingSourceScenarios: + """`run_entrypoint` populates `config.streaming_source` from config.json.""" + + def _run_capturing_streaming_source(self, config_json_body): + """Run an entrypoint that records config.streaming_source and return it.""" + with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp: + entrypoint_content = textwrap.dedent( + """ + from datacustomcode.config import config + with open("streaming_source_output.txt", "w") as f: + f.write(f"streaming_source: {config.streaming_source}") + """ + ) + temp.write(entrypoint_content.encode("utf-8")) + entrypoint_file = temp.name + + entrypoint_dir = os.path.dirname(entrypoint_file) + config_json_path = os.path.join(entrypoint_dir, "config.json") + with open(config_json_path, "w") as f: + json.dump(config_json_body, f) + + try: + run_entrypoint( + entrypoint=entrypoint_file, + config_file=None, + dependencies=[], + profile="default", + ) + with open("streaming_source_output.txt", "r") as f: + return f.read() + finally: + if os.path.exists(entrypoint_file): + os.unlink(entrypoint_file) + if os.path.exists(config_json_path): + os.unlink(config_json_path) + if os.path.exists("streaming_source_output.txt"): + os.unlink("streaming_source_output.txt") + + def test_streaming_source_set_from_permissions_read(self): + content = self._run_capturing_streaming_source( + { + "dataspace": "default", + "permissions": {"read": {"dlo": ["Account_std__dll"]}}, + } + ) + assert "streaming_source: Account_std__dll" in content + + def test_streaming_source_none_for_batch_without_read(self): + content = self._run_capturing_streaming_source({"dataspace": "default"}) + assert "streaming_source: None" in content From f6deec03b573debd5dac9f545962e9f0a2e86300 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 8 Jul 2026 11:40:07 -0400 Subject: [PATCH 4/5] streaming source fix --- src/datacustomcode/client.py | 6 ++-- src/datacustomcode/config.py | 2 +- src/datacustomcode/run.py | 20 +++++-------- tests/test_client.py | 4 +-- tests/test_run.py | 56 ++++++++++++++---------------------- 5 files changed, 35 insertions(+), 53 deletions(-) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index c1cbff6..56c0588 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -55,17 +55,17 @@ def _streaming_source_name() -> str: """Return the streaming transform's read-source name. Resolved from ``config.streaming_source``, which ``run_entrypoint`` - populates from config.json's ``permissions.read`` entry. + populates from config.json's ``streamingSource`` field. Raises: RuntimeError: If no ``streaming_source`` has been configured (e.g. the - transform's config.json has no ``permissions.read`` entry). + transform's config.json has no ``streamingSource`` field). """ source = config.streaming_source if not source: raise RuntimeError( "No streaming source configured. A streaming transform must declare " - "its read source in config.json under 'permissions.read'." + "its read source in config.json under 'streamingSource'." ) return source diff --git a/src/datacustomcode/config.py b/src/datacustomcode/config.py index 779b5c9..1e2bead 100644 --- a/src/datacustomcode/config.py +++ b/src/datacustomcode/config.py @@ -90,7 +90,7 @@ class ClientConfig(BaseConfig): SparkProviderConfig[BaseSparkSessionProvider], None ] = None # Source object name for a streaming (DELTA_SYNC) transform, populated by - # ``run_entrypoint`` from config.json's ``permissions.read`` + # ``run_entrypoint`` from config.json's ``streamingSource`` field streaming_source: Union[str, None] = None def update(self, other: ClientConfig) -> ClientConfig: diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index 6167b0b..45e7749 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -42,20 +42,14 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None: config_obj.options[key] = value -def _read_source_from_permissions(config_json: dict) -> Optional[str]: - """Return the read-source name from config.json ``permissions.read``. +def _read_streaming_source(config_json: dict) -> Optional[str]: + """Return the streaming source name from config.json's ``streamingSource``. """ - permissions = config_json.get("permissions") - if not isinstance(permissions, dict): + source = config_json.get("streamingSource") + if not isinstance(source, dict): return None - read = permissions.get("read") - if not isinstance(read, dict): - return None - for layer in ("dlo", "dmo"): - names = read.get(layer) - if names: - return names[0] - return None + name = source.get("name") + return str(name) if name else None def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]): @@ -141,7 +135,7 @@ def run_entrypoint( _set_config_option(config.reader_config, "dataspace", dataspace) _set_config_option(config.writer_config, "dataspace", dataspace) - config.streaming_source = _read_source_from_permissions(config_json) + config.streaming_source = _read_streaming_source(config_json) _update_config_options(profile, sf_cli_org) diff --git a/tests/test_client.py b/tests/test_client.py index bc6f571..c40a995 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -301,7 +301,7 @@ def test_read_dlo_deltas(self, reset_client, mock_spark): reader.read_dlo_deltas.return_value = mock_df client = StreamingClient(reader=reader, writer=writer) - + with patch("datacustomcode.client.config") as mock_config: mock_config.streaming_source = "Account_std__dll" result = client.read_dlo_deltas() @@ -323,7 +323,7 @@ def test_read_dlo_deltas_without_configured_source_raises( with pytest.raises(RuntimeError) as exc_info: client.read_dlo_deltas() - assert "permissions.read" in str(exc_info.value) + assert "streamingSource" in str(exc_info.value) reader.read_dlo_deltas.assert_not_called() def test_read_dmo_deltas(self, reset_client, mock_spark): diff --git a/tests/test_run.py b/tests/test_run.py index 86c0d8f..154b0bf 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -490,51 +490,39 @@ def test_run_entrypoint_empty_dataspace_value(self): os.unlink(config_json_path) -class TestReadSourceFromPermissions: - """`_read_source_from_permissions` extracts the streaming read source from - config.json's `permissions.read`.""" +class TestReadStreamingSource: + """`_read_streaming_source` extracts the source name from config.json's + `streamingSource` object.""" - def test_returns_single_dlo(self): - from datacustomcode.run import _read_source_from_permissions + def test_returns_dlo_name(self): + from datacustomcode.run import _read_streaming_source - config_json = {"permissions": {"read": {"dlo": ["Account_std__dll"]}}} - assert _read_source_from_permissions(config_json) == "Account_std__dll" + config_json = {"streamingSource": {"type": "dlo", "name": "Account_Home__dll"}} + assert _read_streaming_source(config_json) == "Account_Home__dll" - def test_returns_single_dmo(self): - from datacustomcode.run import _read_source_from_permissions - - config_json = {"permissions": {"read": {"dmo": ["Account_model__dlm"]}}} - assert _read_source_from_permissions(config_json) == "Account_model__dlm" - - def test_dlo_preferred_when_both_present(self): - from datacustomcode.run import _read_source_from_permissions + def test_returns_dmo_name(self): + from datacustomcode.run import _read_streaming_source config_json = { - "permissions": {"read": {"dlo": ["the_dll"], "dmo": ["the_dlm"]}} + "streamingSource": {"type": "dmo", "name": "AccountTransformed__dlm"} } - assert _read_source_from_permissions(config_json) == "the_dll" - - def test_returns_first_of_multiple(self): - from datacustomcode.run import _read_source_from_permissions - - config_json = {"permissions": {"read": {"dlo": ["first__dll", "second__dll"]}}} - assert _read_source_from_permissions(config_json) == "first__dll" + assert _read_streaming_source(config_json) == "AccountTransformed__dlm" @pytest.mark.parametrize( "config_json", [ {}, - {"permissions": None}, - {"permissions": {}}, - {"permissions": {"read": None}}, - {"permissions": {"read": {}}}, - {"permissions": {"read": {"dlo": []}}}, + {"streamingSource": None}, + {"streamingSource": {}}, + {"streamingSource": {"type": "dlo"}}, + {"streamingSource": {"type": "dlo", "name": ""}}, + {"streamingSource": {"type": "dlo", "name": None}}, ], ) def test_returns_none_when_absent_or_empty(self, config_json): - from datacustomcode.run import _read_source_from_permissions + from datacustomcode.run import _read_streaming_source - assert _read_source_from_permissions(config_json) is None + assert _read_streaming_source(config_json) is None class TestStreamingSourceScenarios: @@ -575,15 +563,15 @@ def _run_capturing_streaming_source(self, config_json_body): if os.path.exists("streaming_source_output.txt"): os.unlink("streaming_source_output.txt") - def test_streaming_source_set_from_permissions_read(self): + def test_streaming_source_set_from_streaming_source_field(self): content = self._run_capturing_streaming_source( { "dataspace": "default", - "permissions": {"read": {"dlo": ["Account_std__dll"]}}, + "streamingSource": {"type": "dlo", "name": "Account_Home__dll"}, } ) - assert "streaming_source: Account_std__dll" in content + assert "streaming_source: Account_Home__dll" in content - def test_streaming_source_none_for_batch_without_read(self): + def test_streaming_source_none_for_batch_without_field(self): content = self._run_capturing_streaming_source({"dataspace": "default"}) assert "streaming_source: None" in content From 2cb85d32fa0e994952dbea100065b652e6b0fdf5 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 8 Jul 2026 12:24:07 -0400 Subject: [PATCH 5/5] lint --- src/datacustomcode/run.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index 45e7749..2be3a66 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -43,8 +43,7 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None: def _read_streaming_source(config_json: dict) -> Optional[str]: - """Return the streaming source name from config.json's ``streamingSource``. - """ + """Return the streaming source name from config.json's ``streamingSource``.""" source = config_json.get("streamingSource") if not isinstance(source, dict): return None