From b470208196535d85f7aa2b57be708800bfdaacbe Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 1 Jul 2026 10:54:44 -0400 Subject: [PATCH 1/9] Add streaming delta read/write methods to the SDK --- CHANGELOG.md | 20 +++++ README.md | 38 +++++++++ src/datacustomcode/client.py | 58 ++++++++++++++ src/datacustomcode/io/reader/base.py | 60 ++++++++++++++ src/datacustomcode/io/writer/base.py | 36 +++++++++ .../examples/streaming_deltas/entrypoint.py | 53 ++++++++++++ tests/io/reader/test_query_api.py | 12 +++ tests/io/writer/test_print.py | 6 ++ tests/test_client.py | 80 +++++++++++++++++++ 9 files changed, 363 insertions(+) create mode 100644 src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 03b8ca4..525199f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # Changelog +## 6.1.0 + +### Added + +- **Streaming (delta) read/write methods on `Client` for BYOC streaming transforms.** + + 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. + - `write_dlo_deltas(name, dataframe, write_mode)` – 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") + transformed = deltas.withColumn("description__c", upper(col("description__c"))) + query = client.write_dlo_deltas("Output__dll", transformed, WriteMode.APPEND) + query.awaitTermination() + ``` + + Supported streaming write modes are `WriteMode.APPEND`, `WriteMode.OVERWRITE`, and `WriteMode.MERGE_UPSERT_DELETE`. These methods run only inside the Data Cloud streaming (`DELTA_SYNC`) runtime; locally they raise `NotImplementedError`. See the `examples/streaming_deltas/entrypoint.py` example and the "Streaming (delta) transforms" section of the README. + ## 6.0.0 ### Breaking Changes diff --git a/README.md b/README.md index 3cf73b9..64b3d27 100644 --- a/README.md +++ b/README.md @@ -154,6 +154,11 @@ You should only need the following methods: * `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: +* `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 +* `write_dlo_deltas(name, spark_dataframe, write_mode)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` + For example: ```python from datacustomcode import Client @@ -169,6 +174,39 @@ client.write_to_dlo('output_DLO') > [!WARNING] > Currently we only support reading from DMOs and writing to DMOs or reading from DLOs and writing to DLOs, but they cannot mix. +### 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: + +```python +from pyspark.sql.functions import col, upper + +from datacustomcode import Client +from datacustomcode.io.writer.base import WriteMode + +client = Client() + +# read_dlo_deltas returns a *streaming* DataFrame over the change feed. +deltas = client.read_dlo_deltas("Input__dll") + +# Ordinary PySpark transform. Keep the change-feed metadata columns +# (those starting with "_") — the streaming sink needs them to apply +# inserts, updates, and deletes to the target DLO. +transformed = deltas.withColumn("description__c", upper(col("description__c"))) + +# write_dlo_deltas starts a streaming query and returns the StreamingQuery. +# The runtime owns the trigger and checkpoint location; you choose only the +# target table and write mode. +query = client.write_dlo_deltas("Output__dll", transformed, WriteMode.APPEND) +query.awaitTermination() +``` + +Notes: + +- Supported streaming write modes are `WriteMode.APPEND`, `WriteMode.OVERWRITE`, and `WriteMode.MERGE_UPSERT_DELETE`. +- These methods only run inside the Data Cloud streaming (`DELTA_SYNC`) runtime. Locally (`datacustomcode run`) they raise `NotImplementedError`, since there is no change feed to stream. +- A complete runnable entry point is provided in [`examples/streaming_deltas/entrypoint.py`](src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py). + ### Bundled file resolution Place bundled files (CSVs, prompt files, etc.) under `payload/files/`. The same `client.find_file_path("data.csv")` call resolves consistently across all three runtimes: diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 0d4df78..19b1209 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -35,6 +35,7 @@ from pathlib import Path from pyspark.sql import Column, DataFrame as PySparkDataFrame + from pyspark.sql.streaming import StreamingQuery from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions from datacustomcode.einstein_predictions.types import PredictionType @@ -335,6 +336,39 @@ 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: + """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. + + 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] + + def read_dmo_deltas(self, name: str) -> 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. + + 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] + def write_to_dlo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs ) -> None: @@ -361,6 +395,30 @@ def write_to_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, write_mode: WriteMode, **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. + write_mode: The write mode to use. Supported streaming modes are + ``WriteMode.APPEND``, ``WriteMode.OVERWRITE``, and + ``WriteMode.MERGE_UPSERT_DELETE``. + + Returns: + The started ``StreamingQuery``. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) + return self._writer.write_dlo_deltas(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] + def find_file_path(self, file_name: str) -> Path: """Resolve a bundled file shipped in the package to an absolute path. diff --git a/src/datacustomcode/io/reader/base.py b/src/datacustomcode/io/reader/base.py index ddea31b..2b69ed4 100644 --- a/src/datacustomcode/io/reader/base.py +++ b/src/datacustomcode/io/reader/base.py @@ -41,3 +41,63 @@ def read_dmo( name: str, schema: Union[AtomicType, StructType, str, None] = None, ) -> PySparkDataFrame: ... + + def read_dlo_deltas( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> PySparkDataFrame: + """Read the streaming change feed (deltas) for a Data Lake Object. + + This is the streaming counterpart to :meth:`read_dlo`. It returns a + streaming DataFrame over the change feed the Data Cloud runtime + publishes for a streaming (``DELTA_SYNC``) transform. Concrete + streaming behavior is provided by the deployed Data Cloud runtime; the + 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. + + Raises: + NotImplementedError: If the active reader does not support streaming + deltas (e.g. the local development readers). + """ + raise NotImplementedError( + "read_dlo_deltas is only supported when running in the Data Cloud " + "streaming runtime; the local reader does not support streaming " + "deltas." + ) + + def read_dmo_deltas( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> 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. + + Raises: + NotImplementedError: If the active reader does not support streaming + deltas (e.g. the local development readers). + """ + raise NotImplementedError( + "read_dmo_deltas is only supported when running in the Data Cloud " + "streaming runtime; the local reader does not support streaming " + "deltas." + ) diff --git a/src/datacustomcode/io/writer/base.py b/src/datacustomcode/io/writer/base.py index cb01f76..6455ddb 100644 --- a/src/datacustomcode/io/writer/base.py +++ b/src/datacustomcode/io/writer/base.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession + from pyspark.sql.streaming import StreamingQuery class WriteMode(str, Enum): @@ -57,3 +58,38 @@ def write_to_dlo( def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: ... + + def write_dlo_deltas( + self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a Data Lake Object. + + Streaming counterpart to :meth:`write_to_dlo`. Starts a streaming query + that writes each micro-batch to the target DLO via the Data Cloud + streaming sink and returns the resulting ``StreamingQuery`` handle. The + runtime owns the trigger and checkpoint location; callers pass only the + table name and write mode. Concrete streaming behavior is provided by + the deployed Data Cloud runtime; the base implementation raises + :class:`NotImplementedError`. + + Args: + name: Target Data Lake Object name. + dataframe: Streaming PySpark DataFrame produced from a + ``read_dlo_deltas`` / ``read_dmo_deltas`` source. + write_mode: Write mode for the streaming sink. Supported modes are + ``WriteMode.APPEND``, ``WriteMode.OVERWRITE``, and + ``WriteMode.MERGE_UPSERT_DELETE``. + + Returns: + The started ``StreamingQuery``; the caller drives its lifecycle + (typically ``query.awaitTermination()``). + + Raises: + NotImplementedError: If the active writer does not support streaming + deltas (e.g. the local development writers). + """ + raise NotImplementedError( + "write_dlo_deltas is only supported when running in the Data Cloud " + "streaming runtime; the local writer does not support streaming " + "deltas." + ) diff --git a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py new file mode 100644 index 0000000..01c19f3 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -0,0 +1,53 @@ +"""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: + +* ``client.read_dlo_deltas(name)`` 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, write_mode)`` starts a streaming query + that writes each micro-batch to the target DLO and returns the + ``StreamingQuery`` handle. The runtime owns the trigger and checkpoint + location — the caller only chooses the table and write mode. + +The transform in between is ordinary PySpark. Because the source is a change +feed, keep the metadata columns on the DataFrame you hand to +``write_dlo_deltas`` — the sink relies on them to merge changes correctly. + +This entrypoint only runs inside the Data Cloud streaming (``DELTA_SYNC``) +runtime; the local ``datacustomcode run`` readers/writers raise +``NotImplementedError`` for the delta methods. +""" + +from pyspark.sql.functions import col, upper + +from datacustomcode.client import Client +from datacustomcode.io.writer.base import WriteMode + + +def main(): + client = Client() + + # Streaming DataFrame over the source DLO's change feed. + deltas = client.read_dlo_deltas("Account_std__dll") + + # Ordinary PySpark transform. Note we do NOT drop the change-feed metadata + # columns (those starting with "_") — the streaming sink needs them to apply + # inserts, updates, and deletes to the target DLO. + transformed = deltas.withColumn("description__c", upper(col("description__c"))) + + # Start the streaming write. write_dlo_deltas returns the StreamingQuery; + # the trigger and checkpoint location are provided by the runtime. + query = client.write_dlo_deltas( + "Account_std_copy__dll", transformed, WriteMode.APPEND + ) + + # Drive the query's lifecycle. In the streaming runtime this blocks until + # the job is stopped by the platform. + query.awaitTermination() + + +if __name__ == "__main__": + main() diff --git a/tests/io/reader/test_query_api.py b/tests/io/reader/test_query_api.py index f9baa81..eaa8848 100644 --- a/tests/io/reader/test_query_api.py +++ b/tests/io/reader/test_query_api.py @@ -224,6 +224,18 @@ def test_read_dlo( assert args[0] is mock_pandas_dataframe # First arg is the pandas DataFrame assert isinstance(args[1], StructType) # Second arg is the schema + 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") + 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") + assert "read_dmo_deltas" in str(exc_info.value) + def test_read_dlo_with_schema( self, reader_without_init, mock_connection, mock_pandas_dataframe ): diff --git a/tests/io/writer/test_print.py b/tests/io/writer/test_print.py index bd5b1f0..cb7e52b 100644 --- a/tests/io/writer/test_print.py +++ b/tests/io/writer/test_print.py @@ -59,6 +59,12 @@ def test_write_to_dmo(self, print_writer, mock_dataframe): # Verify show() was called mock_dataframe.show.assert_called_once() + def test_write_dlo_deltas_not_supported_locally(self, print_writer, mock_dataframe): + """Streaming delta writes are not supported by the local writer.""" + with pytest.raises(NotImplementedError) as exc_info: + print_writer.write_dlo_deltas("test_dll", mock_dataframe, WriteMode.APPEND) + assert "write_dlo_deltas" in str(exc_info.value) + def test_config_name(self): """Test that the CONFIG_NAME class variable is set correctly.""" assert PrintDataCloudWriter.CONFIG_NAME == "PrintDataCloudWriter" diff --git a/tests/test_client.py b/tests/test_client.py index 86b72aa..1112b60 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -193,6 +193,86 @@ 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): + 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) + result = client.read_dlo_deltas("test_dlo") + + reader.read_dlo_deltas.assert_called_once_with("test_dlo") + assert result is mock_df + assert "test_dlo" in client._data_layer_history[DataCloudObjectType.DLO] + + 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) + result = client.read_dmo_deltas("test_dmo") + + reader.read_dmo_deltas.assert_called_once_with("test_dmo") + assert result is mock_df + assert "test_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + + def test_write_dlo_deltas(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + mock_query = MagicMock() + writer.write_dlo_deltas.return_value = mock_query + + client = Client(reader=reader, writer=writer) + client._record_dlo_access("some_dlo") + + result = client.write_dlo_deltas( + "test_dlo", mock_df, WriteMode.APPEND, extra_param=True + ) + + writer.write_dlo_deltas.assert_called_once_with( + "test_dlo", mock_df, WriteMode.APPEND, extra_param=True + ) + assert result is mock_query + + def test_write_dlo_deltas_after_dmo_read_raises_exception( + self, reset_client, mock_spark + ): + """Streaming DLO write is subject to the same DLO/DMO mixing guard.""" + 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_dlo_deltas("test_dlo", mock_df, WriteMode.APPEND) + + assert "test_dmo" in str(exc_info.value) + writer.write_dlo_deltas.assert_not_called() + + def test_streaming_read_write_flow(self, reset_client, mock_spark): + """A read_dlo_deltas → write_dlo_deltas flow stays within the DLO layer.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + stream_df = MagicMock(spec=DataFrame) + reader.read_dlo_deltas.return_value = stream_df + + client = Client(reader=reader, writer=writer) + + df = client.read_dlo_deltas("source_dll") + client.write_dlo_deltas("target_dll", df, WriteMode.MERGE_UPSERT_DELETE) + + reader.read_dlo_deltas.assert_called_once_with("source_dll") + writer.write_dlo_deltas.assert_called_once_with( + "target_dll", stream_df, WriteMode.MERGE_UPSERT_DELETE + ) + 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) From 5837e4c8c3d6700fbd6892f145224eff029ccd3a Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 1 Jul 2026 14:34:33 -0400 Subject: [PATCH 2/9] documentation fixes and write mode removal --- CHANGELOG.md | 6 +++--- README.md | 14 +++++--------- src/datacustomcode/client.py | 7 ++----- src/datacustomcode/io/writer/base.py | 9 +++------ .../examples/streaming_deltas/entrypoint.py | 17 ++++++----------- tests/io/writer/test_print.py | 2 +- tests/test_client.py | 14 +++++--------- 7 files changed, 25 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 525199f..93b82ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,16 +9,16 @@ 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. - - `write_dlo_deltas(name, dataframe, write_mode)` – start a streaming query that writes each micro-batch to the target DLO and return the `StreamingQuery` handle. + - `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") transformed = deltas.withColumn("description__c", upper(col("description__c"))) - query = client.write_dlo_deltas("Output__dll", transformed, WriteMode.APPEND) + query = client.write_dlo_deltas("Output__dll", transformed) query.awaitTermination() ``` - Supported streaming write modes are `WriteMode.APPEND`, `WriteMode.OVERWRITE`, and `WriteMode.MERGE_UPSERT_DELETE`. These methods run only inside the Data Cloud streaming (`DELTA_SYNC`) runtime; locally they raise `NotImplementedError`. See the `examples/streaming_deltas/entrypoint.py` example and the "Streaming (delta) transforms" section of the README. + These methods run only inside the Data Cloud streaming (`DELTA_SYNC`) runtime; locally they raise `NotImplementedError`. See the `examples/streaming_deltas/entrypoint.py` example and the "Streaming (delta) transforms" section of the README. ## 6.0.0 diff --git a/README.md b/README.md index 64b3d27..789eef7 100644 --- a/README.md +++ b/README.md @@ -157,7 +157,7 @@ You should only need the following methods: 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 -* `write_dlo_deltas(name, spark_dataframe, write_mode)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` +* `write_dlo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` For example: ```python @@ -182,28 +182,24 @@ Streaming BYOC transforms process a Data Lake Object's Change Data Feed continuo from pyspark.sql.functions import col, upper from datacustomcode import Client -from datacustomcode.io.writer.base import WriteMode client = Client() # read_dlo_deltas returns a *streaming* DataFrame over the change feed. deltas = client.read_dlo_deltas("Input__dll") -# Ordinary PySpark transform. Keep the change-feed metadata columns -# (those starting with "_") — the streaming sink needs them to apply -# inserts, updates, and deletes to the target DLO. +# Ordinary PySpark transform. transformed = deltas.withColumn("description__c", upper(col("description__c"))) # write_dlo_deltas starts a streaming query and returns the StreamingQuery. -# The runtime owns the trigger and checkpoint location; you choose only the -# target table and write mode. -query = client.write_dlo_deltas("Output__dll", transformed, WriteMode.APPEND) +# The runtime owns the trigger and checkpoint location; you +# choose only the target table. +query = client.write_dlo_deltas("Output__dll", transformed) query.awaitTermination() ``` Notes: -- Supported streaming write modes are `WriteMode.APPEND`, `WriteMode.OVERWRITE`, and `WriteMode.MERGE_UPSERT_DELETE`. - These methods only run inside the Data Cloud streaming (`DELTA_SYNC`) runtime. Locally (`datacustomcode run`) they raise `NotImplementedError`, since there is no change feed to stream. - A complete runnable entry point is provided in [`examples/streaming_deltas/entrypoint.py`](src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py). diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 19b1209..143796d 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -396,7 +396,7 @@ def write_to_dmo( 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, write_mode: WriteMode, **kwargs + self, name: str, dataframe: PySparkDataFrame, **kwargs ) -> StreamingQuery: """Write a streaming DataFrame of deltas to a DLO in Data Cloud. @@ -409,15 +409,12 @@ def write_dlo_deltas( Args: name: The name of the DLO to write to. dataframe: The streaming PySpark DataFrame to write. - write_mode: The write mode to use. Supported streaming modes are - ``WriteMode.APPEND``, ``WriteMode.OVERWRITE``, and - ``WriteMode.MERGE_UPSERT_DELETE``. Returns: The started ``StreamingQuery``. """ self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) - return self._writer.write_dlo_deltas(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return] + return self._writer.write_dlo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return] def find_file_path(self, file_name: str) -> Path: """Resolve a bundled file shipped in the package to an absolute path. diff --git a/src/datacustomcode/io/writer/base.py b/src/datacustomcode/io/writer/base.py index 6455ddb..47a7bd2 100644 --- a/src/datacustomcode/io/writer/base.py +++ b/src/datacustomcode/io/writer/base.py @@ -60,7 +60,7 @@ def write_to_dmo( ) -> None: ... def write_dlo_deltas( - self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode + self, name: str, dataframe: PySparkDataFrame ) -> StreamingQuery: """Write a streaming DataFrame of deltas to a Data Lake Object. @@ -68,17 +68,14 @@ def write_dlo_deltas( that writes each micro-batch to the target DLO via the Data Cloud streaming sink and returns the resulting ``StreamingQuery`` handle. The runtime owns the trigger and checkpoint location; callers pass only the - table name and write mode. Concrete streaming behavior is provided by - the deployed Data Cloud runtime; the base implementation raises + table name. Concrete streaming behavior is provided by the deployed + Data Cloud runtime; the base implementation raises :class:`NotImplementedError`. Args: name: Target Data Lake Object name. dataframe: Streaming PySpark DataFrame produced from a ``read_dlo_deltas`` / ``read_dmo_deltas`` source. - write_mode: Write mode for the streaming sink. Supported modes are - ``WriteMode.APPEND``, ``WriteMode.OVERWRITE``, and - ``WriteMode.MERGE_UPSERT_DELETE``. Returns: The started ``StreamingQuery``; the caller drives its lifecycle diff --git a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py index 01c19f3..8919cf7 100644 --- a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -7,10 +7,10 @@ * ``client.read_dlo_deltas(name)`` 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, write_mode)`` starts a streaming query - that writes each micro-batch to the target DLO and returns the - ``StreamingQuery`` handle. The runtime owns the trigger and checkpoint - location — the caller only chooses the table and write mode. +* ``client.write_dlo_deltas(name, df)`` starts a streaming query that writes + each micro-batch to the target DLO and returns the ``StreamingQuery`` handle. + The runtime owns the trigger, and checkpoint location — the caller only + chooses the table. The transform in between is ordinary PySpark. Because the source is a change feed, keep the metadata columns on the DataFrame you hand to @@ -24,7 +24,6 @@ from pyspark.sql.functions import col, upper from datacustomcode.client import Client -from datacustomcode.io.writer.base import WriteMode def main(): @@ -33,16 +32,12 @@ def main(): # Streaming DataFrame over the source DLO's change feed. deltas = client.read_dlo_deltas("Account_std__dll") - # Ordinary PySpark transform. Note we do NOT drop the change-feed metadata - # columns (those starting with "_") — the streaming sink needs them to apply - # inserts, updates, and deletes to the target DLO. + # Ordinary PySpark transform. transformed = deltas.withColumn("description__c", upper(col("description__c"))) # Start the streaming write. write_dlo_deltas returns the StreamingQuery; # the trigger and checkpoint location are provided by the runtime. - query = client.write_dlo_deltas( - "Account_std_copy__dll", transformed, WriteMode.APPEND - ) + query = client.write_dlo_deltas("Account_std_copy__dll", transformed) # Drive the query's lifecycle. In the streaming runtime this blocks until # the job is stopped by the platform. diff --git a/tests/io/writer/test_print.py b/tests/io/writer/test_print.py index cb7e52b..a10f2a3 100644 --- a/tests/io/writer/test_print.py +++ b/tests/io/writer/test_print.py @@ -62,7 +62,7 @@ def test_write_to_dmo(self, print_writer, mock_dataframe): def test_write_dlo_deltas_not_supported_locally(self, print_writer, mock_dataframe): """Streaming delta writes are not supported by the local writer.""" with pytest.raises(NotImplementedError) as exc_info: - print_writer.write_dlo_deltas("test_dll", mock_dataframe, WriteMode.APPEND) + print_writer.write_dlo_deltas("test_dll", mock_dataframe) assert "write_dlo_deltas" in str(exc_info.value) def test_config_name(self): diff --git a/tests/test_client.py b/tests/test_client.py index 1112b60..e661cbc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -229,12 +229,10 @@ def test_write_dlo_deltas(self, reset_client, mock_spark): client = Client(reader=reader, writer=writer) client._record_dlo_access("some_dlo") - result = client.write_dlo_deltas( - "test_dlo", mock_df, WriteMode.APPEND, extra_param=True - ) + result = client.write_dlo_deltas("test_dlo", mock_df, extra_param=True) writer.write_dlo_deltas.assert_called_once_with( - "test_dlo", mock_df, WriteMode.APPEND, extra_param=True + "test_dlo", mock_df, extra_param=True ) assert result is mock_query @@ -250,7 +248,7 @@ def test_write_dlo_deltas_after_dmo_read_raises_exception( client._record_dmo_access("test_dmo") with pytest.raises(DataCloudAccessLayerException) as exc_info: - client.write_dlo_deltas("test_dlo", mock_df, WriteMode.APPEND) + client.write_dlo_deltas("test_dlo", mock_df) assert "test_dmo" in str(exc_info.value) writer.write_dlo_deltas.assert_not_called() @@ -265,12 +263,10 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): client = Client(reader=reader, writer=writer) df = client.read_dlo_deltas("source_dll") - client.write_dlo_deltas("target_dll", df, WriteMode.MERGE_UPSERT_DELETE) + client.write_dlo_deltas("target_dll", df) reader.read_dlo_deltas.assert_called_once_with("source_dll") - writer.write_dlo_deltas.assert_called_once_with( - "target_dll", stream_df, WriteMode.MERGE_UPSERT_DELETE - ) + 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): From 3da41b60d2e64a24d724c8918bf52580ecbcb2c7 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Thu, 2 Jul 2026 15:11:25 -0400 Subject: [PATCH 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 From 4c00881c6e0b3d9b87d389988abf1508fc6cbe41 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 22 Jul 2026 15:02:18 -0400 Subject: [PATCH 8/9] Support creating and deploying streaming script code packages --- src/datacustomcode/cli.py | 33 +++- src/datacustomcode/constants.py | 8 + src/datacustomcode/deploy.py | 71 ++++++-- src/datacustomcode/scan.py | 196 ++++++++++++++++++---- src/datacustomcode/template.py | 14 +- tests/test_cli.py | 97 ++++++++++- tests/test_deploy.py | 286 ++++++++++++++++++++++++++++++++ tests/test_scan.py | 207 +++++++++++++++++++++++ 8 files changed, 860 insertions(+), 52 deletions(-) diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py index 7e2cd00..3a03e6d 100644 --- a/src/datacustomcode/cli.py +++ b/src/datacustomcode/cli.py @@ -283,10 +283,19 @@ def deploy( ) @click.option( "--use-in-feature", - default="SearchIndexChunking", - help="Feature where this function will be used (only applicable for function).", + "-u", + default=None, + help=( + "Invoke option for this package. For scripts: 'BatchTransform' " + "(default) or 'StreamingTransform'. For functions: 'SearchIndexChunking'." + ), ) def init(directory: str, code_type: str, use_in_feature: Optional[str]): + from datacustomcode.constants import ( + SCRIPT_USE_IN_FEATURE_BATCH, + SCRIPT_USE_IN_FEATURE_OPTIONS, + SCRIPT_USE_IN_FEATURE_STREAMING, + ) from datacustomcode.scan import ( dc_config_json_from_file, update_config, @@ -294,9 +303,23 @@ def init(directory: str, code_type: str, use_in_feature: Optional[str]): ) from datacustomcode.template import copy_function_template, copy_script_template + streaming = False + if code_type == "script": + use_in_feature = use_in_feature or SCRIPT_USE_IN_FEATURE_BATCH + if use_in_feature not in SCRIPT_USE_IN_FEATURE_OPTIONS: + click.secho( + f"Error: Invalid --use-in-feature '{use_in_feature}' for a " + f"script. Valid options: {', '.join(SCRIPT_USE_IN_FEATURE_OPTIONS)}.", + fg="red", + ) + raise click.Abort() + streaming = use_in_feature == SCRIPT_USE_IN_FEATURE_STREAMING + else: + use_in_feature = use_in_feature or "SearchIndexChunking" + click.echo("Copying template to " + click.style(directory, fg="blue", bold=True)) if code_type == "script": - copy_script_template(directory) + copy_script_template(directory, streaming=streaming) elif code_type == "function": copy_function_template(directory, use_in_feature) entrypoint_path = os.path.join(directory, PAYLOAD_DIR, ENTRYPOINT_FILE) @@ -306,7 +329,9 @@ def init(directory: str, code_type: str, use_in_feature: Optional[str]): sdk_config = {"type": code_type} write_sdk_config(directory, sdk_config) - config_json = dc_config_json_from_file(entrypoint_path, code_type) + config_json = dc_config_json_from_file( + entrypoint_path, code_type, streaming=streaming + ) with open(config_location, "w") as f: json.dump(config_json, f, indent=2) diff --git a/src/datacustomcode/constants.py b/src/datacustomcode/constants.py index 76b6a7c..fe29a5c 100644 --- a/src/datacustomcode/constants.py +++ b/src/datacustomcode/constants.py @@ -38,6 +38,14 @@ "SearchIndexChunking": "UnstructuredChunking", } +# Script (data transform) invoke options +SCRIPT_USE_IN_FEATURE_BATCH = "BatchTransform" +SCRIPT_USE_IN_FEATURE_STREAMING = "StreamingTransform" +SCRIPT_USE_IN_FEATURE_OPTIONS = [ + SCRIPT_USE_IN_FEATURE_BATCH, + SCRIPT_USE_IN_FEATURE_STREAMING, +] + # Pydantic request/response type names to feature names REQUEST_TYPE_TO_FEATURE = { "SearchIndexChunkingV1Request": "SearchIndexChunking", diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index e8c4ec4..e8555ba 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -14,6 +14,7 @@ # limitations under the License. from __future__ import annotations +import copy from html import unescape import json import os @@ -41,6 +42,7 @@ DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code" DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms" +DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH = "services/data/v67.0/ssot/data-custom-code" WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000 # Available compute types for Data Cloud deployments. @@ -108,6 +110,7 @@ class CodeExtensionMetadata(BaseModel): computeType: str codeType: str functionInvokeOptions: Union[list[str], None] = None + invokeOptions: Union[list[str], None] = None def __init__(self, **data): name = data.get("name", "") @@ -200,7 +203,14 @@ def create_deployment( access_token: AccessTokenResponse, metadata: CodeExtensionMetadata ) -> CreateDeploymentResponse: """Create a custom code deployment in the DataCloud.""" - url = _join_strip_url(access_token.instance_url, DATA_CUSTOM_CODE_PATH) + # invokeOptions only binds at v67.0; route there when it is set so the + # option isn't silently dropped. Everything else stays on v63.0. + code_custom_code_path = ( + DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH + if metadata.invokeOptions + else DATA_CUSTOM_CODE_PATH + ) + url = _join_strip_url(access_token.instance_url, code_custom_code_path) body = dict[str, Any]( { "label": metadata.name, @@ -213,6 +223,8 @@ def create_deployment( ) if metadata.functionInvokeOptions: body["functionInvokeOptions"] = metadata.functionInvokeOptions + if metadata.invokeOptions: + body["invokeOptions"] = metadata.invokeOptions logger.debug(f"Creating deployment {metadata.name}...") try: response = _make_api_call( @@ -388,6 +400,30 @@ class DataTransformConfig(BaseConfig): dataspace: str permissions: Permissions dataObjects: Optional[list[DataObject]] = None + streamingSource: Optional[StreamingSource] = None + + @property + def is_streaming(self) -> bool: + return self.streamingSource is not None + + @model_validator(mode="after") + def _validate_layers(self) -> "DataTransformConfig": + read_is_dlo = isinstance(self.permissions.read, DloPermission) + write_is_dlo = isinstance(self.permissions.write, DloPermission) + if self.is_streaming: + if not write_is_dlo: + raise ValueError( + "A streaming transform must write to a DLO " + "(permissions.write must be a 'dlo' entry)." + ) + elif read_is_dlo != write_is_dlo: + raise ValueError( + "permissions.read and permissions.write must both reference " + "DLOs or both reference DMOs (got " + f"read={type(self.permissions.read).__name__}, " + f"write={type(self.permissions.write).__name__})" + ) + return self class FunctionConfig(BaseConfig): @@ -402,23 +438,15 @@ class DmoPermission(BaseModel): dmo: list[str] +class StreamingSource(BaseModel): + type: str + name: str + + class Permissions(BaseModel): read: Union[DloPermission, DmoPermission] write: Union[DloPermission, DmoPermission] - @model_validator(mode="after") - def _no_mixed_layers(self) -> "Permissions": - read_is_dlo = isinstance(self.read, DloPermission) - write_is_dlo = isinstance(self.write, DloPermission) - if read_is_dlo != write_is_dlo: - raise ValueError( - "permissions.read and permissions.write must both reference " - "DLOs or both reference DMOs (got " - f"read={type(self.read).__name__}, " - f"write={type(self.write).__name__})" - ) - return self - def _permission_entries(perm: Union[DloPermission, DmoPermission]) -> list[str]: """Return the list of object names regardless of layer (DLO or DMO).""" @@ -490,7 +518,9 @@ def create_data_transform( ) -> dict: """Create a data transform in the DataCloud.""" script_name = metadata.name - request_hydrated = DATA_TRANSFORM_REQUEST_TEMPLATE.copy() + # Deep copy: the template's nested nodes/sources/macros dicts would + # otherwise be shared across calls and accumulate entries between deploys. + request_hydrated = copy.deepcopy(DATA_TRANSFORM_REQUEST_TEMPLATE) # Add nodes for each write entry (DLO or DMO) for i, name in enumerate( @@ -533,7 +563,7 @@ def create_data_transform( "definition": definition, "label": f"{metadata.name}", "name": f"{metadata.name}", - "type": "BATCH", + "type": "STREAMING" if data_transform_config.is_streaming else "BATCH", "dataSpaceName": data_transform_config.dataspace, } @@ -616,9 +646,18 @@ def deploy_full( callback=None, ) -> AccessTokenResponse: """Deploy a data transform in the DataCloud.""" + from datacustomcode.constants import SCRIPT_USE_IN_FEATURE_STREAMING + # prepare payload config = get_config(directory) + if ( + isinstance(config, DataTransformConfig) + and config.is_streaming + and not metadata.invokeOptions + ): + metadata.invokeOptions = [SCRIPT_USE_IN_FEATURE_STREAMING] + # create deployment and upload payload deployment = create_deployment(access_token, metadata) zip(directory, docker_network, metadata.codeType) diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index 5e50c5d..cdb6434 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -15,6 +15,7 @@ from __future__ import annotations import ast +import copy import json import os import sys @@ -22,6 +23,7 @@ Any, ClassVar, Dict, + Optional, Set, Union, ) @@ -32,6 +34,8 @@ from datacustomcode.version import get_version DATA_ACCESS_METHODS = ["read_dlo", "read_dmo", "write_to_dlo", "write_to_dmo"] +STREAMING_READ_METHODS = ["read_dlo_deltas", "read_dmo_deltas"] +STREAMING_WRITE_METHODS = ["write_dlo_deltas"] DATA_TRANSFORM_CONFIG_TEMPLATE = { "sdkVersion": get_version(), @@ -43,6 +47,20 @@ }, } +STREAMING_TRANSFORM_CONFIG_TEMPLATE = { + "sdkVersion": get_version(), + "entryPoint": "", + "dataspace": "default", + "streamingSource": { + "type": "dlo", + "name": "", + }, + "permissions": { + "read": {}, + "write": {}, + }, +} + FUNCTION_CONFIG_TEMPLATE = { "entryPoint": "", } @@ -160,6 +178,35 @@ def output_str(self) -> str: return next(iter(self.write_to_dmo)) +class StreamingDataAccessLayerCalls(pydantic.BaseModel): + read_dlo_deltas: bool + read_dmo_deltas: bool + write_dlo_deltas: frozenset[str] + + @pydantic.model_validator(mode="after") + def validate_access_layer(self) -> StreamingDataAccessLayerCalls: + if self.read_dlo_deltas and self.read_dmo_deltas: + raise ValueError( + "Cannot read DLO and DMO deltas in the same streaming transform." + ) + if not self.read_dlo_deltas and not self.read_dmo_deltas: + raise ValueError( + "A streaming transform must read from at least one DLO or DMO " + "delta stream (read_dlo_deltas / read_dmo_deltas)." + ) + if not self.write_dlo_deltas: + raise ValueError( + "A streaming transform must write to at least one DLO via " + "write_dlo_deltas." + ) + return self + + @property + def read_layer(self) -> str: + """Return the read source layer, ``"dlo"`` or ``"dmo"``.""" + return "dlo" if self.read_dlo_deltas else "dmo" + + class ClientMethodVisitor(ast.NodeVisitor): """AST Visitor that finds all instances of Client read/write method calls.""" @@ -168,6 +215,9 @@ def __init__(self) -> None: self._read_dmo_instances: set[str] = set() self._write_to_dlo_instances: set[str] = set() self._write_to_dmo_instances: set[str] = set() + self._read_dlo_deltas: bool = False + self._read_dmo_deltas: bool = False + self._write_dlo_deltas_instances: set[str] = set() self.variable_values: Dict[str, Union[str, None]] = {} def visit_Assign(self, node: ast.Assign) -> None: @@ -189,14 +239,15 @@ def visit_Call(self, node: ast.Call) -> None: node.func.value, ast.Name ): method_name = node.func.attr + + if method_name == "read_dlo_deltas": + self._read_dlo_deltas = True + elif method_name == "read_dmo_deltas": + self._read_dmo_deltas = True + if method_name in DATA_ACCESS_METHODS and node.args: arg = node.args[0] - name = None - - if isinstance(arg, ast.Constant) and isinstance(arg.value, str): - name = arg.value - elif isinstance(arg, ast.Name) and arg.id in self.variable_values: - name = self.variable_values[arg.id] + name = self._resolve_name_arg(arg) if name: if method_name == "read_dlo": @@ -207,8 +258,29 @@ def visit_Call(self, node: ast.Call) -> None: self._write_to_dlo_instances.add(name) elif method_name == "write_to_dmo": self._write_to_dmo_instances.add(name) + elif method_name in STREAMING_WRITE_METHODS and node.args: + name = self._resolve_name_arg(node.args[0]) + if name and method_name == "write_dlo_deltas": + self._write_dlo_deltas_instances.add(name) self.generic_visit(node) + def _resolve_name_arg(self, arg: ast.expr) -> Union[str, None]: + """Resolve a string-literal or tracked-variable first argument.""" + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + return arg.value + if isinstance(arg, ast.Name) and arg.id in self.variable_values: + return self.variable_values[arg.id] + return None + + @property + def is_streaming(self) -> bool: + """Whether any streaming (delta) access method was found.""" + return ( + self._read_dlo_deltas + or self._read_dmo_deltas + or bool(self._write_dlo_deltas_instances) + ) + def found(self) -> DataAccessLayerCalls: return DataAccessLayerCalls( read_dlo=frozenset(self._read_dlo_instances), @@ -217,6 +289,13 @@ def found(self) -> DataAccessLayerCalls: write_to_dmo=frozenset(self._write_to_dmo_instances), ) + def found_streaming(self) -> StreamingDataAccessLayerCalls: + return StreamingDataAccessLayerCalls( + read_dlo_deltas=self._read_dlo_deltas, + read_dmo_deltas=self._read_dmo_deltas, + write_dlo_deltas=frozenset(self._write_dlo_deltas_instances), + ) + class ImportVisitor(ast.NodeVisitor): """AST Visitor that extracts external package imports from Python code.""" @@ -301,23 +380,51 @@ def write_requirements_file(file_path: str) -> str: return requirements_path -def scan_file(file_path: str) -> DataAccessLayerCalls: - """Scan a single Python file for Client read/write method calls.""" +def _visit_file(file_path: str) -> ClientMethodVisitor: + """Parse a Python file and return the populated method visitor.""" with open(file_path, "r") as f: - code = f.read() - tree = ast.parse(code) - visitor = ClientMethodVisitor() - visitor.visit(tree) - return visitor.found() + tree = ast.parse(f.read()) + visitor = ClientMethodVisitor() + visitor.visit(tree) + return visitor + + +def scan_file(file_path: str) -> DataAccessLayerCalls: + """Scan a single Python file for batch Client read/write method calls.""" + return _visit_file(file_path).found() + + +def scan_file_streaming(file_path: str) -> StreamingDataAccessLayerCalls: + """Scan a single Python file for StreamingClient delta method calls.""" + return _visit_file(file_path).found_streaming() + +def file_is_streaming(file_path: str) -> bool: + """Return whether the entrypoint uses streaming (delta) access methods.""" + return _visit_file(file_path).is_streaming -def dc_config_json_from_file(file_path: str, type: str) -> dict[str, Any]: - """Create a Data Cloud Custom Code config JSON from a script.""" + +def dc_config_json_from_file( + file_path: str, type: str, streaming: bool = False +) -> dict[str, Any]: + """Create a Data Cloud Custom Code config JSON from a script. + + Args: + file_path: Path to the entrypoint. + type: Package type, ``"script"`` or ``"function"``. + streaming: For scripts, a streaming + (``streamingSource``) config instead of a batch one. + """ config: dict[str, Any] if type == "script": - config = DATA_TRANSFORM_CONFIG_TEMPLATE.copy() + template = ( + STREAMING_TRANSFORM_CONFIG_TEMPLATE + if streaming + else DATA_TRANSFORM_CONFIG_TEMPLATE + ) + config = copy.deepcopy(template) elif type == "function": - config = FUNCTION_CONFIG_TEMPLATE.copy() + config = copy.deepcopy(FUNCTION_CONFIG_TEMPLATE) config["entryPoint"] = os.path.basename(file_path) return config @@ -372,22 +479,53 @@ def update_config(file_path: str) -> dict[str, Any]: if package_type == "script": existing_config["dataspace"] = get_dataspace(existing_config) - output = scan_file(file_path) - read: dict[str, list[str]] = {} - if output.read_dlo: - read["dlo"] = list(output.read_dlo) - else: - read["dmo"] = list(output.read_dmo) - write: dict[str, list[str]] = {} - if output.write_to_dlo: - write["dlo"] = list(output.write_to_dlo) + if file_is_streaming(file_path): + _update_streaming_config(existing_config, file_path) else: - write["dmo"] = list(output.write_to_dmo) - - existing_config["permissions"] = {"read": read, "write": write} + existing_config.pop("streamingSource", None) + output = scan_file(file_path) + read: dict[str, list[str]] = {} + if output.read_dlo: + read["dlo"] = list(output.read_dlo) + else: + read["dmo"] = list(output.read_dmo) + write: dict[str, list[str]] = {} + if output.write_to_dlo: + write["dlo"] = list(output.write_to_dlo) + else: + write["dmo"] = list(output.write_to_dmo) + + existing_config["permissions"] = {"read": read, "write": write} return existing_config +def _update_streaming_config( + existing_config: dict[str, Any], file_path: str +) -> None: + output = scan_file_streaming(file_path) + read_layer = output.read_layer + + source = existing_config.get("streamingSource") + if not isinstance(source, dict): + source = {} + source_name = source.get("name", "") + existing_config["streamingSource"] = {"type": read_layer, "name": source_name} + + if not source_name: + logger.warning( + "streamingSource.name is empty in config.json. A streaming " + "transform must declare its read source; set streamingSource.name " + "to the DLO/DMO the transform reads from." + ) + + read_names = [source_name] if source_name else [] + write_names = list(output.write_dlo_deltas) + existing_config["permissions"] = { + "read": {read_layer: read_names}, + "write": {"dlo": write_names}, + } + + def get_dataspace(existing_config: dict[str, str]) -> str: if "dataspace" in existing_config: dataspace_value = existing_config["dataspace"] diff --git a/src/datacustomcode/template.py b/src/datacustomcode/template.py index 6807510..a543575 100644 --- a/src/datacustomcode/template.py +++ b/src/datacustomcode/template.py @@ -23,8 +23,12 @@ script_template_dir = os.path.join(os.path.dirname(__file__), "templates", "script") function_template_dir = os.path.join(os.path.dirname(__file__), "templates", "function") +STREAMING_EXAMPLE_ENTRYPOINT = os.path.join( + script_template_dir, "examples", "streaming_deltas", "entrypoint.py" +) -def copy_script_template(target_dir: str) -> None: + +def copy_script_template(target_dir: str, streaming: bool = False) -> None: """Copy the template to the target directory.""" os.makedirs(target_dir, exist_ok=True) @@ -39,6 +43,14 @@ def copy_script_template(target_dir: str) -> None: logger.debug(f"Copying file {source} to {destination}...") shutil.copy2(source, destination) + if streaming: + destination = os.path.join(target_dir, "payload", "entrypoint.py") + logger.debug( + f"Copying streaming example {STREAMING_EXAMPLE_ENTRYPOINT} to " + f"{destination}..." + ) + shutil.copy2(STREAMING_EXAMPLE_ENTRYPOINT, destination) + def copy_function_template(target_dir: str, use_in_feature: Optional[str]) -> None: os.makedirs(target_dir, exist_ok=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 7765560..ad55810 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,11 +46,14 @@ def test_init_command( result = runner.invoke(init, ["test_dir", "--code-type", "script"]) assert result.exit_code == 0 - mock_copy.assert_called_once_with("test_dir") + # A script with no --use-in-feature defaults to batch (streaming=False). + mock_copy.assert_called_once_with("test_dir", streaming=False) # Verify SDK config was written mock_write_sdk.assert_called_once_with("test_dir", {"type": "script"}) mock_scan.assert_called_once_with( - os.path.join("test_dir", "payload", "entrypoint.py"), "script" + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=False, ) mock_update.assert_called_once_with( os.path.join("test_dir", "payload", "entrypoint.py") @@ -69,6 +72,96 @@ def test_init_command( expected_content = json.dumps(mock_update.return_value, indent=2) assert expected_content in written_content + @patch("datacustomcode.template.copy_script_template") + @patch("datacustomcode.scan.update_config") + @patch("datacustomcode.scan.dc_config_json_from_file") + @patch("datacustomcode.scan.write_sdk_config") + @patch("builtins.open", new_callable=mock_open) + def test_init_command_streaming( + self, mock_file, mock_write_sdk, mock_scan, mock_update, mock_copy + ): + """Test init command with --use-in-feature StreamingTransform.""" + mock_scan.return_value = {"streamingSource": {"type": "dlo", "name": ""}} + mock_update.return_value = {"streamingSource": {"type": "dlo", "name": ""}} + + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "--use-in-feature", + "StreamingTransform", + ], + ) + + assert result.exit_code == 0 + mock_copy.assert_called_once_with("test_dir", streaming=True) + mock_scan.assert_called_once_with( + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=True, + ) + + @patch("datacustomcode.template.copy_script_template") + @patch("datacustomcode.scan.update_config") + @patch("datacustomcode.scan.dc_config_json_from_file") + @patch("datacustomcode.scan.write_sdk_config") + @patch("builtins.open", new_callable=mock_open) + def test_init_command_batch_explicit( + self, mock_file, mock_write_sdk, mock_scan, mock_update, mock_copy + ): + """Test init command with explicit --use-in-feature BatchTransform.""" + mock_scan.return_value = {"permissions": {"read": {}, "write": {}}} + mock_update.return_value = {"permissions": {"read": {}, "write": {}}} + + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "-u", + "BatchTransform", + ], + ) + + assert result.exit_code == 0 + mock_copy.assert_called_once_with("test_dir", streaming=False) + mock_scan.assert_called_once_with( + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=False, + ) + + def test_init_command_invalid_use_in_feature(self): + """A bad --use-in-feature value for a script aborts with an error.""" + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "--use-in-feature", + "NotARealOption", + ], + ) + + assert result.exit_code != 0 + assert "Invalid --use-in-feature" in result.output + class TestDeploy: @patch("datacustomcode.deploy.deploy_full") diff --git a/tests/test_deploy.py b/tests/test_deploy.py index af804d3..9c203b8 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -17,6 +17,7 @@ DloPermission, DmoPermission, Permissions, + StreamingSource, get_config, ) @@ -671,6 +672,56 @@ def test_create_deployment_function_invoke_options(self, mock_make_api_call): assert isinstance(result, CreateDeploymentResponse) assert result.fileUploadUrl == "https://upload.example.com" + @patch("datacustomcode.deploy._make_api_call") + def test_create_deployment_default_path_no_invoke_options( + self, mock_make_api_call + ): + """Without invokeOptions, the deployment uses the default v63.0 path.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + mock_make_api_call.return_value = { + "fileUploadUrl": "https://upload.example.com" + } + + create_deployment(access_token, metadata) + + url = mock_make_api_call.call_args[0][0] + assert "v63.0" in url + body = mock_make_api_call.call_args[1]["json"] + assert "invokeOptions" not in body + + @patch("datacustomcode.deploy._make_api_call") + def test_create_deployment_invoke_options_routes_to_v67(self, mock_make_api_call): + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + invokeOptions=["StreamingTransform"], + ) + mock_make_api_call.return_value = { + "fileUploadUrl": "https://upload.example.com" + } + + create_deployment(access_token, metadata) + + url = mock_make_api_call.call_args[0][0] + assert "v67.0" in url + body = mock_make_api_call.call_args[1]["json"] + assert body["invokeOptions"] == ["StreamingTransform"] + class TestZip: @patch("datacustomcode.deploy.has_nonempty_requirements_file") @@ -1353,6 +1404,151 @@ def test_create_data_transform_dmo_missing_data_objects_raises( "/test/dir", access_token, metadata, data_transform_config ) + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy._make_api_call") + def test_create_data_transform_batch_type( + self, mock_make_api_call, mock_get_config + ): + """A non-streaming transform sends type BATCH.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="batch_job", + version="1.0.0", + description="Batch job", + computeType="CPU_M", + codeType="script", + ) + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_make_api_call.return_value = {"id": "transform_id"} + + create_data_transform( + "/test/dir", access_token, metadata, data_transform_config + ) + + request_body = mock_make_api_call.call_args[1]["json"] + assert request_body["type"] == "BATCH" + + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy._make_api_call") + def test_create_data_transform_streaming_type( + self, mock_make_api_call, mock_get_config + ): + """A streaming (streamingSource) transform sends type STREAMING.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="streaming_job", + version="1.0.0", + description="Streaming job", + computeType="CPU_M", + codeType="script", + ) + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_make_api_call.return_value = {"id": "transform_id"} + + create_data_transform( + "/test/dir", access_token, metadata, data_transform_config + ) + + request_body = mock_make_api_call.call_args[1]["json"] + assert request_body["type"] == "STREAMING" + manifest = request_body["definition"]["manifest"] + assert manifest["sources"] == {"source1": {"relation_name": "input_dlo"}} + assert manifest["nodes"]["node1"]["relation_name"] == "output_dlo" + + +class TestDataTransformConfigStreaming: + """The streamingSource field and its effect on layer validation.""" + + def test_is_streaming_true_when_source_present(self): + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is True + + def test_is_streaming_false_when_source_absent(self): + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is False + + def test_streaming_allows_dmo_read_dlo_write(self): + """Streaming may read a DMO change feed and write a DLO (writes are + DLO-only)""" + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), + permissions=Permissions( + read=DmoPermission(dmo=["input_dmo__dlm"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is True + assert isinstance(config.permissions.read, DmoPermission) + assert isinstance(config.permissions.write, DloPermission) + + def test_streaming_rejects_dmo_write(self): + """Streaming writes must target a DLO.""" + with pytest.raises(ValueError, match="must write to a DLO"): + DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), + ), + ) + + def test_batch_still_rejects_mixed_layers(self): + """Without a streamingSource, mixed read/write layers are still rejected.""" + with pytest.raises(ValueError, match="both reference"): + DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), + ), + ) + class TestDeployFull: @patch("datacustomcode.deploy.get_config") @@ -1466,6 +1662,96 @@ def test_deploy_full_client_credentials( ) assert result == access_token + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy.create_data_transform") + @patch("datacustomcode.deploy.wait_for_deployment") + @patch("datacustomcode.deploy.upload_zip") + @patch("datacustomcode.deploy.zip") + @patch("datacustomcode.deploy.create_deployment") + def test_deploy_full_streaming_sets_invoke_options( + self, + mock_create_deployment, + mock_zip, + mock_upload_zip, + mock_wait, + mock_create_transform, + mock_get_config, + ): + """A streaming config makes deploy_full set invokeOptions on the + metadata before creating the deployment.""" + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_get_config.return_value = data_transform_config + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + mock_create_deployment.return_value = CreateDeploymentResponse( + fileUploadUrl="https://upload.example.com" + ) + + deploy_full("/test/dir", metadata, access_token, "default") + + assert metadata.invokeOptions == ["StreamingTransform"] + + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy.create_data_transform") + @patch("datacustomcode.deploy.wait_for_deployment") + @patch("datacustomcode.deploy.upload_zip") + @patch("datacustomcode.deploy.zip") + @patch("datacustomcode.deploy.create_deployment") + def test_deploy_full_batch_leaves_invoke_options_unset( + self, + mock_create_deployment, + mock_zip, + mock_upload_zip, + mock_wait, + mock_create_transform, + mock_get_config, + ): + """A batch config must not set invokeOptions (stays on the v63.0 path).""" + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_get_config.return_value = data_transform_config + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + mock_create_deployment.return_value = CreateDeploymentResponse( + fileUploadUrl="https://upload.example.com" + ) + + deploy_full("/test/dir", metadata, access_token, "default") + + assert metadata.invokeOptions is None + class TestRunDataTransform: @patch("datacustomcode.deploy._make_api_call") diff --git a/tests/test_scan.py b/tests/test_scan.py index 2acbc25..716c233 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -12,8 +12,10 @@ SDK_CONFIG_FILE, DataAccessLayerCalls, dc_config_json_from_file, + file_is_streaming, scan_file, scan_file_for_imports, + scan_file_streaming, update_config, write_requirements_file, write_sdk_config, @@ -339,6 +341,27 @@ def test_dlo_to_dlo_config(self): os.remove(sdk_config_path) os.rmdir(os.path.dirname(sdk_config_path)) + def test_streaming_config_scaffolds_streaming_source(self): + """dc_config_json_from_file(streaming=True) scaffolds streamingSource.""" + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + result = dc_config_json_from_file(temp_path, "script", streaming=True) + assert result["entryPoint"] == os.path.basename(temp_path) + assert result["dataspace"] == "default" + assert result["sdkVersion"] == get_version() + assert result["streamingSource"] == {"type": "dlo", "name": ""} + finally: + os.unlink(temp_path) + + def test_batch_config_has_no_streaming_source(self): + """The default (batch) script template has no streamingSource.""" + temp_path = create_test_script("x = 1\n") + try: + result = dc_config_json_from_file(temp_path, "script") + assert "streamingSource" not in result + finally: + os.unlink(temp_path) + def test_dmo_to_dmo_config(self): """Test generating config JSON for DMO to DMO operations.""" content = textwrap.dedent( @@ -694,6 +717,190 @@ def my_function(event, context): os.rmdir(os.path.dirname(sdk_config_path)) +STREAMING_DLO_ENTRYPOINT = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + transformed = deltas.withColumn("x", deltas.x) + query = client.write_dlo_deltas("Account_copy__dll", transformed) + query.awaitTermination() + """ +) + +STREAMING_DMO_ENTRYPOINT = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dmo_deltas() + query = client.write_dlo_deltas("Account_copy__dll", deltas) + query.awaitTermination() + """ +) + + +class TestStreamingScan: + """Tests for streaming (delta) detection and extraction.""" + + def test_file_is_streaming_true_for_delta_methods(self): + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + assert file_is_streaming(temp_path) is True + finally: + os.unlink(temp_path) + + def test_file_is_streaming_false_for_batch(self): + content = textwrap.dedent( + """ + from datacustomcode.client import Client + + client = Client() + df = client.read_dlo("input_dlo") + client.write_to_dlo("output_dlo", df, "overwrite") + """ + ) + temp_path = create_test_script(content) + try: + assert file_is_streaming(temp_path) is False + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_dlo(self): + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + result = scan_file_streaming(temp_path) + assert result.read_dlo_deltas is True + assert result.read_dmo_deltas is False + assert result.read_layer == "dlo" + assert "Account_copy__dll" in result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_dmo_read(self): + temp_path = create_test_script(STREAMING_DMO_ENTRYPOINT) + try: + result = scan_file_streaming(temp_path) + assert result.read_dmo_deltas is True + assert result.read_layer == "dmo" + assert "Account_copy__dll" in result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_requires_read(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + client.write_dlo_deltas("Account_copy__dll", some_df) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="at least one DLO or DMO delta"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_requires_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="must write to at least one DLO"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + +class TestStreamingUpdateConfig: + """Tests for update_config on streaming entrypoints.""" + + def _run(self, entrypoint_src: str, initial_config: dict) -> dict: + temp_path = create_test_script(entrypoint_src) + file_dir = os.path.dirname(temp_path) + config_path = os.path.join(file_dir, "config.json") + sdk_config_path = create_sdk_config(file_dir, "script") + try: + with open(config_path, "w") as f: + json.dump(initial_config, f) + return update_config(temp_path) + finally: + os.remove(temp_path) + if os.path.exists(config_path): + os.remove(config_path) + if os.path.exists(sdk_config_path): + os.remove(sdk_config_path) + os.rmdir(os.path.dirname(sdk_config_path)) + + def test_preserves_streaming_source_name(self): + """config.json's streamingSource.name is authoritative and preserved.""" + updated = self._run( + STREAMING_DLO_ENTRYPOINT, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "Account_Home__dll"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert updated["streamingSource"] == { + "type": "dlo", + "name": "Account_Home__dll", + } + assert updated["permissions"]["read"]["dlo"] == ["Account_Home__dll"] + assert updated["permissions"]["write"]["dlo"] == ["Account_copy__dll"] + + def test_reasserts_layer_from_code(self): + """A code read layer of DMO overrides a stale DLO type in config.""" + updated = self._run( + STREAMING_DMO_ENTRYPOINT, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "Account__dlm"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert updated["streamingSource"]["type"] == "dmo" + assert updated["streamingSource"]["name"] == "Account__dlm" + assert updated["permissions"]["read"]["dmo"] == ["Account__dlm"] + + def test_batch_entrypoint_drops_stale_streaming_source(self): + """Switching a streaming entrypoint to batch clears streamingSource.""" + batch_src = textwrap.dedent( + """ + from datacustomcode.client import Client + + client = Client() + df = client.read_dlo("input_dlo") + client.write_to_dlo("output_dlo", df, "overwrite") + """ + ) + updated = self._run( + batch_src, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "stale__dll"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert "streamingSource" not in updated + assert updated["permissions"]["read"]["dlo"] == ["input_dlo"] + + class TestDataAccessLayerCalls: """Tests for the DataAccessLayerCalls class directly.""" From 5738e807e4472c6ebc8d16fd7caef894d59f45a4 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Thu, 23 Jul 2026 10:21:16 -0400 Subject: [PATCH 9/9] lint --- src/datacustomcode/scan.py | 5 +---- tests/test_deploy.py | 4 +--- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index cdb6434..059925a 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -23,7 +23,6 @@ Any, ClassVar, Dict, - Optional, Set, Union, ) @@ -499,9 +498,7 @@ def update_config(file_path: str) -> dict[str, Any]: return existing_config -def _update_streaming_config( - existing_config: dict[str, Any], file_path: str -) -> None: +def _update_streaming_config(existing_config: dict[str, Any], file_path: str) -> None: output = scan_file_streaming(file_path) read_layer = output.read_layer diff --git a/tests/test_deploy.py b/tests/test_deploy.py index 9c203b8..05063a4 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -673,9 +673,7 @@ def test_create_deployment_function_invoke_options(self, mock_make_api_call): assert result.fileUploadUrl == "https://upload.example.com" @patch("datacustomcode.deploy._make_api_call") - def test_create_deployment_default_path_no_invoke_options( - self, mock_make_api_call - ): + def test_create_deployment_default_path_no_invoke_options(self, mock_make_api_call): """Without invokeOptions, the deployment uses the default v63.0 path.""" access_token = AccessTokenResponse( access_token="test_token", instance_url="https://instance.example.com"