From d2fc1f5da7f1250a05049e2db16b3a1a532c0472 Mon Sep 17 00:00:00 2001 From: Steffan Byrne Date: Thu, 2 Jul 2026 12:28:18 -0400 Subject: [PATCH 1/2] Added auto_write* methods to writers and clients. These new methods will be used with streaming transforms when running in specific modes where the system must control the write mode. --- src/datacustomcode/client.py | 37 +++++++++++++++++++++++ src/datacustomcode/io/writer/base.py | 14 +++++++++ src/datacustomcode/io/writer/csv.py | 6 ++++ src/datacustomcode/io/writer/print.py | 7 +++++ tests/spark/test_session_provider.py | 6 ++++ tests/test_client.py | 42 +++++++++++++++++++++++++++ 6 files changed, 112 insertions(+) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 143796d..739090f 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, @@ -206,6 +207,20 @@ def __str__(self) -> str: return msg +_VALID_RUN_MODES = frozenset({"BATCH", "INITIAL_SYNC", "REBUILD", "DELTA_SYNC"}) + + +def get_run_mode() -> str: + """Read and validate the BYOC_RUN_MODE env var; default to BATCH when unset.""" + run_mode = os.getenv("BYOC_RUN_MODE", "BATCH").upper() + if run_mode not in _VALID_RUN_MODES: + raise ValueError( + f"Invalid BYOC_RUN_MODE: {run_mode}. Must be one of: " + f"{', '.join(sorted(_VALID_RUN_MODES))}." + ) + return run_mode + + class Client: """Entrypoint for accessing DataCloud objects. @@ -382,6 +397,17 @@ def write_to_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 auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write a PySpark DataFrame to a DLO in Data Cloud automatically picking + the WriteMode. + For use with streaming transforms when running in rebuild or initial sync mode. + Args: + name: The name of the DLO to write to. + dataframe: The PySpark DataFrame to write. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) + return self._writer.auto_write_to_dlo(name, dataframe) + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs ) -> None: @@ -395,6 +421,17 @@ 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 auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write a PySpark DataFrame to a DMO in Data Cloud automatically picking + the WriteMode. + For use with streaming transforms when running in rebuild or initial sync mode. + Args: + name: The name of the DMO to write to. + dataframe: The PySpark DataFrame to write. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) + return self._writer.auto_write_to_dmo(name, dataframe) + def write_dlo_deltas( self, name: str, dataframe: PySparkDataFrame, **kwargs ) -> StreamingQuery: diff --git a/src/datacustomcode/io/writer/base.py b/src/datacustomcode/io/writer/base.py index 47a7bd2..9562912 100644 --- a/src/datacustomcode/io/writer/base.py +++ b/src/datacustomcode/io/writer/base.py @@ -54,11 +54,25 @@ def write_to_dlo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: ... + @abstractmethod + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write to a DLO automatically picking the write mode. + + For use with streaming transforms when running in rebuild or initial sync mode. + """ + @abstractmethod def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: ... + @abstractmethod + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write to a DMO automatically picking the write mode. + + For use with streaming transforms when running in rebuild or initial sync mode. + """ + def write_dlo_deltas( self, name: str, dataframe: PySparkDataFrame ) -> StreamingQuery: diff --git a/src/datacustomcode/io/writer/csv.py b/src/datacustomcode/io/writer/csv.py index 3d037d9..7505026 100644 --- a/src/datacustomcode/io/writer/csv.py +++ b/src/datacustomcode/io/writer/csv.py @@ -36,6 +36,9 @@ def write_to_dlo( name = f"{name}{SUFFIX}" dataframe.write.csv(name, mode=write_mode) + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + self.write_to_dlo(name, dataframe, WriteMode.OVERWRITE) + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: @@ -43,3 +46,6 @@ def write_to_dmo( if not name.lower().endswith(SUFFIX): name = f"{name}{SUFFIX}" dataframe.write.csv(name, mode=write_mode) + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + self.write_to_dmo(name, dataframe, WriteMode.OVERWRITE) diff --git a/src/datacustomcode/io/writer/print.py b/src/datacustomcode/io/writer/print.py index c4d2a75..19ef117 100644 --- a/src/datacustomcode/io/writer/print.py +++ b/src/datacustomcode/io/writer/print.py @@ -122,6 +122,10 @@ def write_to_dlo( dataframe.show() + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + self.validate_dataframe_columns_against_dlo(dataframe, name) + dataframe.show() + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: @@ -130,3 +134,6 @@ def write_to_dmo( # so just show the dataframe. dataframe.show() + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + dataframe.show() diff --git a/tests/spark/test_session_provider.py b/tests/spark/test_session_provider.py index 71f0e70..b0c7d90 100644 --- a/tests/spark/test_session_provider.py +++ b/tests/spark/test_session_provider.py @@ -52,11 +52,17 @@ def write_to_dlo( ) -> None: # type: ignore[override] raise NotImplementedError + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + raise NotImplementedError + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: # type: ignore[override] raise NotImplementedError + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + raise NotImplementedError + class FakeProvider(BaseSparkSessionProvider): CONFIG_NAME = "FakeProvider" diff --git a/tests/test_client.py b/tests/test_client.py index e661cbc..beccd0c 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 @@ -10,6 +11,7 @@ DataCloudAccessLayerException, DataCloudObjectType, einstein_predict_col, + get_run_mode, llm_gateway_generate_text_col, ) from datacustomcode.config import ( @@ -46,11 +48,17 @@ def write_to_dlo( ) -> None: pass + def auto_write_to_dlo(self, name: str, dataframe: DataFrame) -> None: + pass + def write_to_dmo( self, name: str, dataframe: DataFrame, write_mode: WriteMode, **kwargs ) -> None: pass + def auto_write_to_dmo(self, name: str, dataframe: DataFrame) -> None: + pass + @pytest.fixture def mock_spark(): @@ -179,6 +187,18 @@ def test_write_to_dlo(self, reset_client, mock_spark): "test_dlo", mock_df, WriteMode.APPEND, extra_param=True ) + def test_auto_write_to_dlo(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = Client(reader=reader, writer=writer) + client._record_dlo_access("some_dlo") + + client.auto_write_to_dlo("test_dlo", mock_df) + + writer.auto_write_to_dlo.assert_called_once_with("test_dlo", mock_df) + def test_write_to_dmo(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) @@ -193,6 +213,18 @@ def test_write_to_dmo(self, reset_client, mock_spark): "test_dmo", mock_df, WriteMode.OVERWRITE, extra_param=True ) + def test_auto_write_to_dmo(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = Client(reader=reader, writer=writer) + client._record_dmo_access("some_dmo") + + client.auto_write_to_dmo("test_dmo", mock_df) + + writer.auto_write_to_dmo.assert_called_once_with("test_dmo", mock_df) + def test_read_dlo_deltas(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) @@ -331,6 +363,16 @@ def test_read_pattern_flow(self, reset_client, mock_spark): assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + @patch.dict(os.environ, {}, clear=True) + def test_get_run_mode_default_batch(self, reset_client, mock_spark): + + assert get_run_mode() == "BATCH" + + @patch.dict(os.environ, {"BYOC_RUN_MODE": "INITIAL_SYNC"}) + def test_get_run_mode(self, reset_client, mock_spark): + + assert get_run_mode() == "INITIAL_SYNC" + class TestClientLlmGatewayGenerateText: From a0936565a9adacbe271371dffa1ea05d1631f5a9 Mon Sep 17 00:00:00 2001 From: Steffan Byrne Date: Thu, 2 Jul 2026 17:05:50 -0400 Subject: [PATCH 2/2] Made RunMode an Enum. --- src/datacustomcode/client.py | 15 +++--- .../examples/streaming_deltas/entrypoint.py | 47 ++++++++++++------- tests/test_client.py | 5 +- 3 files changed, 40 insertions(+), 27 deletions(-) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 739090f..663c548 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -207,18 +207,17 @@ def __str__(self) -> str: return msg -_VALID_RUN_MODES = frozenset({"BATCH", "INITIAL_SYNC", "REBUILD", "DELTA_SYNC"}) +class RunMode(Enum): + BATCH = "BATCH" + INITIAL_SYNC = "INITIAL_SYNC" + REBUILD = "REBUILD" + DELTA_SYNC = "DELTA_SYNC" -def get_run_mode() -> str: +def get_run_mode() -> RunMode: """Read and validate the BYOC_RUN_MODE env var; default to BATCH when unset.""" run_mode = os.getenv("BYOC_RUN_MODE", "BATCH").upper() - if run_mode not in _VALID_RUN_MODES: - raise ValueError( - f"Invalid BYOC_RUN_MODE: {run_mode}. Must be one of: " - f"{', '.join(sorted(_VALID_RUN_MODES))}." - ) - return run_mode + return RunMode(run_mode) class Client: diff --git a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py index 8919cf7..afdf70c 100644 --- a/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -16,32 +16,45 @@ 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 +The local ``datacustomcode run`` readers/writers raise ``NotImplementedError`` for the delta methods. """ +from pyspark.sql import DataFrame from pyspark.sql.functions import col, upper -from datacustomcode.client import Client +from datacustomcode.client import ( + Client, + RunMode, + get_run_mode, +) def main(): client = Client() - - # Streaming DataFrame over the source DLO's change feed. - deltas = client.read_dlo_deltas("Account_std__dll") - - # 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) - - # Drive the query's lifecycle. In the streaming runtime this blocks until - # the job is stopped by the platform. - query.awaitTermination() + source_dlo = "Account_std__dll" + target_dlo = "Account_std_copy__dll" + if get_run_mode() == RunMode.DELTA_SYNC: + # Streaming DataFrame over the source DLO's change feed. + dataframe = client.read_dlo_deltas(source_dlo) + # Ordinary PySpark transform. + transformed = transform(dataframe) + + # 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(target_dlo, transformed) + + # Drive the query's lifecycle. In the streaming runtime this blocks until + # the job is stopped by the platform. + query.awaitTermination() + else: + dataframe = client.read_dlo(source_dlo) + transformed = transform(dataframe) + client.auto_write_to_dlo(target_dlo, transformed) + + +def transform(dataframe: DataFrame) -> DataFrame: + return dataframe.withColumn("description__c", upper(col("description__c"))) if __name__ == "__main__": diff --git a/tests/test_client.py b/tests/test_client.py index beccd0c..d74319a 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -10,6 +10,7 @@ Client, DataCloudAccessLayerException, DataCloudObjectType, + RunMode, einstein_predict_col, get_run_mode, llm_gateway_generate_text_col, @@ -366,12 +367,12 @@ def test_read_pattern_flow(self, reset_client, mock_spark): @patch.dict(os.environ, {}, clear=True) def test_get_run_mode_default_batch(self, reset_client, mock_spark): - assert get_run_mode() == "BATCH" + assert get_run_mode() == RunMode.BATCH @patch.dict(os.environ, {"BYOC_RUN_MODE": "INITIAL_SYNC"}) def test_get_run_mode(self, reset_client, mock_spark): - assert get_run_mode() == "INITIAL_SYNC" + assert get_run_mode() == RunMode.INITIAL_SYNC class TestClientLlmGatewayGenerateText: