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 3452d04..a9f5ae8 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..b43e423 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -15,13 +15,16 @@ from __future__ import annotations from enum import Enum +import os from typing import ( TYPE_CHECKING, Any, ClassVar, Dict, Optional, + TypeVar, Union, + cast, ) from datacustomcode.config import config @@ -29,21 +32,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 +145,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 +164,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 +222,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 +300,26 @@ def __str__(self) -> str: return msg -class Client: - """Entrypoint for accessing DataCloud objects. +_ClientT = TypeVar("_ClientT", bound="_BaseClient") + - 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. +class _BaseClient: + """Shared machinery for the Data Cloud client singletons. - 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. + Holds the wiring common to :class:`Client` (batch) and + :class:`StreamingClient` + + 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 +329,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 +413,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 +538,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 +586,170 @@ 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(self) -> PySparkDataFrame: + """Read the streamingSource + + Returns: + A standard PySpark DataFrame from the streaming source DLO + """ + self._record_dlo_access(_streaming_source_name()) + return self._reader.read_dlo(_streaming_source_name()) + + 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(self) -> PySparkDataFrame: + """Read the streamingSource + + Returns a standard PySpark DataFrame from the streaming source DMO + """ + self._record_dmo_access(_streaming_source_name()) + return self._reader.read_dmo(_streaming_source_name()) + + 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] + + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write a PySpark DataFrame to a DLO in Data Cloud automatically picking + the WriteMode. + For use with streaming transforms when running in rebuild or initial sync mode. + Args: + name: The name of the DLO to write to. + dataframe: The PySpark DataFrame to write. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) + return self._writer.auto_write_to_dlo(name, dataframe) + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write a PySpark DataFrame to a DMO in Data Cloud automatically picking + the WriteMode. + For use with streaming transforms when running in rebuild or initial sync mode. + Args: + name: The name of the DMO to write to. + dataframe: The PySpark DataFrame to write. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) + return self._writer.auto_write_to_dmo(name, dataframe) + + +class RunMode(Enum): + BATCH = "BATCH" + INITIAL_SYNC = "INITIAL_SYNC" + REBUILD = "REBUILD" + DELTA_SYNC = "DELTA_SYNC" + + +def get_run_mode() -> RunMode: + """Read and validate the BYOC_RUN_MODE env var; default to BATCH when unset.""" + run_mode = os.getenv("BYOC_RUN_MODE", "BATCH").upper() + try: + return RunMode(run_mode) + except ValueError as exc: + raise ValueError("Set BYOC_RUN_MODE to a valid value") from exc 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..d24a33b 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,47 @@ def write_to_dlo( def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: ... + + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write to a DLO automatically picking the write mode. + For use with streaming transforms when running in rebuild or initial sync mode. + """ + raise NotImplementedError + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + """Write to a DMO automatically picking the write mode. + For use with streaming transforms when running in rebuild or initial sync mode. + """ + raise NotImplementedError + + 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/io/writer/csv.py b/src/datacustomcode/io/writer/csv.py index 3d037d9..292d92d 100644 --- a/src/datacustomcode/io/writer/csv.py +++ b/src/datacustomcode/io/writer/csv.py @@ -36,6 +36,10 @@ def write_to_dlo( name = f"{name}{SUFFIX}" dataframe.write.csv(name, mode=write_mode) + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + # use overwrite since this is a local only writer + self.write_to_dlo(name, dataframe, WriteMode.OVERWRITE) + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: @@ -43,3 +47,7 @@ def write_to_dmo( if not name.lower().endswith(SUFFIX): name = f"{name}{SUFFIX}" dataframe.write.csv(name, mode=write_mode) + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + # use overwrite since this is a local only writer + self.write_to_dmo(name, dataframe, WriteMode.OVERWRITE) diff --git a/src/datacustomcode/io/writer/print.py b/src/datacustomcode/io/writer/print.py index c4d2a75..19ef117 100644 --- a/src/datacustomcode/io/writer/print.py +++ b/src/datacustomcode/io/writer/print.py @@ -122,6 +122,10 @@ def write_to_dlo( dataframe.show() + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + self.validate_dataframe_columns_against_dlo(dataframe, name) + dataframe.show() + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: @@ -130,3 +134,6 @@ def write_to_dmo( # so just show the dataframe. dataframe.show() + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + dataframe.show() diff --git a/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..49bb699 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/auth.py @@ -0,0 +1,191 @@ +# 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 +import datetime +import hashlib +import hmac +from typing import ( + TYPE_CHECKING, + Any, + Dict, +) +import urllib.parse + +from requests.auth import AuthBase + +from datacustomcode.named_credential.direct.credentials import AuthType + +if TYPE_CHECKING: + from requests.models import PreparedRequest + +_SIGV4_ALGORITHM = "AWS4-HMAC-SHA256" + + +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}" + + elif self.auth_type == AuthType.AWS_SIG_V4.value: + self._sign_aws_sigv4(request) + + else: + raise ValueError(f"Unsupported auth_type '{self.auth_type}'.") + + return request + + def _sign_aws_sigv4(self, request: PreparedRequest) -> None: + """Sign ``request`` with AWS Signature Version 4. + + Requires ``aws_access_key_id``, ``aws_secret_access_key``, ``aws_region``, + and ``aws_service`` in the credential config; ``aws_session_token`` is + optional (for temporary credentials). The signed date, payload hash, and + (when present) session token are added as ``x-amz-*`` headers so the sent + request matches what was signed. + """ + access_key = self.config.get("aws_access_key_id") + secret_key = self.config.get("aws_secret_access_key") + region = self.config.get("aws_region") + service = self.config.get("aws_service") + session_token = self.config.get("aws_session_token") + missing = [ + name + for name, value in ( + ("aws_access_key_id", access_key), + ("aws_secret_access_key", secret_key), + ("aws_region", region), + ("aws_service", service), + ) + if not value + ] + if missing: + raise ValueError(f"'{self.auth_type}' auth requires {', '.join(missing)}.") + access_key = str(access_key) + secret_key = str(secret_key) + region = str(region) + service = str(service) + + parsed = urllib.parse.urlsplit(str(request.url or "")) + host = parsed.netloc + canonical_uri = urllib.parse.quote(parsed.path or "/", safe="/-_.~") + canonical_query = _canonical_query_string(parsed.query) + + body = request.body or b"" + if isinstance(body, str): + body = body.encode("utf-8") + payload_hash = hashlib.sha256(body).hexdigest() + + now = datetime.datetime.now(datetime.timezone.utc) + amz_date = now.strftime("%Y%m%dT%H%M%SZ") + datestamp = now.strftime("%Y%m%d") + + request.headers["x-amz-date"] = amz_date + request.headers["x-amz-content-sha256"] = payload_hash + if session_token: + request.headers["x-amz-security-token"] = session_token + + signed = { + "host": host, + "x-amz-content-sha256": payload_hash, + "x-amz-date": amz_date, + } + if session_token: + signed["x-amz-security-token"] = session_token + signed_headers = ";".join(sorted(signed)) + canonical_headers = "".join( + f"{name}:{signed[name]}\n" for name in sorted(signed) + ) + + canonical_request = "\n".join( + [ + request.method or "GET", + canonical_uri, + canonical_query, + canonical_headers, + signed_headers, + payload_hash, + ] + ) + credential_scope = f"{datestamp}/{region}/{service}/aws4_request" + string_to_sign = "\n".join( + [ + _SIGV4_ALGORITHM, + amz_date, + credential_scope, + hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(), + ] + ) + signing_key = _derive_signing_key(secret_key, datestamp, region, service) + signature = hmac.new( + signing_key, string_to_sign.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + request.headers["Authorization"] = ( + f"{_SIGV4_ALGORITHM} Credential={access_key}/{credential_scope}, " + f"SignedHeaders={signed_headers}, Signature={signature}" + ) + + +def _canonical_query_string(query: str) -> str: + """Build the AWS Sig V4 canonical query string from a raw query string.""" + pairs = urllib.parse.parse_qsl(query, keep_blank_values=True) + encoded = [ + ( + urllib.parse.quote(key, safe="-_.~"), + urllib.parse.quote(value, safe="-_.~"), + ) + for key, value in pairs + ] + encoded.sort() + return "&".join(f"{key}={value}" for key, value in encoded) + + +def _derive_signing_key( + secret_key: str, datestamp: str, region: str, service: str +) -> bytes: + """Derive the AWS Sig V4 signing key via the chained HMAC-SHA256 sequence.""" + + def _hmac(key: bytes, msg: str) -> bytes: + return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest() + + k_date = _hmac(f"AWS4{secret_key}".encode(), datestamp) + k_region = _hmac(k_date, region) + k_service = _hmac(k_region, service) + return _hmac(k_service, "aws4_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..03a73c9 --- /dev/null +++ b/src/datacustomcode/named_credential/direct/credentials.py @@ -0,0 +1,122 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""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.""" + + AWS_SIG_V4 = "AwsSv4" + 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..cadd1e1 --- /dev/null +++ b/src/datacustomcode/named_credential/types/http_request_builder.py @@ -0,0 +1,88 @@ +# 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, + Optional, + Union, +) + +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request import HTTPRequest + +# Control header carrying the per-request callout response timeout (seconds). +CALLOUT_RESPONSE_TIMEOUT_HEADER = "ctx-callout-response-timeout-seconds" + + +class HTTPRequestBuilder: + def __init__(self) -> None: + self._url = "" + self._method: Union[str, HTTPMethod] = HTTPMethod.GET + self._headers: Dict[str, str] = {} + self._response_timeout_seconds: Optional[int] = None + + 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": + """Set the HTTP headers. + + Args: + headers: a dictionary of HTTP headers + """ + self._headers = headers + return self + + def set_response_timeout_seconds(self, seconds: int) -> "HTTPRequestBuilder": + """Set the per-request callout response timeout, in seconds. + + The ``ctx-callout-response-timeout-seconds`` control header allows to + configure the timeout for the callout response. This is optional. + If not set, the platform will use the default timeout. + """ + if isinstance(seconds, bool) or not isinstance(seconds, int) or seconds <= 0: + raise ValueError( + "response timeout must be a positive integer number of seconds, " + f"got {seconds!r}." + ) + self._response_timeout_seconds = seconds + return self + + def build(self) -> HTTPRequest: + headers = dict(self._headers) + if self._response_timeout_seconds is not None: + headers[CALLOUT_RESPONSE_TIMEOUT_HEADER] = str( + self._response_timeout_seconds + ) + return HTTPRequest( + url=self._url, + method=self._method, # type: ignore[arg-type] + headers=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/script/examples/streaming_deltas/entrypoint.py b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py new file mode 100644 index 0000000..abe7981 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py @@ -0,0 +1,62 @@ +"""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. + +The first run of a streaming job will use the run mode INITIAL_SYNC which behaves +like a batch run on the streaming source. A streaming transform can also use run +mode REBUILD to do the same thing on demand. Note that these will process all +source rows and overwrite the target. + +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 runtime; + the local ``datacustomcode run`` readers/writers raise +``NotImplementedError`` for the delta methods. +""" + +from pyspark.sql import DataFrame +from pyspark.sql.functions import col, upper + +from datacustomcode.client import ( + RunMode, + StreamingClient, + get_run_mode, +) + + +def main(): + target_dlo = "Account_std_copy__dll" + client = StreamingClient() + + if get_run_mode() == RunMode.DELTA_SYNC: + # Streaming DataFrame over the source DLO's change feed. + dataframe = client.read_dlo_deltas() + # Ordinary PySpark transform. + transformed = transform(dataframe) + + # Start the streaming write. write_dlo_deltas returns the StreamingQuery; + # the trigger and checkpoint location are provided by the runtime. + query = client.write_dlo_deltas(target_dlo, transformed) + + # Drive the query's lifecycle. In the streaming runtime this blocks until + # the job is stopped by the platform. + query.awaitTermination() + else: + # initial sync and rebuild read the entire streaming source DLO and + # write using a server-decided mode based on the run mode + dataframe = client.read_dlo() + transformed = transform(dataframe) + client.auto_write_to_dlo(target_dlo, transformed) + + +def transform(dataframe: DataFrame) -> DataFrame: + return dataframe.withColumn("description__c", upper(col("description__c"))) + + +if __name__ == "__main__": + 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/spark/test_session_provider.py b/tests/spark/test_session_provider.py index 71f0e70..b0c7d90 100644 --- a/tests/spark/test_session_provider.py +++ b/tests/spark/test_session_provider.py @@ -52,11 +52,17 @@ def write_to_dlo( ) -> None: # type: ignore[override] raise NotImplementedError + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + raise NotImplementedError + def write_to_dmo( self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode ) -> None: # type: ignore[override] raise NotImplementedError + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + raise NotImplementedError + class FakeProvider(BaseSparkSessionProvider): CONFIG_NAME = "FakeProvider" diff --git a/tests/test_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..4f3f2f4 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from unittest.mock import MagicMock, patch from pyspark.sql import DataFrame, SparkSession @@ -9,8 +10,13 @@ Client, DataCloudAccessLayerException, DataCloudObjectType, + RunMode, + StreamingClient, + _BaseClient, einstein_predict_col, + get_run_mode, llm_gateway_generate_text_col, + named_credential_request_col, ) from datacustomcode.config import ( AccessLayerObjectConfig, @@ -46,11 +52,17 @@ def write_to_dlo( ) -> None: pass + def auto_write_to_dlo(self, name: str, dataframe: DataFrame) -> None: + pass + def write_to_dmo( self, name: str, dataframe: DataFrame, write_mode: WriteMode, **kwargs ) -> None: pass + def auto_write_to_dmo(self, name: str, dataframe: DataFrame) -> None: + pass + @pytest.fixture def mock_spark(): @@ -80,10 +92,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: @@ -255,8 +275,272 @@ def test_read_pattern_flow(self, reset_client, mock_spark): assert "source_dmo" in client._data_layer_history[DataCloudObjectType.DMO] + @patch.dict(os.environ, {}, clear=True) + def test_get_run_mode_default_batch(self, reset_client, mock_spark): + + assert get_run_mode() == RunMode.BATCH + + @patch.dict(os.environ, {"BYOC_RUN_MODE": "INITIAL_SYNC"}) + def test_get_run_mode(self, reset_client, mock_spark): + + assert get_run_mode() == RunMode.INITIAL_SYNC + + @patch.dict(os.environ, {"BYOC_RUN_MODE": "INVALID"}) + def test_get_run_mode_throws(self): + with pytest.raises(ValueError, match="Set BYOC_RUN_MODE to a valid value"): + get_run_mode() + + +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] + + def test_auto_write_to_dlo(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = StreamingClient(reader=reader, writer=writer) + client._record_dlo_access("some_dlo") + + client.auto_write_to_dlo("test_dlo", mock_df) + + writer.auto_write_to_dlo.assert_called_once_with("test_dlo", mock_df) -class TestClientLlmGatewayGenerateText: + def test_auto_write_to_dmo(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = StreamingClient(reader=reader, writer=writer) + client._record_dmo_access("some_dmo") + + client.auto_write_to_dmo("test_dmo", mock_df) + + writer.auto_write_to_dmo.assert_called_once_with("test_dmo", mock_df) + + +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 +710,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..d486a74 --- /dev/null +++ b/tests/test_named_credential.py @@ -0,0 +1,496 @@ +# 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() + + def test_builder_sets_response_timeout_header(self): + request = ( + HTTPRequestBuilder() + .set_url("callout:NC/path") + .set_headers({"Accept": "application/json"}) + .set_response_timeout_seconds(60) + .build() + ) + assert request.headers == { + "Accept": "application/json", + "ctx-callout-response-timeout-seconds": "60", + } + + def test_builder_response_timeout_is_order_independent(self): + # Setting the timeout before set_headers must not be wiped out by the + # header map replacement — build() applies it on top either way. + request = ( + HTTPRequestBuilder() + .set_url("callout:NC/path") + .set_response_timeout_seconds(60) + .set_headers({"Accept": "application/json"}) + .build() + ) + assert request.headers == { + "Accept": "application/json", + "ctx-callout-response-timeout-seconds": "60", + } + + @pytest.mark.parametrize("bad", [0, -5, True, 1.5, "30", None]) + def test_builder_rejects_invalid_response_timeout(self, bad): + with pytest.raises(ValueError): + HTTPRequestBuilder().set_url( + "callout:NC/path" + ).set_response_timeout_seconds(bad) + + +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..1b8a39f --- /dev/null +++ b/tests/test_named_credential_direct.py @@ -0,0 +1,456 @@ +# 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()) + + +def _reference_sigv4(config, method, url, body): + """Independent SigV4 reference used to validate DynamicAuthHandler output.""" + import datetime + import hashlib + import hmac + import urllib.parse + + parsed = urllib.parse.urlsplit(url) + payload = body.encode("utf-8") if isinstance(body, str) else (body or b"") + payload_hash = hashlib.sha256(payload).hexdigest() + now = datetime.datetime(2024, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) + amz_date = now.strftime("%Y%m%dT%H%M%SZ") + datestamp = now.strftime("%Y%m%d") + + signed = { + "host": parsed.netloc, + "x-amz-content-sha256": payload_hash, + "x-amz-date": amz_date, + } + token = config.get("aws_session_token") + if token: + signed["x-amz-security-token"] = token + signed_headers = ";".join(sorted(signed)) + canonical_headers = "".join(f"{n}:{signed[n]}\n" for n in sorted(signed)) + + pairs = sorted( + (urllib.parse.quote(k, safe="-_.~"), urllib.parse.quote(v, safe="-_.~")) + for k, v in urllib.parse.parse_qsl(parsed.query, keep_blank_values=True) + ) + canonical_query = "&".join(f"{k}={v}" for k, v in pairs) + canonical_uri = urllib.parse.quote(parsed.path or "/", safe="/-_.~") + canonical_request = "\n".join( + [ + method, + canonical_uri, + canonical_query, + canonical_headers, + signed_headers, + payload_hash, + ] + ) + scope = f"{datestamp}/{config['aws_region']}/{config['aws_service']}/aws4_request" + string_to_sign = "\n".join( + [ + "AWS4-HMAC-SHA256", + amz_date, + scope, + hashlib.sha256(canonical_request.encode()).hexdigest(), + ] + ) + + def _h(key, msg): + return hmac.new(key, msg.encode(), hashlib.sha256).digest() + + k = _h(f"AWS4{config['aws_secret_access_key']}".encode(), datestamp) + k = _h(k, config["aws_region"]) + k = _h(k, config["aws_service"]) + k = _h(k, "aws4_request") + signature = hmac.new(k, string_to_sign.encode(), hashlib.sha256).hexdigest() + return ( + ( + f"AWS4-HMAC-SHA256 Credential={config['aws_access_key_id']}/{scope}, " + f"SignedHeaders={signed_headers}, Signature={signature}" + ), + amz_date, + payload_hash, + ) + + +class TestAwsSigV4Auth: + _CONFIG: ClassVar[dict] = { + "auth_type": AuthType.AWS_SIG_V4.value, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + "aws_region": "us-east-1", + "aws_service": "s3", + } + + @staticmethod + def _freeze_clock(monkeypatch): + import datetime as _dt + + from datacustomcode.named_credential.direct import auth as auth_mod + + frozen = _dt.datetime(2024, 1, 2, 3, 4, 5, tzinfo=_dt.timezone.utc) + + class _FrozenDatetime(_dt.datetime): + @classmethod + def now(cls, tz=None): + return frozen if tz is None else frozen.astimezone(tz) + + monkeypatch.setattr(auth_mod.datetime, "datetime", _FrozenDatetime) + + def _sign(self, monkeypatch, config, method="GET", url=None, body=None): + self._freeze_clock(monkeypatch) + url = url or "https://bucket.s3.amazonaws.com/key" + request = PreparedRequest() + request.prepare(method=method, url=url, data=body) + return DynamicAuthHandler(config)(request) + + def test_matches_reference_signature(self, monkeypatch): + url = "https://bucket.s3.amazonaws.com/some/key?b=2&a=1" + request = self._sign(monkeypatch, self._CONFIG, "PUT", url, "payload") + expected, amz_date, payload_hash = _reference_sigv4( + self._CONFIG, "PUT", url, "payload" + ) + assert request.headers["Authorization"] == expected + assert request.headers["x-amz-date"] == amz_date + assert request.headers["x-amz-content-sha256"] == payload_hash + assert "x-amz-security-token" not in request.headers + + def test_empty_body_hashes_empty_string(self, monkeypatch): + request = self._sign(monkeypatch, self._CONFIG) + assert ( + request.headers["x-amz-content-sha256"] + == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ) + + def test_session_token_is_signed(self, monkeypatch): + config = {**self._CONFIG, "aws_session_token": "SESSIONTOKEN"} + request = self._sign(monkeypatch, config) + assert request.headers["x-amz-security-token"] == "SESSIONTOKEN" + assert "x-amz-security-token" in request.headers["Authorization"] + expected, _, _ = _reference_sigv4(config, "GET", request.url, None) + assert request.headers["Authorization"] == expected + + @pytest.mark.parametrize( + "missing", + ["aws_access_key_id", "aws_secret_access_key", "aws_region", "aws_service"], + ) + def test_missing_field_raises(self, monkeypatch, missing): + config = {k: v for k, v in self._CONFIG.items() if k != missing} + with pytest.raises(ValueError, match=missing): + self._sign(monkeypatch, config) + + +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."""