From b470208196535d85f7aa2b57be708800bfdaacbe Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 1 Jul 2026 10:54:44 -0400 Subject: [PATCH 01/17] 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 02/17] 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 03/17] 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 04/17] 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 05/17] 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 06/17] 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 07/17] 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 08/17] 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 09/17] 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" From 80261cb5b31d88f2b27cdda2908897a4168be249 Mon Sep 17 00:00:00 2001 From: Diksha Date: Thu, 30 Jul 2026 18:54:42 +0530 Subject: [PATCH 10/17] Add support for external-callout --- src/datacustomcode/client.py | 90 +++ src/datacustomcode/config.yaml | 6 + src/datacustomcode/deploy.py | 4 +- src/datacustomcode/function/runtime.py | 16 + .../named_credential/__init__.py | 28 + src/datacustomcode/named_credential/base.py | 72 +++ .../named_credential/default.py | 122 ++++ .../named_credential/direct/__init__.py | 19 + .../named_credential/direct/auth.py | 56 ++ .../named_credential/direct/credentials.py | 115 ++++ .../named_credential/direct/transport.py | 113 ++++ .../named_credential/direct/url_resolver.py | 109 ++++ src/datacustomcode/named_credential/errors.py | 36 ++ .../named_credential/spark_base.py | 91 +++ .../named_credential/spark_default.py | 169 ++++++ .../named_credential/types/__init__.py | 14 + .../named_credential/types/http_method.py | 29 + .../named_credential/types/http_request.py | 63 +++ .../types/http_request_builder.py | 55 ++ .../named_credential/types/http_response.py | 46 ++ .../types/http_response_builder.py | 24 + src/datacustomcode/named_credential_config.py | 105 ++++ src/datacustomcode/run.py | 7 + .../config.json | 3 + tests/test_named_credential.py | 520 ++++++++++++++++++ tests/test_named_credential_direct.py | 296 ++++++++++ tests/test_runtime_named_credential.py | 128 +++++ 27 files changed, 2334 insertions(+), 2 deletions(-) create mode 100644 src/datacustomcode/named_credential/__init__.py create mode 100644 src/datacustomcode/named_credential/base.py create mode 100644 src/datacustomcode/named_credential/default.py create mode 100644 src/datacustomcode/named_credential/direct/__init__.py create mode 100644 src/datacustomcode/named_credential/direct/auth.py create mode 100644 src/datacustomcode/named_credential/direct/credentials.py create mode 100644 src/datacustomcode/named_credential/direct/transport.py create mode 100644 src/datacustomcode/named_credential/direct/url_resolver.py create mode 100644 src/datacustomcode/named_credential/errors.py create mode 100644 src/datacustomcode/named_credential/spark_base.py create mode 100644 src/datacustomcode/named_credential/spark_default.py create mode 100644 src/datacustomcode/named_credential/types/__init__.py create mode 100644 src/datacustomcode/named_credential/types/http_method.py create mode 100644 src/datacustomcode/named_credential/types/http_request.py create mode 100644 src/datacustomcode/named_credential/types/http_request_builder.py create mode 100644 src/datacustomcode/named_credential/types/http_response.py create mode 100644 src/datacustomcode/named_credential/types/http_response_builder.py create mode 100644 src/datacustomcode/named_credential_config.py create mode 100644 src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json create mode 100644 tests/test_named_credential.py create mode 100644 tests/test_named_credential_direct.py create mode 100644 tests/test_runtime_named_credential.py diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 0d4df78..9776d42 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -29,6 +29,7 @@ from datacustomcode.file.path.default import DefaultFindFilePath from datacustomcode.io.reader.base import BaseDataCloudReader from datacustomcode.llm_gateway_config import spark_llm_gateway_config +from datacustomcode.named_credential_config import spark_named_credential_config from datacustomcode.spark.default import DefaultSparkSessionProvider if TYPE_CHECKING: @@ -41,6 +42,9 @@ from datacustomcode.io.reader.base import BaseDataCloudReader from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode from datacustomcode.llm_gateway.spark_base import SparkLLMGateway + from datacustomcode.named_credential.spark_base import SparkNamedCredential + from datacustomcode.named_credential.types.http_request import HTTPRequest + from datacustomcode.named_credential.types.http_response import HTTPResponse from datacustomcode.spark.base import BaseSparkSessionProvider @@ -118,6 +122,21 @@ def _build_spark_einstein_predictions() -> "SparkEinsteinPredictions": return cfg.to_object() +def _build_spark_named_credential() -> "SparkNamedCredential": + """Instantiate the SDK-configured :class:`SparkNamedCredential`. + + Raises: + RuntimeError: If no ``spark_named_credential_config`` has been loaded. + """ + cfg = spark_named_credential_config.spark_named_credential_config + if cfg is None: + raise RuntimeError( + "spark_named_credential_config is not configured. Add a " + "'spark_named_credential_config' section to config.yaml." + ) + return cfg.to_object() + + def einstein_predict_col( model_api_name: str, prediction_type: "PredictionType", @@ -167,6 +186,39 @@ def einstein_predict_col( ) +def named_credential_request_col( + request: "HTTPRequest", + body: Optional["Column"] = None, +) -> "Column": + """Build a Spark Column that makes one Named Credential callout per row. + + The endpoint, method, and headers are fixed for the call (taken from + ``request``); only ``body`` varies per row. Use this instead of + :meth:`Client.named_credential_request` when the callout runs across a + DataFrame so each row is dispatched independently rather than one-shot on + the driver. + + The returned Column yields a struct ``{status, response, error_code, + error_message}`` for each row. ``response`` is itself a struct + ``{status_code, body, headers}``. Use ``[...]`` to pick a field, e.g. + ``named_credential_request_col(...)["response"]["status_code"]``. Per-row + failures populate ``status`` / ``error_code`` / ``error_message`` so a + single bad row does not abort the whole Spark job. + + Args: + request: The callout template — its symbolic reference, method, and + headers are applied to every row. + body: Optional per-row ``Column`` holding the JSON request body as a + string (or null for no body). + + Returns: + A Spark ``Column`` of ``StructType`` with fields ``status``, + ``response``, ``error_code``, and ``error_message``. + """ + named_credential = Client()._get_spark_named_credential() + return named_credential.request_col(request, body=body) + + class DataCloudObjectType(Enum): DLO = "dlo" DMO = "dmo" @@ -228,6 +280,7 @@ class Client: spark_llm_gateway: Optional custom :class:`SparkLLMGateway`. spark_einstein_predictions: Optional custom :class:`SparkEinsteinPredictions`. + spark_named_credential: Optional custom :class:`SparkNamedCredential`. Example: >>> client = Client() @@ -243,6 +296,7 @@ class Client: _file: DefaultFindFilePath _spark_llm_gateway: Optional[SparkLLMGateway] _spark_einstein_predictions: Optional[SparkEinsteinPredictions] + _spark_named_credential: Optional[SparkNamedCredential] _data_layer_history: dict[DataCloudObjectType, set[str]] _code_type: str @@ -253,6 +307,7 @@ def __new__( spark_provider: Optional[BaseSparkSessionProvider] = None, spark_llm_gateway: Optional[SparkLLMGateway] = None, spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None, + spark_named_credential: Optional[SparkNamedCredential] = None, code_type: str = "script", ) -> Client: @@ -260,6 +315,7 @@ def __new__( cls._instance = super().__new__(cls) cls._instance._spark_llm_gateway = spark_llm_gateway cls._instance._spark_einstein_predictions = spark_einstein_predictions + cls._instance._spark_named_credential = spark_named_credential # Initialize Readers and Writers from config # and/or provided reader and writer if reader is None or writer is None: @@ -474,6 +530,40 @@ def _get_spark_einstein_predictions(self) -> SparkEinsteinPredictions: self._spark_einstein_predictions = _build_spark_einstein_predictions() return self._spark_einstein_predictions + def named_credential_request( + self, + request: "HTTPRequest", + body: Optional[Dict[str, Any]] = None, + ) -> "HTTPResponse": + """Issue a one-shot Named Credential external callout. This is the + scalar counterpart to :func:`named_credential_request_col`: it runs + **once** on the driver — not per row. Use the column helper method + instead when you want to fan a callout out across every row of a + DataFrame. + + Example: + + >>> from datacustomcode.named_credential.types.http_request_builder \\ + ... import HTTPRequestBuilder + >>> request = ( + ... HTTPRequestBuilder().set_url("callout:NC/search").build() + ... ) + >>> response = Client().named_credential_request(request) + + Args: + request: The callout request + body: Optional JSON-serializable request body. + + Returns: + The external service's response. + """ + return self._get_spark_named_credential().request(request, body=body) + + def _get_spark_named_credential(self) -> SparkNamedCredential: + if self._spark_named_credential is None: + self._spark_named_credential = _build_spark_named_credential() + return self._spark_named_credential + def _validate_data_layer_history_does_not_contain( self, data_cloud_object_type: DataCloudObjectType ) -> None: diff --git a/src/datacustomcode/config.yaml b/src/datacustomcode/config.yaml index e7ed0c9..e9a8a61 100644 --- a/src/datacustomcode/config.yaml +++ b/src/datacustomcode/config.yaml @@ -34,3 +34,9 @@ llm_gateway_config: spark_llm_gateway_config: type_config_name: DefaultSparkLLMGateway + +named_credential_config: + type_config_name: DefaultNamedCredential + +spark_named_credential_config: + type_config_name: DefaultSparkNamedCredential diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index e8c4ec4..c65d43c 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -598,9 +598,9 @@ def zip( with zipfile.ZipFile(ZIP_FILE_NAME, "w", zipfile.ZIP_DEFLATED) as zipf: for root, dirs, files in os.walk(directory): - # Skip .DS_Store files when adding to zip + # Skip .DS_Store and local credentials. for file in files: - if file != ".DS_Store": + if file not in (".DS_Store", "credential.json"): abs_path = os.path.join(root, file) arcname = os.path.relpath(abs_path, directory) zipf.write(abs_path, arcname) diff --git a/src/datacustomcode/function/runtime.py b/src/datacustomcode/function/runtime.py index df9cb1e..d8164a2 100644 --- a/src/datacustomcode/function/runtime.py +++ b/src/datacustomcode/function/runtime.py @@ -23,6 +23,8 @@ from datacustomcode.function.base import BaseRuntime from datacustomcode.llm_gateway.base import LLMGateway from datacustomcode.llm_gateway_config import llm_gateway_config +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential_config import named_credential_config class Runtime(BaseRuntime): @@ -69,6 +71,7 @@ def __init__(self) -> None: self._llm_gateway: Optional[LLMGateway] = None self._file = DefaultFindFilePath() self._einstein_predictions: Optional[EinsteinPredictions] = None + self._named_credential: Optional[NamedCredential] = None @property def llm_gateway(self) -> LLMGateway: @@ -98,3 +101,16 @@ def einstein_predictions(self) -> EinsteinPredictions: einstein_predictions_config.einstein_predictions_config.to_object() ) return self._einstein_predictions + + @property + def named_credential(self) -> NamedCredential: + if self._named_credential is None: + if named_credential_config.named_credential_config is None: + raise RuntimeError( + "Named Credential is not configured. Add " + "'named_credential_config' section to config.yaml" + ) + self._named_credential = ( + named_credential_config.named_credential_config.to_object() + ) + return self._named_credential diff --git a/src/datacustomcode/named_credential/__init__.py b/src/datacustomcode/named_credential/__init__.py new file mode 100644 index 0000000..31245db --- /dev/null +++ b/src/datacustomcode/named_credential/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential.default import DefaultNamedCredential +from datacustomcode.named_credential.errors import NamedCredentialCallError +from datacustomcode.named_credential.spark_base import SparkNamedCredential +from datacustomcode.named_credential.spark_default import DefaultSparkNamedCredential + +__all__ = [ + "DefaultNamedCredential", + "DefaultSparkNamedCredential", + "NamedCredential", + "NamedCredentialCallError", + "SparkNamedCredential", +] diff --git a/src/datacustomcode/named_credential/base.py b/src/datacustomcode/named_credential/base.py new file mode 100644 index 0000000..246ae9f --- /dev/null +++ b/src/datacustomcode/named_credential/base.py @@ -0,0 +1,72 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, +) + +from datacustomcode.mixin import UserExtendableNamedConfigMixin + +if TYPE_CHECKING: + from datacustomcode.named_credential.types.http_request import HTTPRequest + from datacustomcode.named_credential.types.http_response import HTTPResponse + + +class NamedCredential(ABC, UserExtendableNamedConfigMixin): + CONFIG_NAME: str + + def __init__(self, **kwargs): + pass + + @abstractmethod + def request( + self, + request: HTTPRequest, + body: Optional[Dict[str, Any]] = None, + ) -> HTTPResponse: + """Make an external callout through a Named Credential. + + The endpoint and its authentication are resolved server-side from the + Named Credential referenced by ``request.url``; the function never sees + the external credential. + + Args: + request: The callout request + body: Optional JSON-serializable request body. + + Returns: + The external service's response. + """ + ... + + def callout_json( + self, + request: HTTPRequest, + body: Optional[str] = None, + ) -> Dict[str, Any]: + """Low-level string-in/string-out callout returning the raw response. + + It is required only by the per-row Spark path which forwards the + raw body per row; the default signals it as unsupported. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement callout_json(); it " + "supports only one-shot request()." + ) diff --git a/src/datacustomcode/named_credential/default.py b/src/datacustomcode/named_credential/default.py new file mode 100644 index 0000000..d3e06f0 --- /dev/null +++ b/src/datacustomcode/named_credential/default.py @@ -0,0 +1,122 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from typing import ( + Any, + Dict, + Optional, +) + +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential.types.http_request import HTTPRequest +from datacustomcode.named_credential.types.http_response import HTTPResponse +from datacustomcode.named_credential.types.http_response_builder import ( + HTTPResponseBuilder, +) + + +class DefaultNamedCredential(NamedCredential): + """ + Executes the callout directly via :class:`DirectCalloutTransport`, resolving + the URL from the Named Credential Connect API (falling back to + ``credential.json``) and injecting auth from ``credential.json``. + """ + + CONFIG_NAME = "DefaultNamedCredential" + + def __init__( + self, + credentials_profile: str = "default", + sf_cli_org: Optional[str] = None, + credential_file: Optional[str] = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self._credentials_profile = credentials_profile + self._sf_cli_org = sf_cli_org + self._credential_file = credential_file + self._transport: Optional[Any] = None + + def request( + self, + request: HTTPRequest, + body: Optional[Dict[str, Any]] = None, + ) -> HTTPResponse: + callout_response = self.callout_json( + request, json.dumps(body) if body is not None else "" + ) + + raw_body = callout_response.get("body") or "" + data: Optional[Any] = None + if raw_body: + try: + # Preserve any JSON value: object, array, or scalar. + data = json.loads(raw_body) + except json.JSONDecodeError: + data = None + + response_dict = { + "status_code": callout_response.get("http_status_code"), + "headers": callout_response.get("headers", {}), + "data": data, + } + return HTTPResponseBuilder.build(response_dict) + + def callout_json( + self, + request: HTTPRequest, + body: Optional[str] = None, + ) -> Dict[str, Any]: + """Raw string-in/string-out callout returning the unparsed response. + + Unlike :meth:`request`, the ``body`` is sent verbatim (never re-parsed) + and the response is returned as a raw ``{http_status_code, headers, + body}`` dict rather than a parsed :class:`HTTPResponse`. + """ + # Callout request shape sent to the transport. + callout_request = { + "path": request.url, + "method": request.method, + "headers": dict(request.headers), + "body": body if body is not None else "", + } + callout_response = self._callout(callout_request) + return { + "http_status_code": callout_response.get("http_status_code"), + "headers": callout_response.get("headers", {}), + "body": callout_response.get("body") or "", + } + + def _callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: + """Execute the callout via the transport. + + Returns a dict with ``http_status_code``, ``headers``, and ``body``. + """ + result: Dict[str, Any] = self._get_transport().callout(callout_request) + return result + + def _get_transport(self) -> Any: + if self._transport is None: + from datacustomcode.named_credential.direct.transport import ( + DirectCalloutTransport, + ) + + self._transport = DirectCalloutTransport( + credentials_profile=self._credentials_profile, + sf_cli_org=self._sf_cli_org, + credential_file=self._credential_file, + ) + return self._transport diff --git a/src/datacustomcode/named_credential/direct/__init__.py b/src/datacustomcode/named_credential/direct/__init__.py new file mode 100644 index 0000000..b8df664 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Named Credential callout path. + +This resolves the endpoint via the NamedCredential Connect API +and injects auth from a developer-provided ``credential.json`` +""" diff --git a/src/datacustomcode/named_credential/direct/auth.py b/src/datacustomcode/named_credential/direct/auth.py new file mode 100644 index 0000000..4893ead --- /dev/null +++ b/src/datacustomcode/named_credential/direct/auth.py @@ -0,0 +1,56 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Inject external credential auth into an outgoing request. +""" + +from __future__ import annotations + +import base64 +from typing import TYPE_CHECKING, Any, Dict + +from requests.auth import AuthBase + +from datacustomcode.named_credential.direct.credentials import AuthType + +if TYPE_CHECKING: + from requests.models import PreparedRequest + + +class DynamicAuthHandler(AuthBase): + def __init__(self, cred_config: Dict[str, Any]) -> None: + self.config = cred_config + self.auth_type = cred_config.get("auth_type") + + def __call__(self, request: PreparedRequest) -> PreparedRequest: + if self.auth_type == AuthType.BASIC.value: + user = self.config.get("username", "") + pwd = self.config.get("password", "") + token = base64.b64encode(f"{user}:{pwd}".encode()).decode() + request.headers["Authorization"] = f"Basic {token}" + + elif self.auth_type == AuthType.CUSTOM.value: + for name, value in self.config.get("custom_headers", {}).items(): + request.headers[name] = value + + elif self.auth_type in (AuthType.OAUTH.value, AuthType.JWT.value): + bearer = self.config.get("access_token") or self.config.get("token") + if bearer: + request.headers["Authorization"] = f"Bearer {bearer}" + + else: + raise ValueError(f"Unsupported auth_type '{self.auth_type}'.") + + return request diff --git a/src/datacustomcode/named_credential/direct/credentials.py b/src/datacustomcode/named_credential/direct/credentials.py new file mode 100644 index 0000000..7c9ce66 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/credentials.py @@ -0,0 +1,115 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Load developer credentials for direct callouts from ``credential.json``. + +The file lives in the parent of the payload folder so it is never packaged in +the deployment zip. Keys are callout references (``callout:``); +each value carries a mandatory ``auth_type`` and an optional ``target_url``. +""" + +from __future__ import annotations + +from enum import Enum +import json +import os +from pathlib import Path +from typing import ( + Any, + Dict, + Optional, +) + +# Default file name; discovered in the parent of the payload folder. +DEFAULT_CREDENTIAL_FILE = "credential.json" +# Absolute-path override, primarily for tests and non-standard layouts. +CREDENTIAL_FILE_ENV_VAR = "DATACUSTOMCODE_CREDENTIAL_FILE" + + +class AuthType(str, Enum): + """External Credential auth types supported by External Services.""" + + BASIC = "Basic" + CUSTOM = "Custom" + JWT = "Jwt" + OAUTH = "OAuth" + + +class CredentialError(RuntimeError): + """Raised when credentials cannot be found or are invalid.""" + + +def _discover_credential_file() -> Optional[Path]: + """Find ``credential.json`` via env override, then by walking up from cwd.""" + override = os.environ.get(CREDENTIAL_FILE_ENV_VAR) + if override: + return Path(override) + + for directory in (Path.cwd(), *Path.cwd().parents): + candidate = directory / DEFAULT_CREDENTIAL_FILE + if candidate.is_file(): + return candidate + return None + + +class CredentialStore: + """Reads ``credential.json`` and returns per-callout configuration.""" + + def __init__(self, credential_file: Optional[str] = None) -> None: + self._explicit_path = Path(credential_file) if credential_file else None + self._credentials: Optional[Dict[str, Dict[str, Any]]] = None + + def _load(self) -> Dict[str, Dict[str, Any]]: + if self._credentials is not None: + return self._credentials + + path = self._explicit_path or _discover_credential_file() + if path is None or not path.is_file(): + raise CredentialError( + f"Could not find '{DEFAULT_CREDENTIAL_FILE}'. Place it in the " + f"parent of your payload folder, or set " + f"${CREDENTIAL_FILE_ENV_VAR} to its path." + ) + try: + with open(path, "r") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise CredentialError(f"Failed to read '{path}': {exc}") from exc + + if not isinstance(data, dict): + raise CredentialError( + f"'{path}' must be a JSON object keyed by callout reference." + ) + self._credentials = data + return data + + def get(self, callout_key: str) -> Dict[str, Any]: + """Return the config for a callout key (e.g. ``callout:AWS_S3_Service``). + + Raises: + CredentialError: if the key is missing or has no ``auth_type``. + """ + credentials = self._load() + config = credentials.get(callout_key) + if config is None: + raise CredentialError( + f"No credential configuration found for '{callout_key}'. " + f"Add it to '{DEFAULT_CREDENTIAL_FILE}'." + ) + if not isinstance(config, dict) or not config.get("auth_type"): + raise CredentialError( + f"Credential for '{callout_key}' is missing the mandatory " + f"'auth_type' field." + ) + return config diff --git a/src/datacustomcode/named_credential/direct/transport.py b/src/datacustomcode/named_credential/direct/transport.py new file mode 100644 index 0000000..c03ffc7 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/transport.py @@ -0,0 +1,113 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Send a Named Credential callout over HTTP. + +Resolves the ``callout:/`` reference to a real endpoint +(via the Named Credential Connect API, falling back to ``target_url`` in +``credential.json``), attaches the credential's auth, and returns the raw +``{http_status_code, headers, body}`` response. +""" + +from __future__ import annotations + +from typing import ( + Any, + Dict, + Optional, +) + +import requests + +from datacustomcode.named_credential.direct.auth import DynamicAuthHandler +from datacustomcode.named_credential.direct.credentials import ( + CredentialError, + CredentialStore, +) +from datacustomcode.named_credential.direct.url_resolver import resolve_base_url +from datacustomcode.token_provider import ( + CredentialsTokenProvider, + SFCLITokenProvider, + TokenProvider, +) + + +class DirectCalloutTransport: + def __init__( + self, + credentials_profile: str = "default", + sf_cli_org: Optional[str] = None, + credential_file: Optional[str] = None, + ) -> None: + self._store = CredentialStore(credential_file) + self._token_provider = self._build_token_provider( + credentials_profile, sf_cli_org + ) + # Resolved base URL per callout key. Stable for the transport's life, so + # cache it to avoid a token fetch + Connect API call on every row + self._base_url_cache: Dict[str, str] = {} + + @staticmethod + def _build_token_provider( + credentials_profile: str, sf_cli_org: Optional[str] + ) -> TokenProvider: + if sf_cli_org: + return SFCLITokenProvider(sf_cli_org) + return CredentialsTokenProvider(credentials_profile) + + def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: + raw_url = callout_request["path"] + if not raw_url.startswith("callout:"): + raise CredentialError( + f"Callout URL must start with 'callout:', got '{raw_url}'." + ) + + # Split the named credential reference at the first '/' or '?'; the remainder + # path or query string is appended to the resolved base URL verbatim. + sep_idx = min( + (i for i in (raw_url.find("/"), raw_url.find("?")) if i != -1), + default=len(raw_url), + ) + callout_key = raw_url[:sep_idx] + path_suffix = raw_url[sep_idx:] + + if callout_key == "callout:": + raise CredentialError( + f"Named Credential name is empty in URL '{raw_url}'." + ) + + cred_config = self._store.get(callout_key) + base_url = self._base_url_cache.get(callout_key) + if base_url is None: + base_url = resolve_base_url(callout_key, cred_config, self._token_provider) + self._base_url_cache[callout_key] = base_url + + body = callout_request.get("body") or None + headers = dict(callout_request.get("headers", {})) + if body and "Content-Type" not in headers: + headers["Content-Type"] = "application/json" + + response = requests.request( + method=callout_request["method"], + url=base_url + path_suffix, + headers=headers, + data=body, + auth=DynamicAuthHandler(cred_config), + timeout=30, + ) + return { + "http_status_code": response.status_code, + "headers": dict(response.headers), + "body": response.text, + } diff --git a/src/datacustomcode/named_credential/direct/url_resolver.py b/src/datacustomcode/named_credential/direct/url_resolver.py new file mode 100644 index 0000000..a22f889 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/url_resolver.py @@ -0,0 +1,109 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Resolve a ``callout:`` reference to a real base URL. + +Prefers the NamedCredential Connect API ``calloutUrl``; falls back to the +developer-supplied ``target_url`` when the API is unavailable or the Named +Credential is not yet onboarded. +""" + +from __future__ import annotations + +import logging +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, +) + +import requests + +from datacustomcode.named_credential.direct.credentials import CredentialError + +if TYPE_CHECKING: + from datacustomcode.token_provider import TokenProvider + +logger = logging.getLogger(__name__) + +NAMED_CREDENTIAL_PATH = "services/data/v63.0/named-credentials/named-credential-setup/{name}" + + +def _callout_url_from_connect_api( + developer_name: str, token_provider: TokenProvider +) -> Optional[str]: + """Fetch ``calloutUrl`` for a Named Credential, or None on any failure.""" + try: + token = token_provider.get_token() + except Exception as exc: + # No usable token (e.g. not logged in) — expected; fall back quietly. + logger.debug("Could not obtain a token for %s: %s", developer_name, exc) + return None + + try: + path = NAMED_CREDENTIAL_PATH.format(name=developer_name) + url = f"{token.instance_url.rstrip('/')}/{path}" + response = requests.get( + url, + headers={"Authorization": f"Bearer {token.access_token}"}, + timeout=30, + ) + response.raise_for_status() + callout_url = response.json().get("calloutUrl") + return callout_url or None + except Exception as exc: + # A token was obtained but the Connect API call failed — likely a + # misconfiguration (named credential not onboarded, missing permission, wrong org). + # Surface it before silently falling back to target_url. + logger.warning( + "Connect API URL resolution failed for %s: %s. " + "Falling back to 'target_url' from credential.json if set.", + developer_name, + exc, + ) + return None + + +def resolve_base_url( + callout_key: str, + cred_config: Dict[str, Any], + token_provider: Optional[TokenProvider], +) -> str: + """Resolve the base URL for a callout key. + + Args: + callout_key: e.g. ``callout:Nominatim_Geocoding``. + cred_config: The callout's ``credential.json`` entry. + token_provider: Provides a token/instance URL for the Connect API; when + None, only ``target_url`` is used. + + Raises: + CredentialError: if no URL can be resolved. + """ + developer_name = callout_key.split(":", 1)[1] if ":" in callout_key else callout_key + + base_url: Optional[str] = None + if token_provider is not None: + base_url = _callout_url_from_connect_api(developer_name, token_provider) + + if not base_url: + base_url = cred_config.get("target_url") or None + + if not base_url: + raise CredentialError( + f"Could not resolve a URL for '{callout_key}'. Ensure the Named " + f"Credential exists (sf login) or set 'target_url' in credential.json." + ) + return base_url.rstrip("/") diff --git a/src/datacustomcode/named_credential/errors.py b/src/datacustomcode/named_credential/errors.py new file mode 100644 index 0000000..e156fe4 --- /dev/null +++ b/src/datacustomcode/named_credential/errors.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Exceptions raised by Named Credential implementations.""" + +from __future__ import annotations + +from typing import Optional + + +class NamedCredentialCallError(RuntimeError): + """Raised when a Named Credential external callout fails.""" + + def __init__( + self, + message: str, + *, + status: Optional[object] = None, + error_code: Optional[str] = None, + error_message: Optional[str] = None, + ) -> None: + super().__init__(message) + self.status = status + self.error_code = error_code + self.error_message = error_message diff --git a/src/datacustomcode/named_credential/spark_base.py b/src/datacustomcode/named_credential/spark_base.py new file mode 100644 index 0000000..71ac800 --- /dev/null +++ b/src/datacustomcode/named_credential/spark_base.py @@ -0,0 +1,91 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, +) + +from datacustomcode.mixin import UserExtendableNamedConfigMixin + +if TYPE_CHECKING: + from pyspark.sql import Column + + from datacustomcode.named_credential.types.http_request import HTTPRequest + from datacustomcode.named_credential.types.http_response import HTTPResponse + + +class SparkNamedCredential(ABC, UserExtendableNamedConfigMixin): + """Named Credential external callout for script (Spark) code. + + The callout is a one-shot request that runs on the driver. The endpoint and + its authentication are resolved from the Named Credential referenced by + ``request.url``. + """ + + CONFIG_NAME: str + + def __init__(self, **kwargs: Any) -> None: + pass + + @abstractmethod + def request( + self, + request: HTTPRequest, + body: Optional[Dict[str, Any]] = None, + ) -> HTTPResponse: + """Make an external callout through a Named Credential. + + Args: + request: The callout request + body: Optional JSON-serializable request body. + + Returns: + The external service's response. + """ + ... + + @abstractmethod + def request_col( + self, + request: HTTPRequest, + body: Optional["Column"] = None, + ) -> "Column": + """Build a Spark ``Column`` that makes one external callout per row. + + The endpoint, method, and headers are fixed for the call (taken from + ``request``); only ``body`` varies per row. Use this instead of + :meth:`request` when the callout runs across a DataFrame so each row is + dispatched independently rather than one-shot on the driver. + + Args: + request: The callout template + body: Optional per-row ``Column`` holding the JSON request body as a + string (or null for no body). + + Returns: + A ``Column`` yielding a struct + ``{status, response, error_code, error_message}``. ``response`` is + itself a struct ``{status_code, body, headers}`` carrying the callout + response. Select a field with ``[...]``, e.g. + ``request_col(...)["response"]["status_code"]``. Returning a struct + means a single failing row does not abort the Spark job — callers can + inspect ``status`` / ``error_code`` per row instead. + """ + ... diff --git a/src/datacustomcode/named_credential/spark_default.py b/src/datacustomcode/named_credential/spark_default.py new file mode 100644 index 0000000..7d8724e --- /dev/null +++ b/src/datacustomcode/named_credential/spark_default.py @@ -0,0 +1,169 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, +) + +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential.spark_base import SparkNamedCredential + +if TYPE_CHECKING: + from pyspark.sql import Column + + from datacustomcode.named_credential.types.http_request import HTTPRequest + from datacustomcode.named_credential.types.http_response import HTTPResponse + + +_STATUS_SUCCESS = "SUCCESS" +_STATUS_ERROR = "ERROR" + + +def _build_underlying_named_credential() -> "NamedCredential": + """Build the callout object from the configured ``named_credential_config``. + + Raises ``RuntimeError`` if no ``named_credential_config`` section is set. + """ + from datacustomcode.named_credential_config import named_credential_config + + cfg = named_credential_config.named_credential_config + if cfg is None: + raise RuntimeError( + "named_credential_config is not configured. Add a " + "'named_credential_config' section to config.yaml." + ) + return cfg.to_object() + + +class DefaultSparkNamedCredential(SparkNamedCredential): + """ + Callout for Spark, delegating to the shared implementation. + """ + + CONFIG_NAME = "DefaultSparkNamedCredential" + + def __init__( + self, + named_credential: Optional["NamedCredential"] = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if named_credential is None: + named_credential = _build_underlying_named_credential() + self._named_credential: "NamedCredential" = named_credential + + def request( + self, + request: "HTTPRequest", + body: Optional[Dict[str, Any]] = None, + ) -> "HTTPResponse": + return self._named_credential.request(request, body) + + def request_col( + self, + request: "HTTPRequest", + body: Optional["Column"] = None, + ) -> "Column": + """Per-row callout via a client-side Spark UDF. + + Returns a struct ``{status, response, error_code, error_message}`` where + ``response`` is itself a struct + ``{status_code, body, headers}`` carrying the callout's HTTP response, so + a script that selects ``request_col(...)["response"]["status_code"]`` + behaves the same during development and in the Data Cloud runtime. Per-row + failures populate the error fields instead of aborting the Spark job. + """ + from pyspark.sql.functions import lit, udf + from pyspark.sql.types import ( + IntegerType, + MapType, + StringType, + StructField, + StructType, + ) + + http_response_schema = StructType( + [ + StructField("status_code", IntegerType(), True), + StructField("body", StringType(), True), + StructField("headers", MapType(StringType(), StringType()), True), + ] + ) + result_schema = StructType( + [ + StructField("status", StringType(), True), + StructField("response", http_response_schema, True), + StructField("error_code", StringType(), True), + StructField("error_message", StringType(), True), + ] + ) + + # Fail at column-build time rather than turning every row into an opaque error. + if ( + getattr(type(self._named_credential), "callout_json", None) + is NamedCredential.callout_json + ): + raise TypeError( + f"{type(self._named_credential).__name__} does not support the " + "per-row callout path; it must override callout_json(). Use " + "named_credential_request() for a one-shot callout, or configure " + "DefaultNamedCredential." + ) + + def _callout(body_str: Optional[str]) -> Dict[str, Any]: + return _invoke_callout_as_struct( + self._named_credential, request, body_str + ) + + body_col = body if body is not None else lit(None).cast(StringType()) + return udf(_callout, result_schema)(body_col) + + +def _invoke_callout_as_struct( + named_credential: "NamedCredential", + request: "HTTPRequest", + body_str: Optional[str], +) -> Dict[str, Any]: + """Run one callout and shape it into the shared result struct. + + ``response`` is a nested struct ``{status_code, body, headers}`` where + ``body`` is the external response verbatim. Transport errors become per-row + ERROR structs rather than aborting the job. + """ + try: + callout_response = named_credential.callout_json(request, body_str) + except Exception as exc: # surface any transport error per row + return { + "status": _STATUS_ERROR, + "response": None, + "error_code": None, + "error_message": str(exc), + } + + status_code = callout_response.get("http_status_code") + return { + "status": _STATUS_SUCCESS, + "response": { + "status_code": int(status_code) if status_code is not None else None, + "body": callout_response.get("body") or "", + "headers": callout_response.get("headers") or {}, + }, + "error_code": None, + "error_message": None, + } diff --git a/src/datacustomcode/named_credential/types/__init__.py b/src/datacustomcode/named_credential/types/__init__.py new file mode 100644 index 0000000..93988ff --- /dev/null +++ b/src/datacustomcode/named_credential/types/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/src/datacustomcode/named_credential/types/http_method.py b/src/datacustomcode/named_credential/types/http_method.py new file mode 100644 index 0000000..954842a --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_method.py @@ -0,0 +1,29 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum + + +class HTTPMethod(str, Enum): + """Callout methods (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS); usable on + Python 3.10 where ``http.HTTPMethod`` is unavailable.""" + + GET = "GET" + POST = "POST" + PUT = "PUT" + DELETE = "DELETE" + PATCH = "PATCH" + HEAD = "HEAD" + OPTIONS = "OPTIONS" diff --git a/src/datacustomcode/named_credential/types/http_request.py b/src/datacustomcode/named_credential/types/http_request.py new file mode 100644 index 0000000..5f88bd8 --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_request.py @@ -0,0 +1,63 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum +from typing import Dict + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_validator, +) + +from datacustomcode.named_credential.types.http_method import HTTPMethod + + +class HTTPRequest(BaseModel): + """External callout request. The endpoint and its auth are resolved + server-side from the Named Credential referenced by ``url``, which uses + ``callout:/`` syntax.""" + + model_config = ConfigDict(populate_by_name=True) + + url: str = Field( + ..., + min_length=1, + description="Symbolic Named Credential reference, " + "e.g. 'callout:/'", + ) + method: str = Field( + default="GET", + description="HTTP method (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS)", + ) + headers: Dict[str, str] = Field(default_factory=dict, description="Request headers") + + @field_validator("method", mode="before") + @classmethod + def _normalize_method(cls, value: object) -> str: + # Accept str, this module's HTTPMethod, or http.HTTPMethod (3.11+). + if isinstance(value, Enum): + method = str(value.value) + else: + method = str(value) + method = method.upper() + if method not in {m.value for m in HTTPMethod}: + supported = ", ".join(m.value for m in HTTPMethod) + raise ValueError( + f"Unsupported HTTP method '{method}'. " + f"Named Credential callouts support {supported}." + ) + return method diff --git a/src/datacustomcode/named_credential/types/http_request_builder.py b/src/datacustomcode/named_credential/types/http_request_builder.py new file mode 100644 index 0000000..23f4304 --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_request_builder.py @@ -0,0 +1,55 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Union + +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request import HTTPRequest + + +class HTTPRequestBuilder: + def __init__(self) -> None: + self._url = "" + self._method: Union[str, HTTPMethod] = HTTPMethod.GET + self._headers: Dict[str, str] = {} + + def set_url(self, url: str) -> "HTTPRequestBuilder": + """Set the symbolic Named Credential reference. + + Args: + url: e.g. 'callout:/' + """ + self._url = url + return self + + def set_method(self, method: Union[str, HTTPMethod]) -> "HTTPRequestBuilder": + """Set the HTTP method. + + Accepts this module's ``HTTPMethod``, ``http.HTTPMethod`` (Python 3.11+), + or a plain string such as ``"GET"``. + """ + self._method = method + return self + + def set_headers(self, headers: Dict[str, str]) -> "HTTPRequestBuilder": + self._headers = headers + return self + + def build(self) -> HTTPRequest: + return HTTPRequest( + url=self._url, + method=self._method, # type: ignore[arg-type] + headers=self._headers, + ) diff --git a/src/datacustomcode/named_credential/types/http_response.py b/src/datacustomcode/named_credential/types/http_response.py new file mode 100644 index 0000000..6e24643 --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_response.py @@ -0,0 +1,46 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import ( + Any, + Dict, + Optional, +) + +from pydantic import BaseModel, Field + + +class HTTPResponse(BaseModel): + """Response from a Named Credential external callout.""" + + status_code: int = Field(..., description="HTTP status code", ge=0) + headers: Dict[str, str] = Field( + default_factory=dict, description="Response headers" + ) + data: Optional[Any] = Field( + default=None, + description="Parsed JSON response body (object, array, or scalar), " + "or None if the body was empty or not JSON", + ) + + @property + def is_success(self) -> bool: + """Check if the request succeeded (2xx).""" + return 200 <= self.status_code < 300 + + @property + def is_error(self) -> bool: + """Check if the request failed.""" + return not self.is_success diff --git a/src/datacustomcode/named_credential/types/http_response_builder.py b/src/datacustomcode/named_credential/types/http_response_builder.py new file mode 100644 index 0000000..b304a1a --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_response_builder.py @@ -0,0 +1,24 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict + +from datacustomcode.named_credential.types.http_response import HTTPResponse + + +class HTTPResponseBuilder: + @staticmethod + def build(response_dict: Dict[str, Any]) -> HTTPResponse: + return HTTPResponse.model_validate(response_dict) diff --git a/src/datacustomcode/named_credential_config.py b/src/datacustomcode/named_credential_config.py new file mode 100644 index 0000000..34fbf88 --- /dev/null +++ b/src/datacustomcode/named_credential_config.py @@ -0,0 +1,105 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import ( + ClassVar, + Generic, + Type, + TypeVar, + Union, +) + +from datacustomcode.common_config import ( + BaseConfig, + BaseObjectConfig, + default_config_file, +) +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential.spark_base import SparkNamedCredential + +_N = TypeVar("_N", bound=NamedCredential) +_S = TypeVar("_S", bound=SparkNamedCredential) + + +class NamedCredentialObjectConfig(BaseObjectConfig, Generic[_N]): + type_to_create: ClassVar[Type[NamedCredential]] = NamedCredential # type: ignore[type-abstract] + + def to_object(self) -> NamedCredential: + type_ = self.type_to_create.subclass_from_config_name(self.type_config_name) + return type_(**self.options) + + +class NamedCredentialConfig(BaseConfig): + named_credential_config: Union[ + NamedCredentialObjectConfig[NamedCredential], None + ] = None + + def update(self, other: "NamedCredentialConfig") -> "NamedCredentialConfig": + def merge( + config_a: Union[NamedCredentialObjectConfig, None], + config_b: Union[NamedCredentialObjectConfig, None], + ) -> Union[NamedCredentialObjectConfig, None]: + if config_a is not None and config_a.force: + return config_a + if config_b: + return config_b + return config_a + + self.named_credential_config = merge( + self.named_credential_config, other.named_credential_config + ) + return self + + +class SparkNamedCredentialObjectConfig(BaseObjectConfig, Generic[_S]): + type_to_create: ClassVar[Type[SparkNamedCredential]] = SparkNamedCredential # type: ignore[type-abstract] + + def to_object(self) -> SparkNamedCredential: + type_ = self.type_to_create.subclass_from_config_name(self.type_config_name) + return type_(**self.options) + + +class SparkNamedCredentialConfig(BaseConfig): + spark_named_credential_config: Union[ + SparkNamedCredentialObjectConfig[SparkNamedCredential], None + ] = None + + def update( + self, other: "SparkNamedCredentialConfig" + ) -> "SparkNamedCredentialConfig": + def merge( + config_a: Union[SparkNamedCredentialObjectConfig, None], + config_b: Union[SparkNamedCredentialObjectConfig, None], + ) -> Union[SparkNamedCredentialObjectConfig, None]: + if config_a is not None and config_a.force: + return config_a + if config_b: + return config_b + return config_a + + self.spark_named_credential_config = merge( + self.spark_named_credential_config, other.spark_named_credential_config + ) + return self + + +# Global Named Credential config instance +named_credential_config = NamedCredentialConfig() +named_credential_config.load(default_config_file()) + + +# Global Spark Named Credential config instance +spark_named_credential_config = SparkNamedCredentialConfig() +spark_named_credential_config.load(default_config_file()) diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index 006055c..c644406 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -27,6 +27,7 @@ from datacustomcode.config import config from datacustomcode.einstein_predictions_config import einstein_predictions_config from datacustomcode.llm_gateway_config import llm_gateway_config +from datacustomcode.named_credential_config import named_credential_config from datacustomcode.scan import find_base_directory, get_package_type @@ -55,6 +56,9 @@ def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]): _set_config_option( llm_gateway_config.llm_gateway_config, config_key, sf_cli_org ) + _set_config_option( + named_credential_config.named_credential_config, config_key, sf_cli_org + ) elif profile != "default": config_key = "credentials_profile" _set_config_option(config.reader_config, config_key, profile) @@ -63,6 +67,9 @@ def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]): einstein_predictions_config.einstein_predictions_config, config_key, profile ) _set_config_option(llm_gateway_config.llm_gateway_config, config_key, profile) + _set_config_option( + named_credential_config.named_credential_config, config_key, profile + ) def run_entrypoint( diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json new file mode 100644 index 0000000..2f911e9 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json @@ -0,0 +1,3 @@ +{ + "entryPoint": "entrypoint.py" +} diff --git a/tests/test_named_credential.py b/tests/test_named_credential.py new file mode 100644 index 0000000..51bd43d --- /dev/null +++ b/tests/test_named_credential.py @@ -0,0 +1,520 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import sys +from unittest.mock import MagicMock, patch + +from pydantic import ValidationError +import pytest + +from datacustomcode.named_credential.default import DefaultNamedCredential +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request import HTTPRequest +from datacustomcode.named_credential.types.http_request_builder import ( + HTTPRequestBuilder, +) +from datacustomcode.named_credential.types.http_response import HTTPResponse +from datacustomcode.named_credential.types.http_response_builder import ( + HTTPResponseBuilder, +) + + +class TestHTTPRequest: + def test_url_required(self): + with pytest.raises(ValidationError): + HTTPRequest() + + def test_url_min_length(self): + with pytest.raises(ValidationError): + HTTPRequest(url="") + + def test_method_defaults_to_get(self): + request = HTTPRequest(url="callout:NC/path") + assert request.method == "GET" + + def test_headers_default_empty(self): + request = HTTPRequest(url="callout:NC/path") + assert request.headers == {} + + def test_method_accepts_local_enum(self): + request = HTTPRequest(url="callout:NC/path", method=HTTPMethod.POST) + assert request.method == "POST" + + def test_method_normalizes_lowercase_string(self): + request = HTTPRequest(url="callout:NC/path", method="post") + assert request.method == "POST" + + @pytest.mark.skipif( + sys.version_info < (3, 11), reason="http.HTTPMethod added in 3.11" + ) + def test_method_accepts_stdlib_httpmethod(self): + from http import HTTPMethod as StdHTTPMethod + + request = HTTPRequest(url="callout:NC/path", method=StdHTTPMethod.GET) + assert request.method == "GET" + + @pytest.mark.parametrize("method", ["PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]) + def test_non_get_post_methods_accepted(self, method): + request = HTTPRequest(url="callout:NC/path", method=method) + assert request.method == method + + def test_method_accepts_local_enum_put(self): + request = HTTPRequest(url="callout:NC/path", method=HTTPMethod.PUT) + assert request.method == "PUT" + + def test_unsupported_method_rejected(self): + with pytest.raises(ValidationError): + HTTPRequest(url="callout:NC/path", method="TRACE") + + def test_query_string_in_url_preserved(self): + request = HTTPRequest(url="callout:NC/v1/accounts?limit=10&active=true") + assert request.url == "callout:NC/v1/accounts?limit=10&active=true" + + +class TestHTTPRequestBuilder: + def test_builder_basic_usage(self): + request = ( + HTTPRequestBuilder() + .set_url("callout:Nominatim_Geocoding/search?q=sf&format=json") + .set_method(HTTPMethod.GET) + .set_headers({"Accept": "application/json"}) + .build() + ) + assert request.url == "callout:Nominatim_Geocoding/search?q=sf&format=json" + assert request.method == "GET" + assert request.headers == {"Accept": "application/json"} + + def test_builder_default_method_is_get(self): + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + assert request.method == "GET" + + def test_builder_accepts_string_method(self): + request = ( + HTTPRequestBuilder().set_url("callout:NC/path").set_method("post").build() + ) + assert request.method == "POST" + + def test_builder_validates_on_build(self): + with pytest.raises(ValidationError): + HTTPRequestBuilder().set_url("").build() + + +class TestHTTPResponse: + def test_response_defaults(self): + response = HTTPResponse(status_code=200) + assert response.headers == {} + assert response.data is None + + def test_is_success_for_2xx(self): + assert HTTPResponse(status_code=200).is_success is True + assert HTTPResponse(status_code=204).is_success is True + assert HTTPResponse(status_code=299).is_success is True + + def test_is_error_for_non_2xx(self): + assert HTTPResponse(status_code=404).is_success is False + assert HTTPResponse(status_code=404).is_error is True + assert HTTPResponse(status_code=500).is_error is True + + def test_status_code_validation(self): + with pytest.raises(ValidationError): + HTTPResponse(status_code=-1) + + +class TestHTTPResponseBuilder: + def test_build_from_dict(self): + response = HTTPResponseBuilder.build( + {"status_code": 200, "headers": {"X": "y"}, "data": {"ok": True}} + ) + assert isinstance(response, HTTPResponse) + assert response.status_code == 200 + assert response.headers == {"X": "y"} + assert response.data == {"ok": True} + + def test_build_requires_status_code(self): + with pytest.raises(ValidationError): + HTTPResponseBuilder.build({"headers": {}}) + + +class TestDefaultNamedCredential: + def test_callout_delegates_to_transport(self, monkeypatch): + nc = DefaultNamedCredential() + + class _FakeTransport: + def callout(self, callout_request): + return {"http_status_code": 200, "headers": {}, "body": "{}"} + + monkeypatch.setattr(nc, "_get_transport", lambda: _FakeTransport()) + result = nc._callout({"path": "callout:NC/x", "method": "GET"}) + assert result["http_status_code"] == 200 + + def test_request_translates_json_response(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return { + "http_status_code": 200, + "headers": {"Content-Type": "application/json"}, + "body": '{"result": "ok"}', + } + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = ( + HTTPRequestBuilder() + .set_url("callout:NC/search?q=sf") + .set_method(HTTPMethod.POST) + .set_headers({"Accept": "application/json"}) + .build() + ) + response = nc.request(request, {"key": "value"}) + + # The query string travels inside the path verbatim. + assert captured["path"] == "callout:NC/search?q=sf" + assert captured["method"] == "POST" + assert captured["headers"] == {"Accept": "application/json"} + assert "query_params" not in captured + assert captured["body"] == '{"key": "value"}' + + assert response.status_code == 200 + assert response.headers == {"Content-Type": "application/json"} + assert response.data == {"result": "ok"} + assert response.is_success is True + + def test_request_without_body_sends_empty_string(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return {"http_status_code": 204, "headers": {}, "body": ""} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + response = nc.request(request) + + assert captured["body"] == "" + assert response.status_code == 204 + assert response.data is None + + def test_request_non_json_body_yields_none_data(self, monkeypatch): + nc = DefaultNamedCredential() + + def fake_callout(callout_request): + return {"http_status_code": 200, "headers": {}, "body": "plain text"} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + response = nc.request(request) + assert response.data is None + + def test_request_json_array_body_preserved(self, monkeypatch): + nc = DefaultNamedCredential() + + def fake_callout(callout_request): + return {"http_status_code": 200, "headers": {}, "body": "[1, 2, 3]"} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + response = nc.request(request) + assert response.data == [1, 2, 3] + + def test_request_json_scalar_body_preserved(self, monkeypatch): + nc = DefaultNamedCredential() + + def fake_callout(callout_request): + return {"http_status_code": 200, "headers": {}, "body": "42"} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + response = nc.request(request) + assert response.data == 42 + + def test_request_empty_dict_body_sends_object(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return {"http_status_code": 200, "headers": {}, "body": ""} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + nc.request(request, {}) + + # An explicit empty dict is a body ("{}"), distinct from None (""). + assert captured["body"] == "{}" + + def test_callout_json_forwards_body_and_returns_raw_shape(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return { + "http_status_code": 200, + "headers": {"X": "y"}, + "body": "plain text", + } + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = ( + HTTPRequestBuilder() + .set_url("callout:NC/path") + .set_method(HTTPMethod.POST) + .build() + ) + + result = nc.callout_json(request, "a=1&b=2") + + # Body is sent verbatim, not re-encoded as JSON. + assert captured["body"] == "a=1&b=2" + # Response is the raw {http_status_code, headers, body} shape, body intact. + assert result == { + "http_status_code": 200, + "headers": {"X": "y"}, + "body": "plain text", + } + + def test_callout_json_none_body_sends_empty_string(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return {"http_status_code": 204, "headers": {}, "body": ""} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + + result = nc.callout_json(request) + + assert captured["body"] == "" + assert result == {"http_status_code": 204, "headers": {}, "body": ""} + + +class TestDefaultSparkNamedCredential: + def test_delegates_to_underlying(self): + from datacustomcode.named_credential.spark_default import ( + DefaultSparkNamedCredential, + ) + + sentinel = object() + + class _Underlying: + def __init__(self): + self.calls = [] + + def request(self, request, body=None): + self.calls.append((request, body)) + return sentinel + + underlying = _Underlying() + spark_nc = DefaultSparkNamedCredential(named_credential=underlying) + + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + result = spark_nc.request(request, {"k": "v"}) + + assert result is sentinel + assert underlying.calls == [(request, {"k": "v"})] + + def test_builds_underlying_from_config_when_absent(self, monkeypatch): + from datacustomcode.named_credential import spark_default + + built = object() + monkeypatch.setattr( + spark_default, "_build_underlying_named_credential", lambda: built + ) + spark_nc = spark_default.DefaultSparkNamedCredential() + assert spark_nc._named_credential is built + + def test_config_resolves_to_spark_default(self): + from datacustomcode.named_credential.spark_base import SparkNamedCredential + from datacustomcode.named_credential.spark_default import ( + DefaultSparkNamedCredential, + ) + + resolved = SparkNamedCredential.subclass_from_config_name( + "DefaultSparkNamedCredential" + ) + assert resolved is DefaultSparkNamedCredential + + +class TestDefaultSparkNamedCredentialRequestCol: + """The client-side per-row UDF path used during development.""" + + @patch("pyspark.sql.functions.udf") + @patch("pyspark.sql.functions.lit") + def test_wraps_callout_in_udf_over_body_column(self, mock_lit, mock_udf): + from datacustomcode.named_credential.spark_default import ( + DefaultSparkNamedCredential, + ) + + sentinel_udf = MagicMock(name="udf") + sentinel_applied = MagicMock(name="udf_applied") + sentinel_udf.return_value = sentinel_applied + mock_udf.return_value = sentinel_udf + + underlying = MagicMock() + underlying.callout_json.return_value = { + "http_status_code": 200, + "headers": {"Content-Type": "application/json"}, + "body": '{"ok":true}', + } + spark_nc = DefaultSparkNamedCredential(named_credential=underlying) + + request = HTTPRequestBuilder().set_url("callout:NC/v1/accounts").build() + body_col = MagicMock(name="body_col") + result = spark_nc.request_col(request, body_col) + + assert result is sentinel_applied + mock_udf.assert_called_once() + # The caller's body column is what the UDF is applied to. + sentinel_udf.assert_called_once_with(body_col) + mock_lit.assert_not_called() + + # Exercise the wrapped callout function on a row body. + udf_fn = mock_udf.call_args.args[0] + out = udf_fn('{"name": "acme"}') + + assert out["status"] == "SUCCESS" + assert out["error_code"] is None + assert out["error_message"] is None + # response mirrors the production UDF: a struct of status_code/body/headers + payload = out["response"] + assert payload["status_code"] == 200 + assert payload["headers"] == {"Content-Type": "application/json"} + # Body is forwarded verbatim, never re-parsed/re-serialized. + assert payload["body"] == '{"ok":true}' + + # The raw body string (NOT a parsed dict) is forwarded to the callout. + sent_request, sent_body = underlying.callout_json.call_args.args + assert sent_request is request + assert sent_body == '{"name": "acme"}' + + @patch("pyspark.sql.functions.udf") + @patch("pyspark.sql.functions.lit") + def test_defaults_body_to_typed_null_column(self, mock_lit, mock_udf): + from datacustomcode.named_credential.spark_default import ( + DefaultSparkNamedCredential, + ) + + null_col = MagicMock(name="null_col") + lit_none = MagicMock(name="lit_none") + lit_none.cast.return_value = null_col + mock_lit.return_value = lit_none + + sentinel_udf = MagicMock(name="udf") + mock_udf.return_value = sentinel_udf + + spark_nc = DefaultSparkNamedCredential(named_credential=MagicMock()) + request = HTTPRequestBuilder().set_url("callout:NC/status").build() + spark_nc.request_col(request) + + # With no body column, a typed null string column is applied instead. + sentinel_udf.assert_called_once_with(null_col) + + +class TestInvokeCalloutAsStruct: + """The callout-to-struct shaping shared by every row.""" + + def test_success_struct_carries_response_fields(self): + from datacustomcode.named_credential.spark_default import ( + _invoke_callout_as_struct, + ) + + underlying = MagicMock() + underlying.callout_json.return_value = { + "http_status_code": 201, + "headers": {"X-Trace": "abc"}, + "body": "[1,2,3]", + } + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + + out = _invoke_callout_as_struct(underlying, request, '{"a": 1}') + + assert out["status"] == "SUCCESS" + # response is a typed struct mirroring the runtime UDF's fields. + assert out["response"] == { + "status_code": 201, + "body": "[1,2,3]", + "headers": {"X-Trace": "abc"}, + } + # The request body is forwarded verbatim (never validated as JSON). + assert underlying.callout_json.call_args.args[1] == '{"a": 1}' + + def test_non_json_request_body_is_forwarded_not_rejected(self): + from datacustomcode.named_credential.spark_default import ( + _invoke_callout_as_struct, + ) + + # Mirrors the runtime: an opaque body (form-encoded / plain text) is sent + # as-is instead of being rejected the way JSON parsing would. + underlying = MagicMock() + underlying.callout_json.return_value = { + "http_status_code": 200, + "headers": {}, + "body": "OK", + } + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + + out = _invoke_callout_as_struct(underlying, request, "a=1&b=2") + + assert out["status"] == "SUCCESS" + assert underlying.callout_json.call_args.args[1] == "a=1&b=2" + # A non-JSON response body is preserved verbatim, not dropped to "". + assert out["response"]["body"] == "OK" + + def test_none_body_forwards_none_and_keeps_all_keys(self): + from datacustomcode.named_credential.spark_default import ( + _invoke_callout_as_struct, + ) + + underlying = MagicMock() + underlying.callout_json.return_value = { + "http_status_code": 200, + "headers": {}, + "body": "", + } + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + + out = _invoke_callout_as_struct(underlying, request, None) + + assert out["status"] == "SUCCESS" + # Empty headers/body keep their keys, matching the runtime UDF's + # includingDefaultValueFields defaults. + assert out["response"] == { + "status_code": 200, + "body": "", + "headers": {}, + } + # A null body column forwards None (not "null") to the callout. + assert underlying.callout_json.call_args.args[1] is None + + def test_transport_error_yields_error_struct(self): + from datacustomcode.named_credential.spark_default import ( + _invoke_callout_as_struct, + ) + + underlying = MagicMock() + underlying.callout_json.side_effect = RuntimeError("proxy down") + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + + out = _invoke_callout_as_struct(underlying, request, '{"a": 1}') + + assert out["status"] == "ERROR" + assert out["response"] is None + assert out["error_message"] == "proxy down" diff --git a/tests/test_named_credential_direct.py b/tests/test_named_credential_direct.py new file mode 100644 index 0000000..045afb5 --- /dev/null +++ b/tests/test_named_credential_direct.py @@ -0,0 +1,296 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import base64 +import json + +import pytest +from requests.models import PreparedRequest + +from datacustomcode.named_credential.direct.auth import DynamicAuthHandler +from datacustomcode.named_credential.direct.credentials import ( + CREDENTIAL_FILE_ENV_VAR, + AuthType, + CredentialError, + CredentialStore, +) +from datacustomcode.named_credential.direct.url_resolver import resolve_base_url + + +def _prepared_request() -> PreparedRequest: + request = PreparedRequest() + request.prepare(method="GET", url="https://example.com") + return request + + +class TestDynamicAuthHandler: + def test_basic_sets_authorization_header(self): + handler = DynamicAuthHandler( + {"auth_type": AuthType.BASIC.value, "username": "u", "password": "p"} + ) + request = handler(_prepared_request()) + expected = base64.b64encode(b"u:p").decode() + assert request.headers["Authorization"] == f"Basic {expected}" + + def test_custom_injects_custom_headers(self): + handler = DynamicAuthHandler( + { + "auth_type": AuthType.CUSTOM.value, + "custom_headers": {"X-Api-Key": "secret"}, + } + ) + request = handler(_prepared_request()) + assert request.headers["X-Api-Key"] == "secret" + + def test_oauth_sets_bearer_from_access_token(self): + handler = DynamicAuthHandler( + {"auth_type": AuthType.OAUTH.value, "access_token": "tok"} + ) + request = handler(_prepared_request()) + assert request.headers["Authorization"] == "Bearer tok" + + def test_jwt_sets_bearer_from_token(self): + handler = DynamicAuthHandler( + {"auth_type": AuthType.JWT.value, "token": "jwt-tok"} + ) + request = handler(_prepared_request()) + assert request.headers["Authorization"] == "Bearer jwt-tok" + + def test_unsupported_auth_type_raises_value_error(self): + handler = DynamicAuthHandler({"auth_type": "Nonsense"}) + with pytest.raises(ValueError): + handler(_prepared_request()) + + +class TestCredentialStore: + def test_get_returns_config(self, tmp_path, monkeypatch): + cred_file = tmp_path / "credential.json" + cred_file.write_text( + json.dumps({"callout:NC": {"auth_type": "Basic", "username": "u"}}) + ) + monkeypatch.setenv(CREDENTIAL_FILE_ENV_VAR, str(cred_file)) + store = CredentialStore() + config = store.get("callout:NC") + assert config["auth_type"] == "Basic" + assert config["username"] == "u" + + def test_explicit_path_takes_precedence(self, tmp_path): + cred_file = tmp_path / "credential.json" + cred_file.write_text(json.dumps({"callout:NC": {"auth_type": "OAuth"}})) + store = CredentialStore(str(cred_file)) + assert store.get("callout:NC")["auth_type"] == "OAuth" + + def test_missing_file_raises(self, tmp_path): + store = CredentialStore(str(tmp_path / "does-not-exist.json")) + with pytest.raises(CredentialError): + store.get("callout:NC") + + def test_missing_key_raises(self, tmp_path): + cred_file = tmp_path / "credential.json" + cred_file.write_text(json.dumps({"callout:Other": {"auth_type": "Basic"}})) + store = CredentialStore(str(cred_file)) + with pytest.raises(CredentialError): + store.get("callout:NC") + + def test_missing_auth_type_raises(self, tmp_path): + cred_file = tmp_path / "credential.json" + cred_file.write_text(json.dumps({"callout:NC": {"username": "u"}})) + store = CredentialStore(str(cred_file)) + with pytest.raises(CredentialError): + store.get("callout:NC") + + def test_non_object_json_raises(self, tmp_path): + cred_file = tmp_path / "credential.json" + cred_file.write_text(json.dumps(["not", "an", "object"])) + store = CredentialStore(str(cred_file)) + with pytest.raises(CredentialError): + store.get("callout:NC") + + +class _FakeToken: + def __init__(self, access_token="tok", instance_url="https://org.example.com"): + self.access_token = access_token + self.instance_url = instance_url + + +class _FakeTokenProvider: + def __init__(self, token=None): + self._token = token or _FakeToken() + + def get_token(self): + return self._token + + +class TestResolveBaseUrl: + def test_prefers_connect_api_callout_url(self, monkeypatch): + monkeypatch.setattr( + "datacustomcode.named_credential.direct.url_resolver." + "_callout_url_from_connect_api", + lambda name, provider: "https://api.example.com/", + ) + url = resolve_base_url( + "callout:NC", + {"target_url": "https://fallback.example.com"}, + _FakeTokenProvider(), + ) + assert url == "https://api.example.com" + + def test_falls_back_to_target_url(self, monkeypatch): + monkeypatch.setattr( + "datacustomcode.named_credential.direct.url_resolver." + "_callout_url_from_connect_api", + lambda name, provider: None, + ) + url = resolve_base_url( + "callout:NC", + {"target_url": "https://fallback.example.com/"}, + _FakeTokenProvider(), + ) + assert url == "https://fallback.example.com" + + def test_no_url_available_raises(self, monkeypatch): + monkeypatch.setattr( + "datacustomcode.named_credential.direct.url_resolver." + "_callout_url_from_connect_api", + lambda name, provider: None, + ) + with pytest.raises(CredentialError): + resolve_base_url("callout:NC", {}, _FakeTokenProvider()) + + def test_no_token_provider_uses_target_url(self): + url = resolve_base_url( + "callout:NC", {"target_url": "https://only.example.com"}, None + ) + assert url == "https://only.example.com" + + +class TestDirectCalloutTransport: + def _make_transport(self, tmp_path, monkeypatch, cred_entry): + from datacustomcode.named_credential.direct import transport as transport_mod + + cred_file = tmp_path / "credential.json" + cred_file.write_text(json.dumps({"callout:NC": cred_entry})) + monkeypatch.setattr( + transport_mod.DirectCalloutTransport, + "_build_token_provider", + staticmethod(lambda profile, org: _FakeTokenProvider()), + ) + monkeypatch.setattr( + transport_mod, + "resolve_base_url", + lambda key, config, provider: "https://api.example.com", + ) + return transport_mod.DirectCalloutTransport(credential_file=str(cred_file)) + + def test_callout_happy_path(self, tmp_path, monkeypatch): + from datacustomcode.named_credential.direct import transport as transport_mod + + transport = self._make_transport( + tmp_path, + monkeypatch, + {"auth_type": "Custom", "custom_headers": {"X-Key": "v"}}, + ) + + captured = {} + + class _Resp: + def __init__(self): + self.status_code = 200 + self.headers = {"Content-Type": "application/json"} + self.text = '{"ok": true}' + + def fake_request(**kwargs): + captured.update(kwargs) + return _Resp() + + monkeypatch.setattr(transport_mod.requests, "request", fake_request) + + result = transport.callout( + { + "path": "callout:NC/search?q=sf", + "method": "POST", + "headers": {}, + "body": '{"a": 1}', + } + ) + + # The query string embedded in the path is passed through verbatim. + assert captured["url"] == "https://api.example.com/search?q=sf" + assert captured["method"] == "POST" + assert "params" not in captured + assert captured["data"] == '{"a": 1}' + # Content-Type auto-added because a body is present. + assert captured["headers"]["Content-Type"] == "application/json" + assert isinstance(captured["auth"], DynamicAuthHandler) + + assert result["http_status_code"] == 200 + assert result["body"] == '{"ok": true}' + + @pytest.mark.parametrize("method", ["PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]) + def test_callout_forwards_non_get_post_methods(self, tmp_path, monkeypatch, method): + from datacustomcode.named_credential.direct import transport as transport_mod + + transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) + + captured = {} + + class _Resp: + def __init__(self): + self.status_code = 200 + self.headers = {} + self.text = "" + + def fake_request(**kwargs): + captured.update(kwargs) + return _Resp() + + monkeypatch.setattr(transport_mod.requests, "request", fake_request) + transport.callout({"path": "callout:NC/x", "method": method, "headers": {}}) + + assert captured["method"] == method + + def test_callout_rejects_non_callout_url(self, tmp_path, monkeypatch): + transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) + with pytest.raises(CredentialError): + transport.callout({"path": "https://example.com/x", "method": "GET"}) + + def test_base_url_resolved_once_per_callout_key(self, tmp_path, monkeypatch): + from datacustomcode.named_credential.direct import transport as transport_mod + + transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) + + calls = {"count": 0} + + def counting_resolve(key, config, provider): + calls["count"] += 1 + return "https://api.example.com" + + monkeypatch.setattr(transport_mod, "resolve_base_url", counting_resolve) + + class _Resp: + status_code = 200 + headers: dict = {} + text = "" + + monkeypatch.setattr( + transport_mod.requests, "request", lambda **kwargs: _Resp() + ) + + # Many per-row callouts, same callout key: base URL is resolved just once. + for _ in range(5): + transport.callout({"path": "callout:NC/x", "method": "GET", "headers": {}}) + + assert calls["count"] == 1 diff --git a/tests/test_runtime_named_credential.py b/tests/test_runtime_named_credential.py new file mode 100644 index 0000000..d5b81c4 --- /dev/null +++ b/tests/test_runtime_named_credential.py @@ -0,0 +1,128 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import pytest + +from datacustomcode.named_credential.base import NamedCredential +from datacustomcode.named_credential.default import DefaultNamedCredential +from datacustomcode.named_credential.types.http_request import HTTPRequest +from datacustomcode.named_credential.types.http_response import HTTPResponse +from datacustomcode.named_credential_config import ( + NamedCredentialConfig, + NamedCredentialObjectConfig, + named_credential_config, +) + + +class TestNamedCredentialConfig: + def test_default_config_resolves_default_impl(self): + assert named_credential_config.named_credential_config is not None + instance = named_credential_config.named_credential_config.to_object() + assert isinstance(instance, DefaultNamedCredential) + + def test_custom_implementation_is_discoverable(self): + class CustomNamedCredential(NamedCredential): + CONFIG_NAME = "CustomNamedCredential" + + def __init__(self, custom_param: str = "default", **kwargs): + super().__init__(**kwargs) + self.custom_param = custom_param + + def request(self, request, body=None): + return HTTPResponse(status_code=200, data={"echo": self.custom_param}) + + assert "CustomNamedCredential" in NamedCredential.available_config_names() + cls = NamedCredential.subclass_from_config_name("CustomNamedCredential") + assert cls is CustomNamedCredential + + config = NamedCredentialObjectConfig( + type_config_name="CustomNamedCredential", + options={"custom_param": "my_value"}, + ) + instance = config.to_object() + assert isinstance(instance, CustomNamedCredential) + + response = instance.request(HTTPRequest(url="callout:NC/path")) + assert response.data == {"echo": "my_value"} + + +class TestRuntimeNamedCredential: + @pytest.fixture(autouse=True) + def _reset_singleton(self): + from datacustomcode.function.runtime import Runtime + + Runtime._instance = None + yield + Runtime._instance = None + + def test_runtime_exposes_default_named_credential(self): + from datacustomcode.function.runtime import Runtime + + runtime = Runtime() + assert isinstance(runtime.named_credential, DefaultNamedCredential) + + def test_named_credential_is_cached(self): + from datacustomcode.function.runtime import Runtime + + runtime = Runtime() + assert runtime.named_credential is runtime.named_credential + + def test_missing_config_raises(self, monkeypatch): + from datacustomcode.function.runtime import Runtime + + monkeypatch.setattr(named_credential_config, "named_credential_config", None) + runtime = Runtime() + with pytest.raises(RuntimeError, match="Named Credential is not configured"): + _ = runtime.named_credential + + +class TestNamedCredentialConfigUpdate: + def test_update_prefers_other(self): + base = NamedCredentialConfig( + named_credential_config=NamedCredentialObjectConfig( + type_config_name="DefaultNamedCredential" + ) + ) + override = NamedCredentialConfig( + named_credential_config=NamedCredentialObjectConfig( + type_config_name="CustomNamedCredential" + ) + ) + base.update(override) + assert base.named_credential_config.type_config_name == "CustomNamedCredential" + + def test_update_keeps_forced_existing(self): + base = NamedCredentialConfig( + named_credential_config=NamedCredentialObjectConfig( + type_config_name="DefaultNamedCredential", force=True + ) + ) + override = NamedCredentialConfig( + named_credential_config=NamedCredentialObjectConfig( + type_config_name="CustomNamedCredential" + ) + ) + base.update(override) + assert base.named_credential_config.type_config_name == "DefaultNamedCredential" + + def test_update_keeps_existing_when_other_empty(self): + base = NamedCredentialConfig( + named_credential_config=NamedCredentialObjectConfig( + type_config_name="DefaultNamedCredential" + ) + ) + base.update(NamedCredentialConfig()) + assert base.named_credential_config.type_config_name == "DefaultNamedCredential" From 10bb3d5e338d52e5b1b8031e600e18d224f2b5a3 Mon Sep 17 00:00:00 2001 From: Diksha Date: Thu, 30 Jul 2026 19:02:28 +0530 Subject: [PATCH 11/17] Fix lint errors --- src/datacustomcode/named_credential/base.py | 2 +- src/datacustomcode/named_credential/default.py | 2 +- src/datacustomcode/named_credential/direct/__init__.py | 2 +- src/datacustomcode/named_credential/direct/auth.py | 6 +++++- src/datacustomcode/named_credential/direct/transport.py | 8 +++----- .../named_credential/direct/url_resolver.py | 8 +++++--- src/datacustomcode/named_credential/spark_base.py | 4 ++-- src/datacustomcode/named_credential/spark_default.py | 4 +--- tests/test_named_credential_direct.py | 7 +++---- 9 files changed, 22 insertions(+), 21 deletions(-) diff --git a/src/datacustomcode/named_credential/base.py b/src/datacustomcode/named_credential/base.py index 246ae9f..ab7e8e9 100644 --- a/src/datacustomcode/named_credential/base.py +++ b/src/datacustomcode/named_credential/base.py @@ -63,7 +63,7 @@ def callout_json( ) -> Dict[str, Any]: """Low-level string-in/string-out callout returning the raw response. - It is required only by the per-row Spark path which forwards the + It is required only by the per-row Spark path which forwards the raw body per row; the default signals it as unsupported. """ raise NotImplementedError( diff --git a/src/datacustomcode/named_credential/default.py b/src/datacustomcode/named_credential/default.py index d3e06f0..da05fd1 100644 --- a/src/datacustomcode/named_credential/default.py +++ b/src/datacustomcode/named_credential/default.py @@ -32,7 +32,7 @@ class DefaultNamedCredential(NamedCredential): """ Executes the callout directly via :class:`DirectCalloutTransport`, resolving the URL from the Named Credential Connect API (falling back to - ``credential.json``) and injecting auth from ``credential.json``. + ``credential.json``) and injecting auth from ``credential.json``. """ CONFIG_NAME = "DefaultNamedCredential" diff --git a/src/datacustomcode/named_credential/direct/__init__.py b/src/datacustomcode/named_credential/direct/__init__.py index b8df664..9dc5930 100644 --- a/src/datacustomcode/named_credential/direct/__init__.py +++ b/src/datacustomcode/named_credential/direct/__init__.py @@ -14,6 +14,6 @@ # limitations under the License. """Named Credential callout path. -This resolves the endpoint via the NamedCredential Connect API +This resolves the endpoint via the NamedCredential Connect API and injects auth from a developer-provided ``credential.json`` """ diff --git a/src/datacustomcode/named_credential/direct/auth.py b/src/datacustomcode/named_credential/direct/auth.py index 4893ead..37eb058 100644 --- a/src/datacustomcode/named_credential/direct/auth.py +++ b/src/datacustomcode/named_credential/direct/auth.py @@ -19,7 +19,11 @@ from __future__ import annotations import base64 -from typing import TYPE_CHECKING, Any, Dict +from typing import ( + TYPE_CHECKING, + Any, + Dict, +) from requests.auth import AuthBase diff --git a/src/datacustomcode/named_credential/direct/transport.py b/src/datacustomcode/named_credential/direct/transport.py index c03ffc7..4217d53 100644 --- a/src/datacustomcode/named_credential/direct/transport.py +++ b/src/datacustomcode/named_credential/direct/transport.py @@ -55,7 +55,7 @@ def __init__( credentials_profile, sf_cli_org ) # Resolved base URL per callout key. Stable for the transport's life, so - # cache it to avoid a token fetch + Connect API call on every row + # cache it to avoid a token fetch + Connect API call on every row self._base_url_cache: Dict[str, str] = {} @staticmethod @@ -73,7 +73,7 @@ def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: f"Callout URL must start with 'callout:', got '{raw_url}'." ) - # Split the named credential reference at the first '/' or '?'; the remainder + # Split the named credential reference at the first '/' or '?'; the remainder # path or query string is appended to the resolved base URL verbatim. sep_idx = min( (i for i in (raw_url.find("/"), raw_url.find("?")) if i != -1), @@ -83,9 +83,7 @@ def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: path_suffix = raw_url[sep_idx:] if callout_key == "callout:": - raise CredentialError( - f"Named Credential name is empty in URL '{raw_url}'." - ) + raise CredentialError(f"Named Credential name is empty in URL '{raw_url}'.") cred_config = self._store.get(callout_key) base_url = self._base_url_cache.get(callout_key) diff --git a/src/datacustomcode/named_credential/direct/url_resolver.py b/src/datacustomcode/named_credential/direct/url_resolver.py index a22f889..991038b 100644 --- a/src/datacustomcode/named_credential/direct/url_resolver.py +++ b/src/datacustomcode/named_credential/direct/url_resolver.py @@ -38,7 +38,9 @@ logger = logging.getLogger(__name__) -NAMED_CREDENTIAL_PATH = "services/data/v63.0/named-credentials/named-credential-setup/{name}" +NAMED_CREDENTIAL_PATH = ( + "services/data/v63.0/named-credentials/named-credential-setup/{name}" +) def _callout_url_from_connect_api( @@ -65,8 +67,8 @@ def _callout_url_from_connect_api( return callout_url or None except Exception as exc: # A token was obtained but the Connect API call failed — likely a - # misconfiguration (named credential not onboarded, missing permission, wrong org). - # Surface it before silently falling back to target_url. + # misconfiguration (named credential not onboarded, missing + # permission, wrong org). Surface it before falling back to target_url. logger.warning( "Connect API URL resolution failed for %s: %s. " "Falling back to 'target_url' from credential.json if set.", diff --git a/src/datacustomcode/named_credential/spark_base.py b/src/datacustomcode/named_credential/spark_base.py index 71ac800..4037b02 100644 --- a/src/datacustomcode/named_credential/spark_base.py +++ b/src/datacustomcode/named_credential/spark_base.py @@ -53,7 +53,7 @@ def request( """Make an external callout through a Named Credential. Args: - request: The callout request + request: The callout request body: Optional JSON-serializable request body. Returns: @@ -75,7 +75,7 @@ def request_col( dispatched independently rather than one-shot on the driver. Args: - request: The callout template + request: The callout template body: Optional per-row ``Column`` holding the JSON request body as a string (or null for no body). diff --git a/src/datacustomcode/named_credential/spark_default.py b/src/datacustomcode/named_credential/spark_default.py index 7d8724e..230c103 100644 --- a/src/datacustomcode/named_credential/spark_default.py +++ b/src/datacustomcode/named_credential/spark_default.py @@ -127,9 +127,7 @@ def request_col( ) def _callout(body_str: Optional[str]) -> Dict[str, Any]: - return _invoke_callout_as_struct( - self._named_credential, request, body_str - ) + return _invoke_callout_as_struct(self._named_credential, request, body_str) body_col = body if body is not None else lit(None).cast(StringType()) return udf(_callout, result_schema)(body_col) diff --git a/tests/test_named_credential_direct.py b/tests/test_named_credential_direct.py index 045afb5..28af85a 100644 --- a/tests/test_named_credential_direct.py +++ b/tests/test_named_credential_direct.py @@ -16,6 +16,7 @@ import base64 import json +from typing import ClassVar import pytest from requests.models import PreparedRequest @@ -282,12 +283,10 @@ def counting_resolve(key, config, provider): class _Resp: status_code = 200 - headers: dict = {} + headers: ClassVar[dict] = {} text = "" - monkeypatch.setattr( - transport_mod.requests, "request", lambda **kwargs: _Resp() - ) + monkeypatch.setattr(transport_mod.requests, "request", lambda **kwargs: _Resp()) # Many per-row callouts, same callout key: base URL is resolved just once. for _ in range(5): From 756fc432992eda6aa4690432a674e061c11e16c4 Mon Sep 17 00:00:00 2001 From: Diksha Date: Mon, 3 Aug 2026 10:31:30 +0530 Subject: [PATCH 12/17] Address review comments --- src/datacustomcode/deploy.py | 7 +- src/datacustomcode/named_credential/base.py | 28 +--- .../named_credential/default.py | 53 ++---- .../named_credential/direct/__init__.py | 2 +- .../named_credential/direct/credentials.py | 34 ++-- .../named_credential/direct/transport.py | 9 +- .../named_credential/direct/url_resolver.py | 7 +- .../named_credential/spark_base.py | 10 +- .../named_credential/spark_default.py | 25 +-- .../named_credential/types/http_response.py | 15 +- tests/test_named_credential.py | 153 ++++++------------ tests/test_named_credential_direct.py | 36 +++-- tests/test_runtime_named_credential.py | 4 +- 13 files changed, 140 insertions(+), 243 deletions(-) diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index c65d43c..391e7b9 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -37,6 +37,9 @@ from datacustomcode.cmd import cmd_output from datacustomcode.constants import REQUEST_TYPE_TO_FEATURE +from datacustomcode.named_credential.direct.credentials import ( + EXTERNAL_CALLOUT_CREDENTIAL, +) from datacustomcode.scan import find_base_directory, get_package_type DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code" @@ -236,6 +239,8 @@ def create_deployment( ) PY_FILES_PATH = os.path.join("payload", "py-files") ZIP_FILE_NAME = "deployment.zip" +# Local-only files that must never be packaged into the deployment zip. +EXCLUDED_FILES = (".DS_Store", EXTERNAL_CALLOUT_CREDENTIAL) def prepare_dependency_archive( @@ -600,7 +605,7 @@ def zip( for root, dirs, files in os.walk(directory): # Skip .DS_Store and local credentials. for file in files: - if file not in (".DS_Store", "credential.json"): + if file not in EXCLUDED_FILES: abs_path = os.path.join(root, file) arcname = os.path.relpath(abs_path, directory) zipf.write(abs_path, arcname) diff --git a/src/datacustomcode/named_credential/base.py b/src/datacustomcode/named_credential/base.py index ab7e8e9..1478234 100644 --- a/src/datacustomcode/named_credential/base.py +++ b/src/datacustomcode/named_credential/base.py @@ -15,12 +15,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Optional, -) +from typing import TYPE_CHECKING, Optional from datacustomcode.mixin import UserExtendableNamedConfigMixin @@ -39,7 +34,7 @@ def __init__(self, **kwargs): def request( self, request: HTTPRequest, - body: Optional[Dict[str, Any]] = None, + body: Optional[str] = None, ) -> HTTPResponse: """Make an external callout through a Named Credential. @@ -49,24 +44,11 @@ def request( Args: request: The callout request - body: Optional JSON-serializable request body. + body: Optional request body. Set the + ``Content-Type`` header to match the format; the SDK + does not assume or inject one. Returns: The external service's response. """ ... - - def callout_json( - self, - request: HTTPRequest, - body: Optional[str] = None, - ) -> Dict[str, Any]: - """Low-level string-in/string-out callout returning the raw response. - - It is required only by the per-row Spark path which forwards the - raw body per row; the default signals it as unsupported. - """ - raise NotImplementedError( - f"{type(self).__name__} does not implement callout_json(); it " - "supports only one-shot request()." - ) diff --git a/src/datacustomcode/named_credential/default.py b/src/datacustomcode/named_credential/default.py index da05fd1..c0e72c2 100644 --- a/src/datacustomcode/named_credential/default.py +++ b/src/datacustomcode/named_credential/default.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from typing import ( Any, Dict, @@ -32,7 +31,7 @@ class DefaultNamedCredential(NamedCredential): """ Executes the callout directly via :class:`DirectCalloutTransport`, resolving the URL from the Named Credential Connect API (falling back to - ``credential.json``) and injecting auth from ``credential.json``. + ``external_callout_config.json``) and injecting auth from the same file. """ CONFIG_NAME = "DefaultNamedCredential" @@ -51,42 +50,12 @@ def __init__( self._transport: Optional[Any] = None def request( - self, - request: HTTPRequest, - body: Optional[Dict[str, Any]] = None, - ) -> HTTPResponse: - callout_response = self.callout_json( - request, json.dumps(body) if body is not None else "" - ) - - raw_body = callout_response.get("body") or "" - data: Optional[Any] = None - if raw_body: - try: - # Preserve any JSON value: object, array, or scalar. - data = json.loads(raw_body) - except json.JSONDecodeError: - data = None - - response_dict = { - "status_code": callout_response.get("http_status_code"), - "headers": callout_response.get("headers", {}), - "data": data, - } - return HTTPResponseBuilder.build(response_dict) - - def callout_json( self, request: HTTPRequest, body: Optional[str] = None, - ) -> Dict[str, Any]: - """Raw string-in/string-out callout returning the unparsed response. - - Unlike :meth:`request`, the ``body`` is sent verbatim (never re-parsed) - and the response is returned as a raw ``{http_status_code, headers, - body}`` dict rather than a parsed :class:`HTTPResponse`. - """ - # Callout request shape sent to the transport. + ) -> HTTPResponse: + # Body and response are treated as opaque strings; the SDK makes no + # assumption about their format (JSON, XML, text, ...). callout_request = { "path": request.url, "method": request.method, @@ -94,16 +63,18 @@ def callout_json( "body": body if body is not None else "", } callout_response = self._callout(callout_request) - return { - "http_status_code": callout_response.get("http_status_code"), - "headers": callout_response.get("headers", {}), - "body": callout_response.get("body") or "", - } + return HTTPResponseBuilder.build( + { + "status_code": callout_response.get("status_code"), + "headers": callout_response.get("headers", {}), + "body": callout_response.get("body") or "", + } + ) def _callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: """Execute the callout via the transport. - Returns a dict with ``http_status_code``, ``headers``, and ``body``. + Returns a dict with ``status_code``, ``headers``, and ``body``. """ result: Dict[str, Any] = self._get_transport().callout(callout_request) return result diff --git a/src/datacustomcode/named_credential/direct/__init__.py b/src/datacustomcode/named_credential/direct/__init__.py index 9dc5930..74bd837 100644 --- a/src/datacustomcode/named_credential/direct/__init__.py +++ b/src/datacustomcode/named_credential/direct/__init__.py @@ -15,5 +15,5 @@ """Named Credential callout path. This resolves the endpoint via the NamedCredential Connect API -and injects auth from a developer-provided ``credential.json`` +and injects auth from a developer-provided ``external_callout_config.json`` """ diff --git a/src/datacustomcode/named_credential/direct/credentials.py b/src/datacustomcode/named_credential/direct/credentials.py index 7c9ce66..8b60864 100644 --- a/src/datacustomcode/named_credential/direct/credentials.py +++ b/src/datacustomcode/named_credential/direct/credentials.py @@ -12,11 +12,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Load developer credentials for direct callouts from ``credential.json``. +"""Load developer credentials for direct callouts from ``external_callout_config.json``. The file lives in the parent of the payload folder so it is never packaged in -the deployment zip. Keys are callout references (``callout:``); -each value carries a mandatory ``auth_type`` and an optional ``target_url``. +the deployment zip. Its ``credentials`` section is keyed by callout reference +(``callout:``); each value carries a mandatory ``auth_type`` +and an optional ``target_url``. """ from __future__ import annotations @@ -32,9 +33,9 @@ ) # Default file name; discovered in the parent of the payload folder. -DEFAULT_CREDENTIAL_FILE = "credential.json" +EXTERNAL_CALLOUT_CREDENTIAL = "external_callout_config.json" # Absolute-path override, primarily for tests and non-standard layouts. -CREDENTIAL_FILE_ENV_VAR = "DATACUSTOMCODE_CREDENTIAL_FILE" +CREDENTIAL_FILE_ENV_VAR = "DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG" class AuthType(str, Enum): @@ -51,20 +52,20 @@ class CredentialError(RuntimeError): def _discover_credential_file() -> Optional[Path]: - """Find ``credential.json`` via env override, then by walking up from cwd.""" + """Find the config file via env override, then by walking up from cwd.""" override = os.environ.get(CREDENTIAL_FILE_ENV_VAR) if override: return Path(override) for directory in (Path.cwd(), *Path.cwd().parents): - candidate = directory / DEFAULT_CREDENTIAL_FILE + candidate = directory / EXTERNAL_CALLOUT_CREDENTIAL if candidate.is_file(): return candidate return None class CredentialStore: - """Reads ``credential.json`` and returns per-callout configuration.""" + """Reads ``external_callout_config.json`` and returns per-callout config.""" def __init__(self, credential_file: Optional[str] = None) -> None: self._explicit_path = Path(credential_file) if credential_file else None @@ -77,7 +78,7 @@ def _load(self) -> Dict[str, Dict[str, Any]]: path = self._explicit_path or _discover_credential_file() if path is None or not path.is_file(): raise CredentialError( - f"Could not find '{DEFAULT_CREDENTIAL_FILE}'. Place it in the " + f"Could not find '{EXTERNAL_CALLOUT_CREDENTIAL}'. Place it in the " f"parent of your payload folder, or set " f"${CREDENTIAL_FILE_ENV_VAR} to its path." ) @@ -87,12 +88,16 @@ def _load(self) -> Dict[str, Dict[str, Any]]: except (OSError, json.JSONDecodeError) as exc: raise CredentialError(f"Failed to read '{path}': {exc}") from exc - if not isinstance(data, dict): + # Per-callout entries live under the ``credentials`` section, leaving + # room for other config sections alongside them in the future. + credentials = data.get("credentials") if isinstance(data, dict) else None + if not isinstance(credentials, dict): raise CredentialError( - f"'{path}' must be a JSON object keyed by callout reference." + f"'{path}' must be a JSON object with a 'credentials' section " + f"keyed by callout reference." ) - self._credentials = data - return data + self._credentials = credentials + return credentials def get(self, callout_key: str) -> Dict[str, Any]: """Return the config for a callout key (e.g. ``callout:AWS_S3_Service``). @@ -105,7 +110,8 @@ def get(self, callout_key: str) -> Dict[str, Any]: if config is None: raise CredentialError( f"No credential configuration found for '{callout_key}'. " - f"Add it to '{DEFAULT_CREDENTIAL_FILE}'." + f"Add it to the 'credentials' section of " + f"'{EXTERNAL_CALLOUT_CREDENTIAL}'." ) if not isinstance(config, dict) or not config.get("auth_type"): raise CredentialError( diff --git a/src/datacustomcode/named_credential/direct/transport.py b/src/datacustomcode/named_credential/direct/transport.py index 4217d53..38dda2d 100644 --- a/src/datacustomcode/named_credential/direct/transport.py +++ b/src/datacustomcode/named_credential/direct/transport.py @@ -16,8 +16,8 @@ Resolves the ``callout:/`` reference to a real endpoint (via the Named Credential Connect API, falling back to ``target_url`` in -``credential.json``), attaches the credential's auth, and returns the raw -``{http_status_code, headers, body}`` response. +``external_callout_config.json``), attaches the credential's auth, and returns the raw +``{status_code, headers, body}`` response. """ from __future__ import annotations @@ -92,9 +92,8 @@ def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: self._base_url_cache[callout_key] = base_url body = callout_request.get("body") or None + # Headers are passed; the SDK assumes no Content-Type. headers = dict(callout_request.get("headers", {})) - if body and "Content-Type" not in headers: - headers["Content-Type"] = "application/json" response = requests.request( method=callout_request["method"], @@ -105,7 +104,7 @@ def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: timeout=30, ) return { - "http_status_code": response.status_code, + "status_code": response.status_code, "headers": dict(response.headers), "body": response.text, } diff --git a/src/datacustomcode/named_credential/direct/url_resolver.py b/src/datacustomcode/named_credential/direct/url_resolver.py index 991038b..6044ac2 100644 --- a/src/datacustomcode/named_credential/direct/url_resolver.py +++ b/src/datacustomcode/named_credential/direct/url_resolver.py @@ -71,7 +71,7 @@ def _callout_url_from_connect_api( # permission, wrong org). Surface it before falling back to target_url. logger.warning( "Connect API URL resolution failed for %s: %s. " - "Falling back to 'target_url' from credential.json if set.", + "Falling back to 'target_url' from external_callout_config.json if set.", developer_name, exc, ) @@ -87,7 +87,7 @@ def resolve_base_url( Args: callout_key: e.g. ``callout:Nominatim_Geocoding``. - cred_config: The callout's ``credential.json`` entry. + cred_config: The callout's ``external_callout_config.json`` entry. token_provider: Provides a token/instance URL for the Connect API; when None, only ``target_url`` is used. @@ -106,6 +106,7 @@ def resolve_base_url( if not base_url: raise CredentialError( f"Could not resolve a URL for '{callout_key}'. Ensure the Named " - f"Credential exists (sf login) or set 'target_url' in credential.json." + f"Credential exists (sf login) or set 'target_url' in " + f"external_callout_config.json." ) return base_url.rstrip("/") diff --git a/src/datacustomcode/named_credential/spark_base.py b/src/datacustomcode/named_credential/spark_base.py index 4037b02..6cc12f5 100644 --- a/src/datacustomcode/named_credential/spark_base.py +++ b/src/datacustomcode/named_credential/spark_base.py @@ -18,7 +18,6 @@ from typing import ( TYPE_CHECKING, Any, - Dict, Optional, ) @@ -48,13 +47,14 @@ def __init__(self, **kwargs: Any) -> None: def request( self, request: HTTPRequest, - body: Optional[Dict[str, Any]] = None, + body: Optional[str] = None, ) -> HTTPResponse: """Make an external callout through a Named Credential. Args: request: The callout request - body: Optional JSON-serializable request body. + body: Optional request body. Set the + ``Content-Type`` header yourself; the SDK does not assume one. Returns: The external service's response. @@ -76,8 +76,8 @@ def request_col( Args: request: The callout template - body: Optional per-row ``Column`` holding the JSON request body as a - string (or null for no body). + body: Optional per-row ``Column`` holding the request body as a + string, sent verbatim (or null for no body). Returns: A ``Column`` yielding a struct diff --git a/src/datacustomcode/named_credential/spark_default.py b/src/datacustomcode/named_credential/spark_default.py index 230c103..8f3ea23 100644 --- a/src/datacustomcode/named_credential/spark_default.py +++ b/src/datacustomcode/named_credential/spark_default.py @@ -21,12 +21,12 @@ Optional, ) -from datacustomcode.named_credential.base import NamedCredential from datacustomcode.named_credential.spark_base import SparkNamedCredential if TYPE_CHECKING: from pyspark.sql import Column + from datacustomcode.named_credential.base import NamedCredential from datacustomcode.named_credential.types.http_request import HTTPRequest from datacustomcode.named_credential.types.http_response import HTTPResponse @@ -71,7 +71,7 @@ def __init__( def request( self, request: "HTTPRequest", - body: Optional[Dict[str, Any]] = None, + body: Optional[str] = None, ) -> "HTTPResponse": return self._named_credential.request(request, body) @@ -114,18 +114,6 @@ def request_col( ] ) - # Fail at column-build time rather than turning every row into an opaque error. - if ( - getattr(type(self._named_credential), "callout_json", None) - is NamedCredential.callout_json - ): - raise TypeError( - f"{type(self._named_credential).__name__} does not support the " - "per-row callout path; it must override callout_json(). Use " - "named_credential_request() for a one-shot callout, or configure " - "DefaultNamedCredential." - ) - def _callout(body_str: Optional[str]) -> Dict[str, Any]: return _invoke_callout_as_struct(self._named_credential, request, body_str) @@ -145,7 +133,7 @@ def _invoke_callout_as_struct( ERROR structs rather than aborting the job. """ try: - callout_response = named_credential.callout_json(request, body_str) + response = named_credential.request(request, body_str) except Exception as exc: # surface any transport error per row return { "status": _STATUS_ERROR, @@ -154,13 +142,12 @@ def _invoke_callout_as_struct( "error_message": str(exc), } - status_code = callout_response.get("http_status_code") return { "status": _STATUS_SUCCESS, "response": { - "status_code": int(status_code) if status_code is not None else None, - "body": callout_response.get("body") or "", - "headers": callout_response.get("headers") or {}, + "status_code": response.status_code, + "body": response.body, + "headers": response.headers, }, "error_code": None, "error_message": None, diff --git a/src/datacustomcode/named_credential/types/http_response.py b/src/datacustomcode/named_credential/types/http_response.py index 6e24643..3f98ae1 100644 --- a/src/datacustomcode/named_credential/types/http_response.py +++ b/src/datacustomcode/named_credential/types/http_response.py @@ -13,11 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import ( - Any, - Dict, - Optional, -) +from typing import Dict from pydantic import BaseModel, Field @@ -29,10 +25,11 @@ class HTTPResponse(BaseModel): headers: Dict[str, str] = Field( default_factory=dict, description="Response headers" ) - data: Optional[Any] = Field( - default=None, - description="Parsed JSON response body (object, array, or scalar), " - "or None if the body was empty or not JSON", + body: str = Field( + default="", + description="Raw response body, verbatim (any format). The SDK does not " + "parse it; the caller decodes as needed (e.g. json.loads). Empty string " + "if the response had no body.", ) @property diff --git a/tests/test_named_credential.py b/tests/test_named_credential.py index 51bd43d..60470a2 100644 --- a/tests/test_named_credential.py +++ b/tests/test_named_credential.py @@ -116,7 +116,7 @@ class TestHTTPResponse: def test_response_defaults(self): response = HTTPResponse(status_code=200) assert response.headers == {} - assert response.data is None + assert response.body == "" def test_is_success_for_2xx(self): assert HTTPResponse(status_code=200).is_success is True @@ -136,12 +136,12 @@ def test_status_code_validation(self): class TestHTTPResponseBuilder: def test_build_from_dict(self): response = HTTPResponseBuilder.build( - {"status_code": 200, "headers": {"X": "y"}, "data": {"ok": True}} + {"status_code": 200, "headers": {"X": "y"}, "body": '{"ok": true}'} ) assert isinstance(response, HTTPResponse) assert response.status_code == 200 assert response.headers == {"X": "y"} - assert response.data == {"ok": True} + assert response.body == '{"ok": true}' def test_build_requires_status_code(self): with pytest.raises(ValidationError): @@ -154,20 +154,20 @@ def test_callout_delegates_to_transport(self, monkeypatch): class _FakeTransport: def callout(self, callout_request): - return {"http_status_code": 200, "headers": {}, "body": "{}"} + return {"status_code": 200, "headers": {}, "body": "{}"} monkeypatch.setattr(nc, "_get_transport", lambda: _FakeTransport()) result = nc._callout({"path": "callout:NC/x", "method": "GET"}) - assert result["http_status_code"] == 200 + assert result["status_code"] == 200 - def test_request_translates_json_response(self, monkeypatch): + def test_request_forwards_raw_body_and_response(self, monkeypatch): nc = DefaultNamedCredential() captured = {} def fake_callout(callout_request): captured.update(callout_request) return { - "http_status_code": 200, + "status_code": 200, "headers": {"Content-Type": "application/json"}, "body": '{"result": "ok"}', } @@ -180,7 +180,7 @@ def fake_callout(callout_request): .set_headers({"Accept": "application/json"}) .build() ) - response = nc.request(request, {"key": "value"}) + response = nc.request(request, '{"key": "value"}') # The query string travels inside the path verbatim. assert captured["path"] == "callout:NC/search?q=sf" @@ -191,7 +191,7 @@ def fake_callout(callout_request): assert response.status_code == 200 assert response.headers == {"Content-Type": "application/json"} - assert response.data == {"result": "ok"} + assert response.body == '{"result": "ok"}' assert response.is_success is True def test_request_without_body_sends_empty_string(self, monkeypatch): @@ -200,7 +200,7 @@ def test_request_without_body_sends_empty_string(self, monkeypatch): def fake_callout(callout_request): captured.update(callout_request) - return {"http_status_code": 204, "headers": {}, "body": ""} + return {"status_code": 204, "headers": {}, "body": ""} monkeypatch.setattr(nc, "_callout", fake_callout) request = HTTPRequestBuilder().set_url("callout:NC/path").build() @@ -208,102 +208,43 @@ def fake_callout(callout_request): assert captured["body"] == "" assert response.status_code == 204 - assert response.data is None + assert response.body == "" - def test_request_non_json_body_yields_none_data(self, monkeypatch): + def test_request_non_json_body_returned_verbatim(self, monkeypatch): nc = DefaultNamedCredential() def fake_callout(callout_request): - return {"http_status_code": 200, "headers": {}, "body": "plain text"} + return {"status_code": 200, "headers": {}, "body": "plain text"} monkeypatch.setattr(nc, "_callout", fake_callout) request = HTTPRequestBuilder().set_url("callout:NC/path").build() response = nc.request(request) - assert response.data is None + assert response.body == "plain text" - def test_request_json_array_body_preserved(self, monkeypatch): + def test_request_json_array_body_returned_verbatim(self, monkeypatch): nc = DefaultNamedCredential() def fake_callout(callout_request): - return {"http_status_code": 200, "headers": {}, "body": "[1, 2, 3]"} + return {"status_code": 200, "headers": {}, "body": "[1, 2, 3]"} monkeypatch.setattr(nc, "_callout", fake_callout) request = HTTPRequestBuilder().set_url("callout:NC/path").build() response = nc.request(request) - assert response.data == [1, 2, 3] + assert response.body == "[1, 2, 3]" - def test_request_json_scalar_body_preserved(self, monkeypatch): - nc = DefaultNamedCredential() - - def fake_callout(callout_request): - return {"http_status_code": 200, "headers": {}, "body": "42"} - - monkeypatch.setattr(nc, "_callout", fake_callout) - request = HTTPRequestBuilder().set_url("callout:NC/path").build() - response = nc.request(request) - assert response.data == 42 - - def test_request_empty_dict_body_sends_object(self, monkeypatch): + def test_request_opaque_string_body_sent_verbatim(self, monkeypatch): nc = DefaultNamedCredential() captured = {} def fake_callout(callout_request): captured.update(callout_request) - return {"http_status_code": 200, "headers": {}, "body": ""} + return {"status_code": 200, "headers": {}, "body": ""} monkeypatch.setattr(nc, "_callout", fake_callout) request = HTTPRequestBuilder().set_url("callout:NC/path").build() - nc.request(request, {}) - - # An explicit empty dict is a body ("{}"), distinct from None (""). - assert captured["body"] == "{}" - - def test_callout_json_forwards_body_and_returns_raw_shape(self, monkeypatch): - nc = DefaultNamedCredential() - captured = {} - - def fake_callout(callout_request): - captured.update(callout_request) - return { - "http_status_code": 200, - "headers": {"X": "y"}, - "body": "plain text", - } + nc.request(request, "a=1&b=2") - monkeypatch.setattr(nc, "_callout", fake_callout) - request = ( - HTTPRequestBuilder() - .set_url("callout:NC/path") - .set_method(HTTPMethod.POST) - .build() - ) - - result = nc.callout_json(request, "a=1&b=2") - - # Body is sent verbatim, not re-encoded as JSON. assert captured["body"] == "a=1&b=2" - # Response is the raw {http_status_code, headers, body} shape, body intact. - assert result == { - "http_status_code": 200, - "headers": {"X": "y"}, - "body": "plain text", - } - - def test_callout_json_none_body_sends_empty_string(self, monkeypatch): - nc = DefaultNamedCredential() - captured = {} - - def fake_callout(callout_request): - captured.update(callout_request) - return {"http_status_code": 204, "headers": {}, "body": ""} - - monkeypatch.setattr(nc, "_callout", fake_callout) - request = HTTPRequestBuilder().set_url("callout:NC/path").build() - - result = nc.callout_json(request) - - assert captured["body"] == "" - assert result == {"http_status_code": 204, "headers": {}, "body": ""} class TestDefaultSparkNamedCredential: @@ -326,10 +267,10 @@ def request(self, request, body=None): spark_nc = DefaultSparkNamedCredential(named_credential=underlying) request = HTTPRequestBuilder().set_url("callout:NC/path").build() - result = spark_nc.request(request, {"k": "v"}) + result = spark_nc.request(request, '{"k": "v"}') assert result is sentinel - assert underlying.calls == [(request, {"k": "v"})] + assert underlying.calls == [(request, '{"k": "v"}')] def test_builds_underlying_from_config_when_absent(self, monkeypatch): from datacustomcode.named_credential import spark_default @@ -369,11 +310,11 @@ def test_wraps_callout_in_udf_over_body_column(self, mock_lit, mock_udf): mock_udf.return_value = sentinel_udf underlying = MagicMock() - underlying.callout_json.return_value = { - "http_status_code": 200, - "headers": {"Content-Type": "application/json"}, - "body": '{"ok":true}', - } + underlying.request.return_value = HTTPResponse( + status_code=200, + headers={"Content-Type": "application/json"}, + body='{"ok":true}', + ) spark_nc = DefaultSparkNamedCredential(named_credential=underlying) request = HTTPRequestBuilder().set_url("callout:NC/v1/accounts").build() @@ -401,7 +342,7 @@ def test_wraps_callout_in_udf_over_body_column(self, mock_lit, mock_udf): assert payload["body"] == '{"ok":true}' # The raw body string (NOT a parsed dict) is forwarded to the callout. - sent_request, sent_body = underlying.callout_json.call_args.args + sent_request, sent_body = underlying.request.call_args.args assert sent_request is request assert sent_body == '{"name": "acme"}' @@ -437,11 +378,11 @@ def test_success_struct_carries_response_fields(self): ) underlying = MagicMock() - underlying.callout_json.return_value = { - "http_status_code": 201, - "headers": {"X-Trace": "abc"}, - "body": "[1,2,3]", - } + underlying.request.return_value = HTTPResponse( + status_code=201, + headers={"X-Trace": "abc"}, + body="[1,2,3]", + ) request = HTTPRequestBuilder().set_url("callout:NC/path").build() out = _invoke_callout_as_struct(underlying, request, '{"a": 1}') @@ -454,7 +395,7 @@ def test_success_struct_carries_response_fields(self): "headers": {"X-Trace": "abc"}, } # The request body is forwarded verbatim (never validated as JSON). - assert underlying.callout_json.call_args.args[1] == '{"a": 1}' + assert underlying.request.call_args.args[1] == '{"a": 1}' def test_non_json_request_body_is_forwarded_not_rejected(self): from datacustomcode.named_credential.spark_default import ( @@ -464,17 +405,17 @@ def test_non_json_request_body_is_forwarded_not_rejected(self): # Mirrors the runtime: an opaque body (form-encoded / plain text) is sent # as-is instead of being rejected the way JSON parsing would. underlying = MagicMock() - underlying.callout_json.return_value = { - "http_status_code": 200, - "headers": {}, - "body": "OK", - } + underlying.request.return_value = HTTPResponse( + status_code=200, + headers={}, + body="OK", + ) request = HTTPRequestBuilder().set_url("callout:NC/path").build() out = _invoke_callout_as_struct(underlying, request, "a=1&b=2") assert out["status"] == "SUCCESS" - assert underlying.callout_json.call_args.args[1] == "a=1&b=2" + assert underlying.request.call_args.args[1] == "a=1&b=2" # A non-JSON response body is preserved verbatim, not dropped to "". assert out["response"]["body"] == "OK" @@ -484,11 +425,11 @@ def test_none_body_forwards_none_and_keeps_all_keys(self): ) underlying = MagicMock() - underlying.callout_json.return_value = { - "http_status_code": 200, - "headers": {}, - "body": "", - } + underlying.request.return_value = HTTPResponse( + status_code=200, + headers={}, + body="", + ) request = HTTPRequestBuilder().set_url("callout:NC/path").build() out = _invoke_callout_as_struct(underlying, request, None) @@ -502,7 +443,7 @@ def test_none_body_forwards_none_and_keeps_all_keys(self): "headers": {}, } # A null body column forwards None (not "null") to the callout. - assert underlying.callout_json.call_args.args[1] is None + assert underlying.request.call_args.args[1] is None def test_transport_error_yields_error_struct(self): from datacustomcode.named_credential.spark_default import ( @@ -510,7 +451,7 @@ def test_transport_error_yields_error_struct(self): ) underlying = MagicMock() - underlying.callout_json.side_effect = RuntimeError("proxy down") + underlying.request.side_effect = RuntimeError("proxy down") request = HTTPRequestBuilder().set_url("callout:NC/path").build() out = _invoke_callout_as_struct(underlying, request, '{"a": 1}') diff --git a/tests/test_named_credential_direct.py b/tests/test_named_credential_direct.py index 28af85a..790f60f 100644 --- a/tests/test_named_credential_direct.py +++ b/tests/test_named_credential_direct.py @@ -78,9 +78,11 @@ def test_unsupported_auth_type_raises_value_error(self): class TestCredentialStore: def test_get_returns_config(self, tmp_path, monkeypatch): - cred_file = tmp_path / "credential.json" + cred_file = tmp_path / "external_callout_config.json" cred_file.write_text( - json.dumps({"callout:NC": {"auth_type": "Basic", "username": "u"}}) + json.dumps( + {"credentials": {"callout:NC": {"auth_type": "Basic", "username": "u"}}} + ) ) monkeypatch.setenv(CREDENTIAL_FILE_ENV_VAR, str(cred_file)) store = CredentialStore() @@ -89,8 +91,10 @@ def test_get_returns_config(self, tmp_path, monkeypatch): assert config["username"] == "u" def test_explicit_path_takes_precedence(self, tmp_path): - cred_file = tmp_path / "credential.json" - cred_file.write_text(json.dumps({"callout:NC": {"auth_type": "OAuth"}})) + cred_file = tmp_path / "external_callout_config.json" + cred_file.write_text( + json.dumps({"credentials": {"callout:NC": {"auth_type": "OAuth"}}}) + ) store = CredentialStore(str(cred_file)) assert store.get("callout:NC")["auth_type"] == "OAuth" @@ -100,21 +104,25 @@ def test_missing_file_raises(self, tmp_path): store.get("callout:NC") def test_missing_key_raises(self, tmp_path): - cred_file = tmp_path / "credential.json" - cred_file.write_text(json.dumps({"callout:Other": {"auth_type": "Basic"}})) + cred_file = tmp_path / "external_callout_config.json" + cred_file.write_text( + json.dumps({"credentials": {"callout:Other": {"auth_type": "Basic"}}}) + ) store = CredentialStore(str(cred_file)) with pytest.raises(CredentialError): store.get("callout:NC") def test_missing_auth_type_raises(self, tmp_path): - cred_file = tmp_path / "credential.json" - cred_file.write_text(json.dumps({"callout:NC": {"username": "u"}})) + cred_file = tmp_path / "external_callout_config.json" + cred_file.write_text( + json.dumps({"credentials": {"callout:NC": {"username": "u"}}}) + ) store = CredentialStore(str(cred_file)) with pytest.raises(CredentialError): store.get("callout:NC") def test_non_object_json_raises(self, tmp_path): - cred_file = tmp_path / "credential.json" + cred_file = tmp_path / "external_callout_config.json" cred_file.write_text(json.dumps(["not", "an", "object"])) store = CredentialStore(str(cred_file)) with pytest.raises(CredentialError): @@ -182,8 +190,8 @@ class TestDirectCalloutTransport: def _make_transport(self, tmp_path, monkeypatch, cred_entry): from datacustomcode.named_credential.direct import transport as transport_mod - cred_file = tmp_path / "credential.json" - cred_file.write_text(json.dumps({"callout:NC": cred_entry})) + cred_file = tmp_path / "external_callout_config.json" + cred_file.write_text(json.dumps({"credentials": {"callout:NC": cred_entry}})) monkeypatch.setattr( transport_mod.DirectCalloutTransport, "_build_token_provider", @@ -233,11 +241,11 @@ def fake_request(**kwargs): assert captured["method"] == "POST" assert "params" not in captured assert captured["data"] == '{"a": 1}' - # Content-Type auto-added because a body is present. - assert captured["headers"]["Content-Type"] == "application/json" + # No Content-Type is assumed; headers are passed through verbatim. + assert "Content-Type" not in captured["headers"] assert isinstance(captured["auth"], DynamicAuthHandler) - assert result["http_status_code"] == 200 + assert result["status_code"] == 200 assert result["body"] == '{"ok": true}' @pytest.mark.parametrize("method", ["PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"]) diff --git a/tests/test_runtime_named_credential.py b/tests/test_runtime_named_credential.py index d5b81c4..78e9f5a 100644 --- a/tests/test_runtime_named_credential.py +++ b/tests/test_runtime_named_credential.py @@ -42,7 +42,7 @@ def __init__(self, custom_param: str = "default", **kwargs): self.custom_param = custom_param def request(self, request, body=None): - return HTTPResponse(status_code=200, data={"echo": self.custom_param}) + return HTTPResponse(status_code=200, body=self.custom_param) assert "CustomNamedCredential" in NamedCredential.available_config_names() cls = NamedCredential.subclass_from_config_name("CustomNamedCredential") @@ -56,7 +56,7 @@ def request(self, request, body=None): assert isinstance(instance, CustomNamedCredential) response = instance.request(HTTPRequest(url="callout:NC/path")) - assert response.data == {"echo": "my_value"} + assert response.body == "my_value" class TestRuntimeNamedCredential: From 1f2096d708b8c9b7f000131fbf1db4687bb402aa Mon Sep 17 00:00:00 2001 From: Diksha Date: Mon, 3 Aug 2026 12:06:58 +0530 Subject: [PATCH 13/17] Fix build error and add example --- src/datacustomcode/client.py | 14 +- .../named_credential/__init__.py | 2 - .../named_credential/direct/auth.py | 7 +- src/datacustomcode/named_credential/errors.py | 36 ---- .../named_credential/spark_base.py | 6 +- .../chunking_with_external_callout/README.md | 119 +++++++++++++ .../entrypoint.py | 161 ++++++++++++++++++ .../external_callout_config.json | 11 ++ .../tests/test.json | 16 ++ tests/test_client.py | 60 +++++++ tests/test_deploy.py | 30 ++++ tests/test_named_credential_direct.py | 16 ++ 12 files changed, 430 insertions(+), 48 deletions(-) delete mode 100644 src/datacustomcode/named_credential/errors.py create mode 100644 src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md create mode 100644 src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py create mode 100644 src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json create mode 100644 src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 9776d42..2e306f6 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -201,14 +201,15 @@ def named_credential_request_col( The returned Column yields a struct ``{status, response, error_code, error_message}`` for each row. ``response`` is itself a struct ``{status_code, body, headers}``. Use ``[...]`` to pick a field, e.g. - ``named_credential_request_col(...)["response"]["status_code"]``. Per-row - failures populate ``status`` / ``error_code`` / ``error_message`` so a - single bad row does not abort the whole Spark job. + ``named_credential_request_col(...)["response"]["status_code"]``. A transport + failure sets ``status`` to ``ERROR`` and populates ``error_message`` (a non-2xx + HTTP response is still ``SUCCESS`` with its code in ``response.status_code``), + so a single bad row does not abort the whole Spark job. Args: request: The callout template — its symbolic reference, method, and headers are applied to every row. - body: Optional per-row ``Column`` holding the JSON request body as a + body: Optional per-row ``Column`` holding the request body as a string (or null for no body). Returns: @@ -533,7 +534,7 @@ def _get_spark_einstein_predictions(self) -> SparkEinsteinPredictions: def named_credential_request( self, request: "HTTPRequest", - body: Optional[Dict[str, Any]] = None, + body: Optional[str] = None, ) -> "HTTPResponse": """Issue a one-shot Named Credential external callout. This is the scalar counterpart to :func:`named_credential_request_col`: it runs @@ -552,7 +553,8 @@ def named_credential_request( Args: request: The callout request - body: Optional JSON-serializable request body. + body: Optional request body. Set the ``Content-Type`` header to + match the format; the SDK does not assume or inject one. Returns: The external service's response. diff --git a/src/datacustomcode/named_credential/__init__.py b/src/datacustomcode/named_credential/__init__.py index 31245db..0f97b7e 100644 --- a/src/datacustomcode/named_credential/__init__.py +++ b/src/datacustomcode/named_credential/__init__.py @@ -15,7 +15,6 @@ from datacustomcode.named_credential.base import NamedCredential from datacustomcode.named_credential.default import DefaultNamedCredential -from datacustomcode.named_credential.errors import NamedCredentialCallError from datacustomcode.named_credential.spark_base import SparkNamedCredential from datacustomcode.named_credential.spark_default import DefaultSparkNamedCredential @@ -23,6 +22,5 @@ "DefaultNamedCredential", "DefaultSparkNamedCredential", "NamedCredential", - "NamedCredentialCallError", "SparkNamedCredential", ] diff --git a/src/datacustomcode/named_credential/direct/auth.py b/src/datacustomcode/named_credential/direct/auth.py index 37eb058..24c59c4 100644 --- a/src/datacustomcode/named_credential/direct/auth.py +++ b/src/datacustomcode/named_credential/direct/auth.py @@ -51,8 +51,11 @@ def __call__(self, request: PreparedRequest) -> PreparedRequest: elif self.auth_type in (AuthType.OAUTH.value, AuthType.JWT.value): bearer = self.config.get("access_token") or self.config.get("token") - if bearer: - request.headers["Authorization"] = f"Bearer {bearer}" + if not bearer: + raise ValueError( + f"'{self.auth_type}' auth requires an 'access_token' or 'token'." + ) + request.headers["Authorization"] = f"Bearer {bearer}" else: raise ValueError(f"Unsupported auth_type '{self.auth_type}'.") diff --git a/src/datacustomcode/named_credential/errors.py b/src/datacustomcode/named_credential/errors.py deleted file mode 100644 index e156fe4..0000000 --- a/src/datacustomcode/named_credential/errors.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright (c) 2025, Salesforce, Inc. -# SPDX-License-Identifier: Apache-2 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Exceptions raised by Named Credential implementations.""" - -from __future__ import annotations - -from typing import Optional - - -class NamedCredentialCallError(RuntimeError): - """Raised when a Named Credential external callout fails.""" - - def __init__( - self, - message: str, - *, - status: Optional[object] = None, - error_code: Optional[str] = None, - error_message: Optional[str] = None, - ) -> None: - super().__init__(message) - self.status = status - self.error_code = error_code - self.error_message = error_message diff --git a/src/datacustomcode/named_credential/spark_base.py b/src/datacustomcode/named_credential/spark_base.py index 6cc12f5..2c47743 100644 --- a/src/datacustomcode/named_credential/spark_base.py +++ b/src/datacustomcode/named_credential/spark_base.py @@ -85,7 +85,9 @@ def request_col( itself a struct ``{status_code, body, headers}`` carrying the callout response. Select a field with ``[...]``, e.g. ``request_col(...)["response"]["status_code"]``. Returning a struct - means a single failing row does not abort the Spark job — callers can - inspect ``status`` / ``error_code`` per row instead. + means a single failing row does not abort the Spark job — a transport + failure sets ``status`` to ``ERROR`` with ``error_message``, while a + non-2xx HTTP response stays ``SUCCESS`` with its code in + ``response.status_code``. """ ... diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md b/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md new file mode 100644 index 0000000..2dc4633 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md @@ -0,0 +1,119 @@ +# Chunking with a Gemini Named Credential Callout + +Splits each input document into paragraph-sized chunks and calls Google's +**Gemini** `generateContent` API for every chunk. The model returns a summary, +category, sentiment, and topics, which are attached to the chunk as citations so +the search index can filter and rank on them. Gemini is reached through a +**Named Credential**, so this code never handles the endpoint URL or the API key. + +## How the callout works + +```python +CALLOUT_URL = "callout:gemini" # callout:[/] + +request = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json"}) + .build() +) +# Body is sent verbatim (serialize it yourself); the response body is a raw string. +response = runtime.named_credential.request(request, json.dumps(payload)) +if response.is_success: + envelope = json.loads(response.body) + text = envelope["candidates"][0]["content"]["parts"][0]["text"] +``` + +The request asks for `responseMimeType: application/json` with a `responseSchema`, +so Gemini returns the classification as a JSON string in +`candidates[0].content.parts[0].text` — decode it, then decode that text again. + +The `gemini` Named Credential's URL already includes the full +`/v1beta/models/:generateContent` path, so the callout is just +`callout:gemini` with **no path suffix** (anything after the name is appended to +the credential's URL). + +## Configure the Named Credential + +1. Create an **External Credential** (e.g. `google_api_key`) that injects your + Gemini API key as the `X-goog-api-key` header. +2. Create a **Named Credential** named `gemini`: + - **URL**: `https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent` + - **Enabled for Callouts** + **Generate Authorization Header**: on + - **External Credential**: `google_api_key` + +## Test locally + +```bash +DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG=/abs/path/to/external_callout_config.json \ + sf data-code-extension function run \ + --entrypoint payload/entrypoint.py \ + --test-with payload/tests/test.json \ + --target-org +``` + +With `--target-org` the SDK fetches only the **URL** from the org's Named +Credential; **auth is always taken from `external_callout_config.json`** locally +(the org's External Credential is used only in the Data Cloud runtime). So the +`X-goog-api-key` must be in the local config for a local test. Omit +`--target-org` to run fully offline using `target_url`. + +```json +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_GEMINI_API_KEY" }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} +``` + +Place `external_callout_config.json` in the **parent of your payload folder** (or +point `DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG` at it). It is never packaged into +the deployment zip. Get a key from [Google AI Studio](https://aistudio.google.com/apikey); +**do not commit it.** + +## Auth types + +`auth_type` selects how auth is injected for local testing. It should mirror the +External Credential your Named Credential uses in the org, so local and deployed +runs behave the same. This example uses `Custom` (Gemini's `X-goog-api-key`); +all four supported types: + +```json +{ + "credentials": { + "callout:my_custom_api": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_API_KEY" } + }, + "callout:my_basic_api": { + "auth_type": "Basic", + "username": "svc_user", + "password": "YOUR_PASSWORD" + }, + "callout:my_oauth_api": { + "auth_type": "OAuth", + "access_token": "YOUR_ACCESS_TOKEN" + }, + "callout:my_jwt_api": { + "auth_type": "Jwt", + "token": "YOUR_JWT" + } + } +} +``` + +| `auth_type` | Fields read | Header sent | +| ----------- | ------------------------------- | --------------------------------------- | +| `Basic` | `username`, `password` | `Authorization: Basic ` | +| `Custom` | `custom_headers` (sent verbatim)| the headers you list | +| `OAuth` | `access_token` or `token` | `Authorization: Bearer ` | +| `Jwt` | `access_token` or `token` | `Authorization: Bearer ` | + +`OAuth`/`Jwt` take a token you supply for the local run — the SDK does not fetch +or refresh it. In the Data Cloud runtime the Named Credential handles token +acquisition; this local config only stands in for that during testing. diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py b/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py new file mode 100644 index 0000000..934b32e --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 + +""" +Document Chunking with a Gemini Named Credential Callout + +Splits each input document into paragraph-sized chunks and classifies every +chunk via Google's Gemini ``generateContent`` API, reached through a Named +Credential (``callout:gemini``) so the endpoint URL and API key are resolved +outside this code. The classification is attached to each chunk as citations. +""" + +import json +import logging + +from datacustomcode.function import Runtime +from datacustomcode.function.feature_types.chunking import ( + ChunkType, + SearchIndexChunkingV1Output, + SearchIndexChunkingV1Request, + SearchIndexChunkingV1Response, +) +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request_builder import ( + HTTPRequestBuilder, +) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +CALLOUT_URL = "callout:gemini" + +_ANALYSIS_FIELDS = ("summary", "category", "sentiment") + +_PROMPT = ( + "Analyze the following document chunk and classify it. Respond with its " + "one-sentence summary, a single-word category, overall sentiment " + "(positive, negative, or neutral), and up to five key topics.\n\nChunk:\n" +) + +# Force Gemini to return the classification as JSON in a fixed shape. +_GENERATION_CONFIG = { + "responseMimeType": "application/json", + "responseSchema": { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "category": {"type": "string"}, + "sentiment": {"type": "string"}, + "topics": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["summary", "category", "sentiment", "topics"], + }, +} + + +def _chunk_text(text: str, max_words: int = 80) -> list[str]: + """Split text into paragraph-aligned chunks of at most ``max_words`` words.""" + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + + chunks: list[str] = [] + current: list[str] = [] + current_words = 0 + + for paragraph in paragraphs: + paragraph_words = len(paragraph.split()) + if current and current_words + paragraph_words > max_words: + chunks.append("\n\n".join(current)) + current = [] + current_words = 0 + current.append(paragraph) + current_words += paragraph_words + + if current: + chunks.append("\n\n".join(current)) + + return chunks + + +def _extract_model_json(body: str) -> dict: + """Decode the model's JSON classification from a Gemini response. + + The generated text sits at ``candidates[0].content.parts[0].text`` and is + itself a JSON string, so decode twice. Any malformed layer yields ``{}``. + """ + try: + envelope = json.loads(body) if body else {} + except json.JSONDecodeError: + return {} + + try: + text = envelope["candidates"][0]["content"]["parts"][0]["text"] + except (KeyError, IndexError, TypeError): + return {} + + try: + payload = json.loads(text) + except json.JSONDecodeError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _analyze_chunk(chunk_text: str, runtime: Runtime) -> dict[str, str]: + """Classify one chunk via the Gemini callout and return it as citations.""" + request = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json", "Accept": "application/json"}) + .build() + ) + + payload = { + "contents": [{"parts": [{"text": _PROMPT + chunk_text}]}], + "generationConfig": _GENERATION_CONFIG, + } + response = runtime.named_credential.request(request, json.dumps(payload)) + + # Don't raise: a single failed callout shouldn't abort the whole job. + if not response.is_success: + logger.error(f"Gemini callout failed with status {response.status_code}") + return {"analysis_status": "failed", "http_status": str(response.status_code)} + + data = _extract_model_json(response.body) + citations = {"analysis_status": "success"} + for field in _ANALYSIS_FIELDS: + value = data.get(field) + citations[field] = str(value) if value is not None else "unavailable" + + topics = data.get("topics") + if isinstance(topics, list): + citations["topics"] = ", ".join(str(topic) for topic in topics) + + return citations + + +def function( + request: SearchIndexChunkingV1Request, runtime: Runtime +) -> SearchIndexChunkingV1Response: + """Chunk each input document and classify every chunk via the Gemini API.""" + logger.info(f"Received {len(request.input)} documents to chunk") + + chunks = [] + chunk_id = 1 + + for doc in request.input: + for chunk_text in _chunk_text(doc.text): + citations = _analyze_chunk(chunk_text, runtime) + + chunk = SearchIndexChunkingV1Output( + text=chunk_text, + seq_no=chunk_id, + chunk_type=ChunkType.TEXT, + citations=citations, + ) + chunks.append(chunk) + chunk_id += 1 + + logger.info(f"Produced {len(chunks)} classified chunks") + return SearchIndexChunkingV1Response(output=chunks) diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json new file mode 100644 index 0000000..2f979ea --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json @@ -0,0 +1,11 @@ +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { + "X-goog-api-key": "YOUR_API_KEY" + }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json new file mode 100644 index 0000000..c4d9960 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json @@ -0,0 +1,16 @@ +{ + "input": [ + { + "text": "Product Review: Northstar Analytics\n\nWe rolled Northstar out to our whole revenue team last quarter and the difference has been night and day. Dashboards that used to take our analysts a full day to assemble now refresh in seconds, and the natural-language query box means our account executives can answer their own questions without filing a ticket.\n\nOnboarding was smoother than any tool we have adopted in years. The guided setup imported our Salesforce data on the first try and the sample templates gave us something useful on day one. Support answered our two questions within the hour. Easily the best purchase decision we made this year." + }, + { + "text": "Support Ticket #48210: Repeated timeouts on scheduled exports\n\nFor the third week running our nightly export to the data warehouse has failed silently. There is no alert, no email, nothing in the activity log, and we only find out when the morning report is empty and the leadership meeting has no numbers.\n\nI have raised this twice already and both times the ticket was closed as resolved without anyone actually contacting me. This is costing us real credibility internally and I am extremely frustrated. If the connector cannot handle our volume we need to know now so we can plan a migration, because right now the product is not doing the one job we bought it for." + }, + { + "text": "Renewal Feedback: mixed feelings heading into year two\n\nThe core product is genuinely good. The reporting engine is fast, the permissions model is granular enough for our compliance team, and our analysts like working in it. On the functionality alone I would renew without hesitation.\n\nWhat gives me pause is the pricing. The per-seat cost jumped noticeably at renewal and several add-ons that used to be included are now separate line items. The value is still there, but the conversation with my finance team was harder than it should have been, and I would like more transparency before the next cycle." + }, + { + "text": "Feature Request: scheduled report subscriptions\n\nWe would like the ability to subscribe internal stakeholders to a report on a recurring schedule so a PDF lands in their inbox every Monday morning. Today we export manually and forward it, which is workable but easy to forget.\n\nA few teams have asked whether subscriptions could support filtered views per recipient, for example each regional manager receiving only their own territory. Not urgent for us, but it would remove a recurring bit of manual work and is something a couple of competing tools already offer." + } + ] +} diff --git a/tests/test_client.py b/tests/test_client.py index 86b72aa..9ea7d64 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -11,6 +11,7 @@ DataCloudObjectType, einstein_predict_col, llm_gateway_generate_text_col, + named_credential_request_col, ) from datacustomcode.config import ( AccessLayerObjectConfig, @@ -426,6 +427,65 @@ def test_delegates_to_spark_predictions(self, mock_build): ) +class TestNamedCredentialRequestCol: + """The module-level ``named_credential_request_col`` is a thin wrapper that + resolves the client-owned :class:`SparkNamedCredential` and delegates. + """ + + @patch("datacustomcode.client._build_spark_named_credential") + def test_delegates_to_spark_named_credential(self, mock_build, reset_client): + mock_nc = MagicMock() + sentinel_col = MagicMock(name="col") + mock_nc.request_col.return_value = sentinel_col + mock_build.return_value = mock_nc + + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + Client(reader=reader, writer=writer) + + request = MagicMock(name="request") + body_col = MagicMock(name="body_col") + result = named_credential_request_col(request, body_col) + + assert result is sentinel_col + mock_nc.request_col.assert_called_once_with(request, body=body_col) + + +class TestClientNamedCredentialRequest: + + @patch("datacustomcode.client._build_spark_named_credential") + def test_forwards_request_and_body(self, mock_build, reset_client): + mock_nc = MagicMock() + sentinel = MagicMock(name="response") + mock_nc.request.return_value = sentinel + mock_build.return_value = mock_nc + + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + client = Client(reader=reader, writer=writer) + + request = MagicMock(name="request") + result = client.named_credential_request(request, '{"k": "v"}') + + assert result is sentinel + mock_nc.request.assert_called_once_with(request, body='{"k": "v"}') + + @patch("datacustomcode.client._build_spark_named_credential") + def test_named_credential_built_lazily_and_cached(self, mock_build, reset_client): + mock_build.return_value = MagicMock() + + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + client = Client(reader=reader, writer=writer) + + mock_build.assert_not_called() + + client.named_credential_request(MagicMock()) + client.named_credential_request(MagicMock()) + + mock_build.assert_called_once_with() + + # Add tests for DefaultSparkSessionProvider class TestDefaultSparkSessionProvider: diff --git a/tests/test_deploy.py b/tests/test_deploy.py index af804d3..7fc91f8 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -761,6 +761,36 @@ def test_zip_with_function_package_type( ) assert mock_zipfile_instance.write.call_count == 3 # One call per file + @patch("datacustomcode.deploy.has_nonempty_requirements_file") + @patch("datacustomcode.deploy.prepare_dependency_archive") + @patch("zipfile.ZipFile") + @patch("os.walk") + def test_zip_excludes_ds_store_and_external_credentials( + self, mock_walk, mock_zipfile, mock_prepare, mock_has_requirements + ): + """Local credentials and .DS_Store must never be packaged into the zip.""" + mock_has_requirements.return_value = False + mock_zipfile_instance = MagicMock() + mock_zipfile.return_value.__enter__.return_value = mock_zipfile_instance + mock_zipfile_instance.write = MagicMock() + + mock_walk.return_value = [ + ( + "/test/dir", + [], + ["entrypoint.py", ".DS_Store", "external_callout_config.json"], + ), + ] + + zip("/test/dir", "default", "script") + + # Only entrypoint.py is written; excluded files are skipped. + assert mock_zipfile_instance.write.call_count == 1 + written = [call.args[0] for call in mock_zipfile_instance.write.call_args_list] + assert written == ["/test/dir/entrypoint.py"] + assert not any("external_callout_config.json" in path for path in written) + assert not any(".DS_Store" in path for path in written) + class TestUploadZip: @patch("datacustomcode.deploy.requests.put") diff --git a/tests/test_named_credential_direct.py b/tests/test_named_credential_direct.py index 790f60f..ab48271 100644 --- a/tests/test_named_credential_direct.py +++ b/tests/test_named_credential_direct.py @@ -70,6 +70,16 @@ def test_jwt_sets_bearer_from_token(self): request = handler(_prepared_request()) assert request.headers["Authorization"] == "Bearer jwt-tok" + def test_oauth_missing_token_raises(self): + handler = DynamicAuthHandler({"auth_type": AuthType.OAUTH.value}) + with pytest.raises(ValueError): + handler(_prepared_request()) + + def test_jwt_missing_token_raises(self): + handler = DynamicAuthHandler({"auth_type": AuthType.JWT.value}) + with pytest.raises(ValueError): + handler(_prepared_request()) + def test_unsupported_auth_type_raises_value_error(self): handler = DynamicAuthHandler({"auth_type": "Nonsense"}) with pytest.raises(ValueError): @@ -276,6 +286,12 @@ def test_callout_rejects_non_callout_url(self, tmp_path, monkeypatch): with pytest.raises(CredentialError): transport.callout({"path": "https://example.com/x", "method": "GET"}) + @pytest.mark.parametrize("path", ["callout:", "callout:/path"]) + def test_callout_rejects_empty_named_credential(self, tmp_path, monkeypatch, path): + transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) + with pytest.raises(CredentialError): + transport.callout({"path": path, "method": "GET", "headers": {}}) + def test_base_url_resolved_once_per_callout_key(self, tmp_path, monkeypatch): from datacustomcode.named_credential.direct import transport as transport_mod From bcfa03d80fb95d3873d5b737532401b8161fa3c0 Mon Sep 17 00:00:00 2001 From: Diksha Date: Mon, 3 Aug 2026 16:21:51 +0530 Subject: [PATCH 14/17] Fix error on deploy --- src/datacustomcode/cli.py | 2 +- src/datacustomcode/constants.py | 2 +- src/datacustomcode/deploy.py | 8 ++--- .../named_credential/__init__.py | 2 ++ .../named_credential/direct/transport.py | 12 +++---- src/datacustomcode/named_credential/errors.py | 36 +++++++++++++++++++ tests/test_deploy.py | 2 +- tests/test_named_credential_direct.py | 5 +-- 8 files changed, 54 insertions(+), 15 deletions(-) create mode 100644 src/datacustomcode/named_credential/errors.py diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py index 7e2cd00..b31deb7 100644 --- a/src/datacustomcode/cli.py +++ b/src/datacustomcode/cli.py @@ -262,7 +262,7 @@ def deploy( mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get( use_in_feature, use_in_feature ) - metadata.functionInvokeOptions = [mapped_feature] + metadata.invokeOptions = [mapped_feature] try: if sf_cli_org: diff --git a/src/datacustomcode/constants.py b/src/datacustomcode/constants.py index 76b6a7c..dbef13e 100644 --- a/src/datacustomcode/constants.py +++ b/src/datacustomcode/constants.py @@ -35,7 +35,7 @@ # Feature name to Connect API name mapping USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = { - "SearchIndexChunking": "UnstructuredChunking", + "SearchIndexChunking": "SearchIndexChunking", } # Pydantic request/response type names to feature names diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index 391e7b9..f13cdc1 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -42,7 +42,7 @@ ) from datacustomcode.scan import find_base_directory, get_package_type -DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code" +DATA_CUSTOM_CODE_PATH = "services/data/v67.0/ssot/data-custom-code" DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms" WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000 @@ -110,7 +110,7 @@ class CodeExtensionMetadata(BaseModel): description: str computeType: str codeType: str - functionInvokeOptions: Union[list[str], None] = None + invokeOptions: Union[list[str], None] = None def __init__(self, **data): name = data.get("name", "") @@ -214,8 +214,8 @@ def create_deployment( "codeType": metadata.codeType, } ) - 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( diff --git a/src/datacustomcode/named_credential/__init__.py b/src/datacustomcode/named_credential/__init__.py index 0f97b7e..31245db 100644 --- a/src/datacustomcode/named_credential/__init__.py +++ b/src/datacustomcode/named_credential/__init__.py @@ -15,6 +15,7 @@ from datacustomcode.named_credential.base import NamedCredential from datacustomcode.named_credential.default import DefaultNamedCredential +from datacustomcode.named_credential.errors import NamedCredentialCallError from datacustomcode.named_credential.spark_base import SparkNamedCredential from datacustomcode.named_credential.spark_default import DefaultSparkNamedCredential @@ -22,5 +23,6 @@ "DefaultNamedCredential", "DefaultSparkNamedCredential", "NamedCredential", + "NamedCredentialCallError", "SparkNamedCredential", ] diff --git a/src/datacustomcode/named_credential/direct/transport.py b/src/datacustomcode/named_credential/direct/transport.py index 38dda2d..a81a909 100644 --- a/src/datacustomcode/named_credential/direct/transport.py +++ b/src/datacustomcode/named_credential/direct/transport.py @@ -31,11 +31,9 @@ import requests from datacustomcode.named_credential.direct.auth import DynamicAuthHandler -from datacustomcode.named_credential.direct.credentials import ( - CredentialError, - CredentialStore, -) +from datacustomcode.named_credential.direct.credentials import CredentialStore from datacustomcode.named_credential.direct.url_resolver import resolve_base_url +from datacustomcode.named_credential.errors import NamedCredentialCallError from datacustomcode.token_provider import ( CredentialsTokenProvider, SFCLITokenProvider, @@ -69,7 +67,7 @@ def _build_token_provider( def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: raw_url = callout_request["path"] if not raw_url.startswith("callout:"): - raise CredentialError( + raise NamedCredentialCallError( f"Callout URL must start with 'callout:', got '{raw_url}'." ) @@ -83,7 +81,9 @@ def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]: path_suffix = raw_url[sep_idx:] if callout_key == "callout:": - raise CredentialError(f"Named Credential name is empty in URL '{raw_url}'.") + raise NamedCredentialCallError( + f"Named Credential name is empty in URL '{raw_url}'." + ) cred_config = self._store.get(callout_key) base_url = self._base_url_cache.get(callout_key) diff --git a/src/datacustomcode/named_credential/errors.py b/src/datacustomcode/named_credential/errors.py new file mode 100644 index 0000000..e156fe4 --- /dev/null +++ b/src/datacustomcode/named_credential/errors.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Exceptions raised by Named Credential implementations.""" + +from __future__ import annotations + +from typing import Optional + + +class NamedCredentialCallError(RuntimeError): + """Raised when a Named Credential external callout fails.""" + + def __init__( + self, + message: str, + *, + status: Optional[object] = None, + error_code: Optional[str] = None, + error_message: Optional[str] = None, + ) -> None: + super().__init__(message) + self.status = status + self.error_code = error_code + self.error_message = error_message diff --git a/tests/test_deploy.py b/tests/test_deploy.py index 7fc91f8..f42b3bc 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -657,7 +657,7 @@ def test_create_deployment_function_invoke_options(self, mock_make_api_call): version="1.0.0", description="Test job", computeType="CPU_M", - functionInvokeOptions=["option1", "option2"], + invokeOptions=["option1", "option2"], codeType="function", ) diff --git a/tests/test_named_credential_direct.py b/tests/test_named_credential_direct.py index ab48271..fab0588 100644 --- a/tests/test_named_credential_direct.py +++ b/tests/test_named_credential_direct.py @@ -29,6 +29,7 @@ CredentialStore, ) from datacustomcode.named_credential.direct.url_resolver import resolve_base_url +from datacustomcode.named_credential.errors import NamedCredentialCallError def _prepared_request() -> PreparedRequest: @@ -283,13 +284,13 @@ def fake_request(**kwargs): def test_callout_rejects_non_callout_url(self, tmp_path, monkeypatch): transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) - with pytest.raises(CredentialError): + with pytest.raises(NamedCredentialCallError): transport.callout({"path": "https://example.com/x", "method": "GET"}) @pytest.mark.parametrize("path", ["callout:", "callout:/path"]) def test_callout_rejects_empty_named_credential(self, tmp_path, monkeypatch, path): transport = self._make_transport(tmp_path, monkeypatch, {"auth_type": "Custom"}) - with pytest.raises(CredentialError): + with pytest.raises(NamedCredentialCallError): transport.callout({"path": path, "method": "GET", "headers": {}}) def test_base_url_resolved_once_per_callout_key(self, tmp_path, monkeypatch): From fa645b5e2dfd45bc958b7e4718e3a50a5c57001a Mon Sep 17 00:00:00 2001 From: Diksha Date: Mon, 3 Aug 2026 16:38:53 +0530 Subject: [PATCH 15/17] Version bump in mock file --- scripts/mock_sf_server.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/mock_sf_server.py b/scripts/mock_sf_server.py index b296577..19c6ab5 100644 --- a/scripts/mock_sf_server.py +++ b/scripts/mock_sf_server.py @@ -18,11 +18,11 @@ Returns fake rows with the columns expected by the default script template (Account_std__dll: description__c, sfdcorganizationid__c, kq_id__c). -POST /services/data/v63.0/ssot/data-custom-code +POST /services/data/v67.0/ssot/data-custom-code Called by deploy_full() → create_deployment(). Returns a fake fileUploadUrl pointing back at this server. -GET /services/data/v63.0/ssot/data-custom-code/* +GET /services/data/v67.0/ssot/data-custom-code/* Called by deploy_full() → wait_for_deployment() → get_deployments(). Returns deploymentStatus=Deployed immediately so the poll loop exits. @@ -87,7 +87,7 @@ ], } -_DATA_CUSTOM_CODE_PATH = "/services/data/v63.0/ssot/data-custom-code" +_DATA_CUSTOM_CODE_PATH = "/services/data/v67.0/ssot/data-custom-code" _DATA_TRANSFORMS_PATH = "/services/data/v63.0/ssot/data-transforms" From 75ad623d0c67f01fcb6b657ec3dd3b7e2ff6a3d3 Mon Sep 17 00:00:00 2001 From: Diksha Date: Mon, 3 Aug 2026 20:49:54 +0530 Subject: [PATCH 16/17] revert change to mock server --- scripts/mock_sf_server.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/mock_sf_server.py b/scripts/mock_sf_server.py index 19c6ab5..b296577 100644 --- a/scripts/mock_sf_server.py +++ b/scripts/mock_sf_server.py @@ -18,11 +18,11 @@ Returns fake rows with the columns expected by the default script template (Account_std__dll: description__c, sfdcorganizationid__c, kq_id__c). -POST /services/data/v67.0/ssot/data-custom-code +POST /services/data/v63.0/ssot/data-custom-code Called by deploy_full() → create_deployment(). Returns a fake fileUploadUrl pointing back at this server. -GET /services/data/v67.0/ssot/data-custom-code/* +GET /services/data/v63.0/ssot/data-custom-code/* Called by deploy_full() → wait_for_deployment() → get_deployments(). Returns deploymentStatus=Deployed immediately so the poll loop exits. @@ -87,7 +87,7 @@ ], } -_DATA_CUSTOM_CODE_PATH = "/services/data/v67.0/ssot/data-custom-code" +_DATA_CUSTOM_CODE_PATH = "/services/data/v63.0/ssot/data-custom-code" _DATA_TRANSFORMS_PATH = "/services/data/v63.0/ssot/data-transforms" From 12eabe1e6f547e29d507337e14b389e3cc6bc32d Mon Sep 17 00:00:00 2001 From: Diksha Date: Fri, 7 Aug 2026 17:58:02 +0530 Subject: [PATCH 17/17] Add example for script --- .../examples/external_callout/README.md | 135 +++++++++++++++++ .../examples/external_callout/entrypoint.py | 142 ++++++++++++++++++ .../external_callout_config.json | 11 ++ 3 files changed, 288 insertions(+) create mode 100644 src/datacustomcode/templates/script/examples/external_callout/README.md create mode 100644 src/datacustomcode/templates/script/examples/external_callout/entrypoint.py create mode 100644 src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json diff --git a/src/datacustomcode/templates/script/examples/external_callout/README.md b/src/datacustomcode/templates/script/examples/external_callout/README.md new file mode 100644 index 0000000..dababe6 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/README.md @@ -0,0 +1,135 @@ +# Transform with a Gemini Named Credential Callout + +The **transform** calls Google's **Gemini** `generateContent` API to summarize text, and +write the result back to a DLO. Gemini is reached through a **Named Credential** +(`callout:gemini`), so this code never handles the endpoint URL or the API key. + +It shows **both** callout paths against the same Named Credential. + +## Shared request template + +```python +from datacustomcode.client import Client, named_credential_request_col + +# URL, method and headers apply to every callout on this template. +request = ( + HTTPRequestBuilder() + .set_url("callout:gemini") # callout:[/] + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json"}) + .build() +) +``` + +## Driver path — one-shot on the driver + +`Client.named_credential_request` runs the callout **once** on the driver and +returns an `HTTPResponse` (`.status_code`, `.body`, `.headers`, `.is_success`). +Use it for a lookup or a shared value you reuse across the job. + +```python +response = client.named_credential_request(request, body=json.dumps(payload)) +if response.is_success: + envelope = json.loads(response.body) + text = envelope["candidates"][0]["content"]["parts"][0]["text"] +``` + +## Per-row path — fan out across the DataFrame + +`named_credential_request_col` dispatches one callout per row; only the body +Column varies. + +```python +# One callout per row; body is a Column built from the row's data. +df = df.withColumn("_callout", named_credential_request_col(request, body=body_col)) +``` + +It returns a struct Column: + +``` +{status, response: {status_code, body, headers}, error_code, error_message} +``` + +- A non-2xx response is still `status = "SUCCESS"` with the HTTP code in + `response.status_code` — extracting the model text just yields null for that row. +- A transport failure sets `status = "ERROR"`; the row survives, the job does not + abort. Pull the model text out of `response.body` with `get_json_object(...)`. + +Both paths resolve the same Named Credential and read auth from the same local +`external_callout_config.json`. + +The `gemini` Named Credential's URL already includes the full +`/v1beta/models/:generateContent` path, so the callout is just +`callout:gemini` with **no path suffix** (anything after the name is appended to +the credential's URL). + +## Configure the Named Credential + +1. Create an **External Credential** (e.g. `google_api_key`) that injects your + Gemini API key as the `X-goog-api-key` header. +2. Create a **Named Credential** named `gemini`: + - **URL**: `https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent` + - **Enabled for Callouts** + **Generate Authorization Header**: on + - **External Credential**: `google_api_key` + +## Test locally + +Copy `entrypoint.py` into your `payload/` folder (or point the run at it), then: + +```bash +DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG=/abs/path/to/external_callout_config.json \ + sf data-code-extension script run --entrypoint entrypoint.py --target-org +``` + +With `----target-org` the SDK fetches only the **URL** from the org's Named +Credential; **auth is always taken from `external_callout_config.json`** locally +(the org's External Credential is used only in the Data Cloud runtime). So the +`X-goog-api-key` must be in the local config for a local test. Omit +`--target-org` to run fully offline using `target_url`. + +```json +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_GEMINI_API_KEY" }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} +``` + +Place `external_callout_config.json` in the **parent of your payload folder** (or +point `DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG` at it). It is never packaged into +the deployment zip. Get a key from [Google AI Studio](https://aistudio.google.com/apikey); +**do not commit it.** + +## What it reads / writes + +`config.json` declares the DLO permissions for deployment: + +| | DLO | Notes | +| ------ | ------------------- | --------------------------------------------- | +| read | `Account_std__dll` | source rows; `description__c` is summarized | +| write | `Account_std_copy__dll` | adds `summary__c`, `callout_status__c`, `callout_http_code__c` | + +Adjust `_TEXT_COLUMN`, `_SOURCE_DLO` and `_TARGET_DLO` in `entrypoint.py` (and the +matching entries in `config.json`) to point at your own DLOs. + +## Auth types + +`auth_type` selects how auth is injected for local testing. It should mirror the +External Credential your Named Credential uses in the org, so local and deployed +runs behave the same. This example uses `Custom` (Gemini's `X-goog-api-key`); +all four supported types: + +| `auth_type` | Fields read | Header sent | +| ----------- | ------------------------------- | --------------------------------------- | +| `Basic` | `username`, `password` | `Authorization: Basic ` | +| `Custom` | `custom_headers` (sent verbatim)| the headers you list | +| `OAuth` | `access_token` or `token` | `Authorization: Bearer ` | +| `Jwt` | `access_token` or `token` | `Authorization: Bearer ` | + +`OAuth`/`Jwt` take a token you supply for the local run — the SDK does not fetch +or refresh it. In the Data Cloud runtime the Named Credential handles token +acquisition; this local config only stands in for that during testing. diff --git a/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py b/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py new file mode 100644 index 0000000..94ac6d1 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 + +""" +Data Transform with a Gemini Named Credential Callout + +This is a batch transform: read a DLO, enrich it via Google's Gemini ``generateContent`` +API, and write the result back to a DLO. + +It shows **both** callout paths against the same Named Credential +(``callout:gemini``); the request template (URL, method, headers) is shared and +only the body differs: + +- **Driver path** — :meth:`Client.named_credential_request` runs **once** on the + driver and returns an :class:`HTTPResponse`. Use it for a single job-level + callout whose result you reuse across the job. +- **Per-row path** — :func:`datacustomcode.client.named_credential_request_col` + fans the callout out across the DataFrame, one call per row, returning a struct + Column ``{status, response, error_code, error_message}`` where ``response`` is + itself ``{status_code, body, headers}``. A non-2xx response is still + ``SUCCESS`` with its code in ``response.status_code``; a transport failure sets + ``status`` to ``ERROR`` without aborting the whole Spark job. +""" + +import json +import logging + +from pyspark.sql.functions import ( + array, + col, + concat, + get_json_object, + lit, + struct, + to_json, +) + +from datacustomcode.client import Client, named_credential_request_col +from datacustomcode.io.writer.base import WriteMode +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request_builder import ( + HTTPRequestBuilder, +) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +CALLOUT_URL = "callout:gemini" + +_TEXT_COLUMN = "Description__c" +_SOURCE_DLO = "Account_std__dll" +_TARGET_DLO = "Account_std_copy__dll" + +_PROMPT = ( + "Summarize the following text in one sentence. Respond with the summary " + "only, no preamble.\n\nText:\n" +) + +_REQUEST = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json", "Accept": "application/json"}) + .build() +) + + +def _gemini_body_col(text_col: "col") -> "col": # noqa: F821 + """Build a per-row Gemini ``generateContent`` request body as a JSON string. + + Using ``to_json(struct(...))`` keeps the row text properly escaped inside the + JSON payload rather than string-concatenating it. + """ + prompt = concat(lit(_PROMPT), text_col) + contents = array(struct(array(struct(prompt.alias("text"))).alias("parts"))) + return to_json(struct(contents.alias("contents"))) + + +def _gemini_body(text: str) -> str: + """Build a Gemini ``generateContent`` request body as a JSON string (driver).""" + return json.dumps({"contents": [{"parts": [{"text": _PROMPT + text}]}]}) + + +def _summarize_on_driver(client: Client, text: str) -> str: + """One-shot driver callout: summarize a single string once, not per row. + + The scalar counterpart to the per-row column path — same request template, + but dispatched once on the driver and returning an ``HTTPResponse``. + """ + response = client.named_credential_request(_REQUEST, body=_gemini_body(text)) + + # Don't raise: a failed driver callout shouldn't abort the whole job. + if not response.is_success: + logger.error(f"Driver Gemini callout failed: HTTP {response.status_code}") + return "" + + envelope = json.loads(response.body) if response.body else {} + try: + return envelope["candidates"][0]["content"]["parts"][0]["text"] + except (KeyError, IndexError, TypeError): + return "" + + +def main(): + client = Client() + + df = client.read_dlo(_SOURCE_DLO) + + # Driver path: one callout on the driver over a single representative row. + sample = df.select(_TEXT_COLUMN).first() + if sample and sample[0]: + driver_summary = _summarize_on_driver(client, sample[0]) + logger.info(f"Driver-path sample summary: {driver_summary}") + + # Per-row path: one Gemini callout per row; the result struct is a column. + callout = named_credential_request_col( + _REQUEST, body=_gemini_body_col(col(_TEXT_COLUMN)) + ) + df = df.withColumn("_callout", callout) + + # Pull the model's text out of the response body. A row whose callout failed + # (non-2xx or transport error) yields null here rather than failing the job. + summary = get_json_object( + col("_callout")["response"]["body"], + "$.candidates[0].content.parts[0].text", + ) + + df = df.select( + col(_ID_COLUMN).alias("id__c"), + col(_TEXT_COLUMN).alias("description__c"), + col(_KQ_ID_COLUMN).alias("kq_id__c"), + summary.alias("summary__c"), + col("_callout")["status"].alias("callout_status__c"), + col("_callout")["response"]["status_code"].alias("callout_http_code__c"), + ) + + client.write_to_dlo(_TARGET_DLO, df, write_mode=WriteMode.APPEND) + + +if __name__ == "__main__": + main() diff --git a/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json b/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json new file mode 100644 index 0000000..2f979ea --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json @@ -0,0 +1,11 @@ +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { + "X-goog-api-key": "YOUR_API_KEY" + }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +}