diff --git a/CHANGELOG.md b/CHANGELOG.md index 03b8ca4..fb41eee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,30 @@ # Changelog +## 6.1.0 + +### Added + +- **`StreamingClient` for BYOC streaming (delta) transforms.** + + 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) + query.awaitTermination() + ``` + + 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..a5d77fc 100644 --- a/README.md +++ b/README.md @@ -145,15 +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 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 @@ -169,6 +176,36 @@ 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 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 StreamingClient + +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. +deltas = client.read_dlo_deltas() + +# 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. +query = client.write_dlo_deltas("Output__dll", transformed) +query.awaitTermination() +``` + +Notes: + +- 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/__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/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/client.py b/src/datacustomcode/client.py index 0d4df78..5a03b84 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -21,7 +21,9 @@ ClassVar, Dict, Optional, + TypeVar, Union, + cast, ) from datacustomcode.config import config @@ -29,21 +31,64 @@ 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: 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 from datacustomcode.einstein_predictions.types import PredictionType 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 +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 ``streamingSource`` field. + + Raises: + RuntimeError: If no ``streaming_source`` has been configured (e.g. the + 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 'streamingSource'." + ) + 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": """Instantiate the SDK-configured :class:`SparkLLMGateway`. @@ -99,7 +144,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) @@ -118,6 +163,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", @@ -161,12 +221,46 @@ 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 ) +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"]``. 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 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" @@ -205,29 +299,26 @@ 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`. + spark_named_credential: Optional custom :class:`SparkNamedCredential`. Example: >>> client = Client() @@ -237,47 +328,68 @@ class Client: >>> 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 _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 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, + spark_named_credential: Optional[SparkNamedCredential] = 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 + 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: - # 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( @@ -300,66 +412,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 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] + 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. @@ -474,6 +537,41 @@ 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[str] = 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 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. + """ + 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: @@ -487,3 +585,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/config.py b/src/datacustomcode/config.py index 901b295..1e2bead 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 ``streamingSource`` field + 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/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/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..18a4f86 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 @@ -37,10 +38,14 @@ 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" 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 +113,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 +206,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 +226,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( @@ -236,6 +251,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( @@ -388,6 +405,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 +443,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 +523,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 +568,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, } @@ -598,9 +633,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 EXCLUDED_FILES: abs_path = os.path.join(root, file) arcname = os.path.relpath(abs_path, directory) zipf.write(abs_path, arcname) @@ -616,9 +651,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/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/io/reader/base.py b/src/datacustomcode/io/reader/base.py index ddea31b..c19769b 100644 --- a/src/datacustomcode/io/reader/base.py +++ b/src/datacustomcode/io/reader/base.py @@ -41,3 +41,45 @@ def read_dmo( 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 + 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. + + 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) -> 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. + + 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..47a7bd2 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,35 @@ 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 + ) -> 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. 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. + + 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/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..1478234 --- /dev/null +++ b/src/datacustomcode/named_credential/base.py @@ -0,0 +1,54 @@ +# 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, 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[str] = 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 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/default.py b/src/datacustomcode/named_credential/default.py new file mode 100644 index 0000000..c0e72c2 --- /dev/null +++ b/src/datacustomcode/named_credential/default.py @@ -0,0 +1,93 @@ +# 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 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 + ``external_callout_config.json``) and injecting auth from the same file. + """ + + 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[str] = None, + ) -> 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, + "headers": dict(request.headers), + "body": body if body is not None else "", + } + callout_response = self._callout(callout_request) + 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 ``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..74bd837 --- /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 ``external_callout_config.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..24c59c4 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/auth.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. +""" +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 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}'.") + + 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..8b60864 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/credentials.py @@ -0,0 +1,121 @@ +# 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 ``external_callout_config.json``. + +The file lives in the parent of the payload folder so it is never packaged in +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 + +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. +EXTERNAL_CALLOUT_CREDENTIAL = "external_callout_config.json" +# Absolute-path override, primarily for tests and non-standard layouts. +CREDENTIAL_FILE_ENV_VAR = "DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG" + + +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 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 / EXTERNAL_CALLOUT_CREDENTIAL + if candidate.is_file(): + return candidate + return None + + +class CredentialStore: + """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 + 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 '{EXTERNAL_CALLOUT_CREDENTIAL}'. 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 + + # 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 with a 'credentials' section " + f"keyed by callout reference." + ) + 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``). + + 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 the 'credentials' section of " + f"'{EXTERNAL_CALLOUT_CREDENTIAL}'." + ) + 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..a81a909 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/transport.py @@ -0,0 +1,110 @@ +# 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 +``external_callout_config.json``), attaches the credential's auth, and returns the raw +``{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 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, + 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 NamedCredentialCallError( + 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 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) + 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 are passed; the SDK assumes no Content-Type. + headers = dict(callout_request.get("headers", {})) + + response = requests.request( + method=callout_request["method"], + url=base_url + path_suffix, + headers=headers, + data=body, + auth=DynamicAuthHandler(cred_config), + timeout=30, + ) + return { + "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..6044ac2 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/url_resolver.py @@ -0,0 +1,112 @@ +# 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 falling back to target_url. + logger.warning( + "Connect API URL resolution failed for %s: %s. " + "Falling back to 'target_url' from external_callout_config.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 ``external_callout_config.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 " + f"external_callout_config.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..2c47743 --- /dev/null +++ b/src/datacustomcode/named_credential/spark_base.py @@ -0,0 +1,93 @@ +# 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, + 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[str] = None, + ) -> HTTPResponse: + """Make an external callout through a Named Credential. + + Args: + request: The callout request + body: Optional request body. Set the + ``Content-Type`` header yourself; the SDK does not assume one. + + 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 request body as a + string, sent verbatim (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 — 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/named_credential/spark_default.py b/src/datacustomcode/named_credential/spark_default.py new file mode 100644 index 0000000..8f3ea23 --- /dev/null +++ b/src/datacustomcode/named_credential/spark_default.py @@ -0,0 +1,154 @@ +# 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.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 + + +_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[str] = 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), + ] + ) + + 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: + response = named_credential.request(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), + } + + return { + "status": _STATUS_SUCCESS, + "response": { + "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/__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..3f98ae1 --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_response.py @@ -0,0 +1,43 @@ +# 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 + +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" + ) + 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 + 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..f615af5 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 @@ -42,6 +43,15 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None: config_obj.options[key] = value +def _read_streaming_source(config_json: dict) -> Optional[str]: + """Return the streaming source name from config.json's ``streamingSource``.""" + source = config_json.get("streamingSource") + if not isinstance(source, dict): + 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]): if sf_cli_org: config_key = "sf_cli_org" @@ -55,6 +65,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 +76,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( @@ -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_streaming_source(config_json) + _update_config_options(profile, sf_cli_org) for dependency in dependencies: diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index 5e50c5d..059925a 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 @@ -32,6 +33,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 +46,20 @@ }, } +STREAMING_TRANSFORM_CONFIG_TEMPLATE = { + "sdkVersion": get_version(), + "entryPoint": "", + "dataspace": "default", + "streamingSource": { + "type": "dlo", + "name": "", + }, + "permissions": { + "read": {}, + "write": {}, + }, +} + FUNCTION_CONFIG_TEMPLATE = { "entryPoint": "", } @@ -160,6 +177,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 +214,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 +238,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 +257,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 +288,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 +379,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 +478,51 @@ 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/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/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/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/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" + } + } +} 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..97dea40 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -0,0 +1,49 @@ +"""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 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 + change-feed metadata columns (``_record_type``, ``_commit_*``). +* ``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 +``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 StreamingClient + + +def main(): + client = StreamingClient() + + # Streaming DataFrame over the source DLO's change feed. + deltas = client.read_dlo_deltas() + + # Ordinary PySpark transform. + transformed = deltas.withColumn("description__c", upper(col("description__c"))) + + # Start the streaming write. write_dlo_deltas returns the StreamingQuery; + # the trigger and checkpoint location are provided by the runtime. + query = client.write_dlo_deltas("Account_std_copy__dll", transformed) + + # Drive the query's lifecycle. In the streaming runtime this blocks until + # the job is stopped by the platform. + query.awaitTermination() + + +if __name__ == "__main__": + main() diff --git a/tests/io/reader/test_query_api.py b/tests/io/reader/test_query_api.py index f9baa81..0b19081 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() + 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() + 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..a10f2a3 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) + 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_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_client.py b/tests/test_client.py index 86b72aa..3037c50 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -9,8 +9,11 @@ Client, DataCloudAccessLayerException, DataCloudObjectType, + StreamingClient, + _BaseClient, einstein_predict_col, llm_gateway_generate_text_col, + named_credential_request_col, ) from datacustomcode.config import ( AccessLayerObjectConfig, @@ -80,10 +83,18 @@ 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.""" + 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: @@ -256,7 +267,232 @@ def test_read_pattern_flow(self, reset_client, mock_spark): assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO] -class TestClientLlmGatewayGenerateText: +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 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) + mock_df = MagicMock(spec=DataFrame) + 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() + + 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_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("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = None + with pytest.raises(RuntimeError) as exc_info: + client.read_dlo_deltas() + + assert "streamingSource" 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 = StreamingClient(reader=reader, writer=writer) + 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() + assert result is mock_df + 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) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + mock_query = MagicMock() + writer.write_dlo_deltas.return_value = mock_query + + client = StreamingClient(reader=reader, writer=writer) + client._record_dlo_access("some_dlo") + + result = client.write_dlo_deltas("test_dlo", mock_df, extra_param=True) + + writer.write_dlo_deltas.assert_called_once_with( + "test_dlo", mock_df, 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 = StreamingClient(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) + + 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 = StreamingClient(reader=reader, writer=writer) + + 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) + + 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] + + +class TestSharedSparkSession: + """Both client types must share a single Spark session (one connection).""" + + 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 + + with ( + patch("datacustomcode.client.config") as mock_config, + patch.object( + DefaultSparkSessionProvider, "get_session" + ) as mock_get_session, + ): + mock_get_session.return_value = mock_spark + + mock_reader_config = MagicMock() + mock_reader_config.to_object.return_value = MagicMock( + spec=BaseDataCloudReader + ) + mock_reader_config.force = False + + mock_writer_config = MagicMock() + mock_writer_config.to_object.return_value = MagicMock( + spec=BaseDataCloudWriter + ) + 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) + + client_cls() + return mock_get_session + + 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) + + # 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() + + assert _BaseClient._shared_spark is mock_spark + assert Client._instance is not StreamingClient._instance + + 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) + + from datacustomcode.spark.default import DefaultSparkSessionProvider + + 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) + + 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) + + StreamingClient() + + 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) + + 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): @@ -426,6 +662,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..0cbaf15 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -17,6 +17,7 @@ DloPermission, DmoPermission, Permissions, + StreamingSource, get_config, ) @@ -657,7 +658,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", ) @@ -671,6 +672,54 @@ 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") @@ -761,6 +810,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") @@ -1353,6 +1432,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 +1690,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_named_credential.py b/tests/test_named_credential.py new file mode 100644 index 0000000..60470a2 --- /dev/null +++ b/tests/test_named_credential.py @@ -0,0 +1,461 @@ +# 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.body == "" + + 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"}, "body": '{"ok": true}'} + ) + assert isinstance(response, HTTPResponse) + assert response.status_code == 200 + assert response.headers == {"X": "y"} + assert response.body == '{"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 {"status_code": 200, "headers": {}, "body": "{}"} + + monkeypatch.setattr(nc, "_get_transport", lambda: _FakeTransport()) + result = nc._callout({"path": "callout:NC/x", "method": "GET"}) + assert result["status_code"] == 200 + + def test_request_forwards_raw_body_and_response(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return { + "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.body == '{"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 {"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.body == "" + + def test_request_non_json_body_returned_verbatim(self, monkeypatch): + nc = DefaultNamedCredential() + + def fake_callout(callout_request): + 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.body == "plain text" + + def test_request_json_array_body_returned_verbatim(self, monkeypatch): + nc = DefaultNamedCredential() + + def fake_callout(callout_request): + 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.body == "[1, 2, 3]" + + def test_request_opaque_string_body_sent_verbatim(self, monkeypatch): + nc = DefaultNamedCredential() + captured = {} + + def fake_callout(callout_request): + captured.update(callout_request) + return {"status_code": 200, "headers": {}, "body": ""} + + monkeypatch.setattr(nc, "_callout", fake_callout) + request = HTTPRequestBuilder().set_url("callout:NC/path").build() + nc.request(request, "a=1&b=2") + + assert captured["body"] == "a=1&b=2" + + +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.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() + 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.request.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.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}') + + 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.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 ( + _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.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.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" + + 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.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) + + 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.request.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.request.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..fab0588 --- /dev/null +++ b/tests/test_named_credential_direct.py @@ -0,0 +1,320 @@ +# 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 +from typing import ClassVar + +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 +from datacustomcode.named_credential.errors import NamedCredentialCallError + + +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_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): + handler(_prepared_request()) + + +class TestCredentialStore: + def test_get_returns_config(self, tmp_path, monkeypatch): + cred_file = tmp_path / "external_callout_config.json" + cred_file.write_text( + json.dumps( + {"credentials": {"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 / "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" + + 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 / "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 / "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 / "external_callout_config.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 / "external_callout_config.json" + cred_file.write_text(json.dumps({"credentials": {"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}' + # No Content-Type is assumed; headers are passed through verbatim. + assert "Content-Type" not in captured["headers"] + assert isinstance(captured["auth"], DynamicAuthHandler) + + assert result["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(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(NamedCredentialCallError): + 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 + + 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: ClassVar[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_run.py b/tests/test_run.py index 1eace88..154b0bf 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -488,3 +488,90 @@ 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 TestReadStreamingSource: + """`_read_streaming_source` extracts the source name from config.json's + `streamingSource` object.""" + + def test_returns_dlo_name(self): + from datacustomcode.run import _read_streaming_source + + config_json = {"streamingSource": {"type": "dlo", "name": "Account_Home__dll"}} + assert _read_streaming_source(config_json) == "Account_Home__dll" + + def test_returns_dmo_name(self): + from datacustomcode.run import _read_streaming_source + + config_json = { + "streamingSource": {"type": "dmo", "name": "AccountTransformed__dlm"} + } + assert _read_streaming_source(config_json) == "AccountTransformed__dlm" + + @pytest.mark.parametrize( + "config_json", + [ + {}, + {"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_streaming_source + + assert _read_streaming_source(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_streaming_source_field(self): + content = self._run_capturing_streaming_source( + { + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "Account_Home__dll"}, + } + ) + assert "streaming_source: Account_Home__dll" in content + + def test_streaming_source_none_for_batch_without_field(self): + content = self._run_capturing_streaming_source({"dataspace": "default"}) + assert "streaming_source: None" in content diff --git a/tests/test_runtime_named_credential.py b/tests/test_runtime_named_credential.py new file mode 100644 index 0000000..78e9f5a --- /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, body=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.body == "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" 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."""