diff --git a/CHANGELOG.md b/CHANGELOG.md index fb41eee..5d927af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## 6.2.0 + +### Added + +- **`StreamingClient.write_dmo_deltas(name, dataframe)`** – start a streaming query that writes each micro-batch to the target DMO and return the `StreamingQuery` handle. It is the DMO counterpart to `write_dlo_deltas` and pairs with `read_dmo_deltas()`. + ## 6.1.0 ### Added diff --git a/README.md b/README.md index 1f90472..1381c6c 100644 --- a/README.md +++ b/README.md @@ -158,6 +158,7 @@ For a streaming (delta) transform, use `StreamingClient`, which exposes the stre * `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` +* `write_dmo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Model Object; returns the started `StreamingQuery` `find_file_path`, `llm_gateway_generate_text`, and `einstein_predict` are available on both clients. @@ -222,6 +223,15 @@ query = client.write_dlo_deltas("Output__dll", transformed) query.awaitTermination() ``` +A streaming transform can also read and write Data Model Objects with `read_dmo_deltas()` and `write_dmo_deltas()`: + +```python +deltas = client.read_dmo_deltas() +transformed = deltas.withColumn("description__c", upper(col("description__c"))) +query = client.write_dmo_deltas("Output__dlm", 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. diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 49d1407..69f7ea1 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -726,6 +726,25 @@ def write_dlo_deltas( 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 write_dmo_deltas( + self, name: str, dataframe: PySparkDataFrame, **kwargs + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a DMO in Data Cloud. + + Starts a streaming query that writes each micro-batch to the + target DMO and returns the ``StreamingQuery`` handle; the caller + typically calls ``query.awaitTermination()``. + + Args: + name: The name of the DMO to write to. + dataframe: The streaming PySpark DataFrame to write. + + Returns: + The started ``StreamingQuery``. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) + return self._writer.write_dmo_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. diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index 881be26..b48ab7f 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -415,13 +415,7 @@ def is_streaming(self) -> bool: 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: + if read_is_dlo != write_is_dlo: raise ValueError( "permissions.read and permissions.write must both reference " "DLOs or both reference DMOs (got " diff --git a/src/datacustomcode/io/writer/base.py b/src/datacustomcode/io/writer/base.py index d24a33b..56b89f9 100644 --- a/src/datacustomcode/io/writer/base.py +++ b/src/datacustomcode/io/writer/base.py @@ -87,7 +87,7 @@ def write_dlo_deltas( Args: name: Target Data Lake Object name. dataframe: Streaming PySpark DataFrame produced from a - ``read_dlo_deltas`` / ``read_dmo_deltas`` source. + ``read_dlo_deltas`` source. Returns: The started ``StreamingQuery``; the caller drives its lifecycle @@ -102,3 +102,35 @@ def write_dlo_deltas( "streaming runtime; the local writer does not support streaming " "deltas." ) + + def write_dmo_deltas( + self, name: str, dataframe: PySparkDataFrame + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a Data Model Object. + + Streaming counterpart to :meth:`write_to_dmo`. Starts a streaming query + that writes each micro-batch to the target DMO 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 Model Object name. + dataframe: Streaming PySpark DataFrame produced from a + ``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_dmo_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/scan.py b/src/datacustomcode/scan.py index 059925a..a3bf718 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -34,7 +34,7 @@ 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"] +STREAMING_WRITE_METHODS = ["write_dlo_deltas", "write_dmo_deltas"] DATA_TRANSFORM_CONFIG_TEMPLATE = { "sdkVersion": get_version(), @@ -181,6 +181,7 @@ class StreamingDataAccessLayerCalls(pydantic.BaseModel): read_dlo_deltas: bool read_dmo_deltas: bool write_dlo_deltas: frozenset[str] + write_dmo_deltas: frozenset[str] @pydantic.model_validator(mode="after") def validate_access_layer(self) -> StreamingDataAccessLayerCalls: @@ -193,10 +194,19 @@ def validate_access_layer(self) -> StreamingDataAccessLayerCalls: "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: + if not self.write_dlo_deltas and not self.write_dmo_deltas: raise ValueError( - "A streaming transform must write to at least one DLO via " - "write_dlo_deltas." + "A streaming transform must write to at least one DLO or DMO " + "via write_dlo_deltas / write_dmo_deltas." + ) + if self.write_dlo_deltas and self.write_dmo_deltas: + raise ValueError( + "Cannot write DLO and DMO deltas in the same streaming transform." + ) + if self.read_layer != self.write_layer: + raise ValueError( + "A streaming transform must read and write the same layer " + f"(got read={self.read_layer}, write={self.write_layer})." ) return self @@ -205,6 +215,11 @@ def read_layer(self) -> str: """Return the read source layer, ``"dlo"`` or ``"dmo"``.""" return "dlo" if self.read_dlo_deltas else "dmo" + @property + def write_layer(self) -> str: + """Return the write target layer, ``"dlo"`` or ``"dmo"``.""" + return "dlo" if self.write_dlo_deltas else "dmo" + class ClientMethodVisitor(ast.NodeVisitor): """AST Visitor that finds all instances of Client read/write method calls.""" @@ -217,6 +232,7 @@ def __init__(self) -> None: self._read_dlo_deltas: bool = False self._read_dmo_deltas: bool = False self._write_dlo_deltas_instances: set[str] = set() + self._write_dmo_deltas_instances: set[str] = set() self.variable_values: Dict[str, Union[str, None]] = {} def visit_Assign(self, node: ast.Assign) -> None: @@ -261,6 +277,8 @@ def visit_Call(self, node: ast.Call) -> None: name = self._resolve_name_arg(node.args[0]) if name and method_name == "write_dlo_deltas": self._write_dlo_deltas_instances.add(name) + elif name and method_name == "write_dmo_deltas": + self._write_dmo_deltas_instances.add(name) self.generic_visit(node) def _resolve_name_arg(self, arg: ast.expr) -> Union[str, None]: @@ -278,6 +296,7 @@ def is_streaming(self) -> bool: self._read_dlo_deltas or self._read_dmo_deltas or bool(self._write_dlo_deltas_instances) + or bool(self._write_dmo_deltas_instances) ) def found(self) -> DataAccessLayerCalls: @@ -293,6 +312,7 @@ def found_streaming(self) -> StreamingDataAccessLayerCalls: read_dlo_deltas=self._read_dlo_deltas, read_dmo_deltas=self._read_dmo_deltas, write_dlo_deltas=frozenset(self._write_dlo_deltas_instances), + write_dmo_deltas=frozenset(self._write_dmo_deltas_instances), ) @@ -516,10 +536,10 @@ def _update_streaming_config(existing_config: dict[str, Any], file_path: str) -> ) read_names = [source_name] if source_name else [] - write_names = list(output.write_dlo_deltas) + write_names = list(output.write_dlo_deltas or output.write_dmo_deltas) existing_config["permissions"] = { "read": {read_layer: read_names}, - "write": {"dlo": write_names}, + "write": {output.write_layer: write_names}, } diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md b/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md new file mode 100644 index 0000000..205c004 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/README.md @@ -0,0 +1,120 @@ +# Chunking with a Gemini Named Credential Callout + +Splits each input document into paragraph-sized chunks and calls Google's +**Gemini** `generateContent` API for every chunk. The model returns a summary, +category, sentiment, and topics, which are attached to the chunk as citations so +the search index can filter and rank on them. Gemini is reached through a +**Named Credential**, so this code never handles the endpoint URL or the API key. + +## How the callout works + +```python +CALLOUT_URL = "callout:gemini" # callout:[/] + +request = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json"}) + .set_response_timeout_seconds(60) # optional: per-callout timeout (seconds) + .build() +) +# Body is sent verbatim (serialize it yourself); the response body is a raw string. +response = runtime.named_credential.request(request, json.dumps(payload)) +if response.is_success: + envelope = json.loads(response.body) + text = envelope["candidates"][0]["content"]["parts"][0]["text"] +``` + +The request asks for `responseMimeType: application/json` with a `responseSchema`, +so Gemini returns the classification as a JSON string in +`candidates[0].content.parts[0].text` — decode it, then decode that text again. + +The `gemini` Named Credential's URL already includes the full +`/v1beta/models/:generateContent` path, so the callout is just +`callout:gemini` with **no path suffix** (anything after the name is appended to +the credential's URL). + +## Configure the Named Credential + +1. Create an **External Credential** (e.g. `google_api_key`) that injects your + Gemini API key as the `X-goog-api-key` header. +2. Create a **Named Credential** named `gemini`: + - **URL**: `https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent` + - **Enabled for Callouts** + **Generate Authorization Header**: on + - **External Credential**: `google_api_key` + +## Test locally + +```bash +DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG=/abs/path/to/external_callout_config.json \ + sf data-code-extension function run \ + --entrypoint payload/entrypoint.py \ + --test-with payload/tests/test.json \ + --target-org +``` + +With `--target-org` the SDK fetches only the **URL** from the org's Named +Credential; **auth is always taken from `external_callout_config.json`** locally +(the org's External Credential is used only in the Data Cloud runtime). So the +`X-goog-api-key` must be in the local config for a local test. Omit +`--target-org` to run fully offline using `target_url`. + +```json +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_GEMINI_API_KEY" }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} +``` + +Place `external_callout_config.json` in the **parent of your payload folder** (or +point `DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG` at it). It is never packaged into +the deployment zip. Get a key from [Google AI Studio](https://aistudio.google.com/apikey); +**do not commit it.** + +## Auth types + +`auth_type` selects how auth is injected for local testing. It should mirror the +External Credential your Named Credential uses in the org, so local and deployed +runs behave the same. This example uses `Custom` (Gemini's `X-goog-api-key`); +all four supported types: + +```json +{ + "credentials": { + "callout:my_custom_api": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_API_KEY" } + }, + "callout:my_basic_api": { + "auth_type": "Basic", + "username": "svc_user", + "password": "YOUR_PASSWORD" + }, + "callout:my_oauth_api": { + "auth_type": "OAuth", + "access_token": "YOUR_ACCESS_TOKEN" + }, + "callout:my_jwt_api": { + "auth_type": "Jwt", + "token": "YOUR_JWT" + } + } +} +``` + +| `auth_type` | Fields read | Header sent | +| ----------- | ------------------------------- | --------------------------------------- | +| `Basic` | `username`, `password` | `Authorization: Basic ` | +| `Custom` | `custom_headers` (sent verbatim)| the headers you list | +| `OAuth` | `access_token` or `token` | `Authorization: Bearer ` | +| `Jwt` | `access_token` or `token` | `Authorization: Bearer ` | + +`OAuth`/`Jwt` take a token you supply for the local run — the SDK does not fetch +or refresh it. In the Data Cloud runtime the Named Credential handles token +acquisition; this local config only stands in for that during testing. diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json new file mode 100644 index 0000000..2f911e9 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/config.json @@ -0,0 +1,3 @@ +{ + "entryPoint": "entrypoint.py" +} diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py b/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py new file mode 100644 index 0000000..ab1b00c --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/entrypoint.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 + +""" +Document Chunking with a Gemini Named Credential Callout + +Splits each input document into paragraph-sized chunks and classifies every +chunk via Google's Gemini ``generateContent`` API, reached through a Named +Credential (``callout:gemini``) so the endpoint URL and API key are resolved +outside this code. The classification is attached to each chunk as citations. +""" + +import json +import logging + +from datacustomcode.function import Runtime +from datacustomcode.function.feature_types.chunking import ( + ChunkType, + SearchIndexChunkingV1Output, + SearchIndexChunkingV1Request, + SearchIndexChunkingV1Response, +) +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request_builder import ( + HTTPRequestBuilder, +) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +CALLOUT_URL = "callout:gemini" + +_ANALYSIS_FIELDS = ("summary", "category", "sentiment") + +_PROMPT = ( + "Analyze the following document chunk and classify it. Respond with its " + "one-sentence summary, a single-word category, overall sentiment " + "(positive, negative, or neutral), and up to five key topics.\n\nChunk:\n" +) + +# Force Gemini to return the classification as JSON in a fixed shape. +_GENERATION_CONFIG = { + "responseMimeType": "application/json", + "responseSchema": { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "category": {"type": "string"}, + "sentiment": {"type": "string"}, + "topics": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["summary", "category", "sentiment", "topics"], + }, +} + + +def _chunk_text(text: str, max_words: int = 80) -> list[str]: + """Split text into paragraph-aligned chunks of at most ``max_words`` words.""" + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + + chunks: list[str] = [] + current: list[str] = [] + current_words = 0 + + for paragraph in paragraphs: + paragraph_words = len(paragraph.split()) + if current and current_words + paragraph_words > max_words: + chunks.append("\n\n".join(current)) + current = [] + current_words = 0 + current.append(paragraph) + current_words += paragraph_words + + if current: + chunks.append("\n\n".join(current)) + + return chunks + + +def _extract_model_json(body: str) -> dict: + """Decode the model's JSON classification from a Gemini response. + + The generated text sits at ``candidates[0].content.parts[0].text`` and is + itself a JSON string, so decode twice. Any malformed layer yields ``{}``. + """ + try: + envelope = json.loads(body) if body else {} + except json.JSONDecodeError: + return {} + + try: + text = envelope["candidates"][0]["content"]["parts"][0]["text"] + except (KeyError, IndexError, TypeError): + return {} + + try: + payload = json.loads(text) + except json.JSONDecodeError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _analyze_chunk(chunk_text: str, runtime: Runtime) -> dict[str, str]: + """Classify one chunk via the Gemini callout and return it as citations.""" + request = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json", "Accept": "application/json"}) + .set_response_timeout_seconds(60) + .build() + ) + + payload = { + "contents": [{"parts": [{"text": _PROMPT + chunk_text}]}], + "generationConfig": _GENERATION_CONFIG, + } + response = runtime.named_credential.request(request, json.dumps(payload)) + + # Don't raise: a single failed callout shouldn't abort the whole job. + if not response.is_success: + logger.error(f"Gemini callout failed with status {response.status_code}") + return {"analysis_status": "failed", "http_status": str(response.status_code)} + + data = _extract_model_json(response.body) + citations = {"analysis_status": "success"} + for field in _ANALYSIS_FIELDS: + value = data.get(field) + citations[field] = str(value) if value is not None else "unavailable" + + topics = data.get("topics") + if isinstance(topics, list): + citations["topics"] = ", ".join(str(topic) for topic in topics) + + return citations + + +def function( + request: SearchIndexChunkingV1Request, runtime: Runtime +) -> SearchIndexChunkingV1Response: + """Chunk each input document and classify every chunk via the Gemini API.""" + logger.info(f"Received {len(request.input)} documents to chunk") + + chunks = [] + chunk_id = 1 + + for doc in request.input: + for chunk_text in _chunk_text(doc.text): + citations = _analyze_chunk(chunk_text, runtime) + + chunk = SearchIndexChunkingV1Output( + text=chunk_text, + seq_no=chunk_id, + chunk_type=ChunkType.TEXT, + citations=citations, + ) + chunks.append(chunk) + chunk_id += 1 + + logger.info(f"Produced {len(chunks)} classified chunks") + return SearchIndexChunkingV1Response(output=chunks) diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json new file mode 100644 index 0000000..2f979ea --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/external_callout_config.json @@ -0,0 +1,11 @@ +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { + "X-goog-api-key": "YOUR_API_KEY" + }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} diff --git a/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json b/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json new file mode 100644 index 0000000..c4d9960 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_external_callout/tests/test.json @@ -0,0 +1,16 @@ +{ + "input": [ + { + "text": "Product Review: Northstar Analytics\n\nWe rolled Northstar out to our whole revenue team last quarter and the difference has been night and day. Dashboards that used to take our analysts a full day to assemble now refresh in seconds, and the natural-language query box means our account executives can answer their own questions without filing a ticket.\n\nOnboarding was smoother than any tool we have adopted in years. The guided setup imported our Salesforce data on the first try and the sample templates gave us something useful on day one. Support answered our two questions within the hour. Easily the best purchase decision we made this year." + }, + { + "text": "Support Ticket #48210: Repeated timeouts on scheduled exports\n\nFor the third week running our nightly export to the data warehouse has failed silently. There is no alert, no email, nothing in the activity log, and we only find out when the morning report is empty and the leadership meeting has no numbers.\n\nI have raised this twice already and both times the ticket was closed as resolved without anyone actually contacting me. This is costing us real credibility internally and I am extremely frustrated. If the connector cannot handle our volume we need to know now so we can plan a migration, because right now the product is not doing the one job we bought it for." + }, + { + "text": "Renewal Feedback: mixed feelings heading into year two\n\nThe core product is genuinely good. The reporting engine is fast, the permissions model is granular enough for our compliance team, and our analysts like working in it. On the functionality alone I would renew without hesitation.\n\nWhat gives me pause is the pricing. The per-seat cost jumped noticeably at renewal and several add-ons that used to be included are now separate line items. The value is still there, but the conversation with my finance team was harder than it should have been, and I would like more transparency before the next cycle." + }, + { + "text": "Feature Request: scheduled report subscriptions\n\nWe would like the ability to subscribe internal stakeholders to a report on a recurring schedule so a PDF lands in their inbox every Monday morning. Today we export manually and forward it, which is workable but easy to forget.\n\nA few teams have asked whether subscriptions could support filtered views per recipient, for example each regional manager receiving only their own territory. Not urgent for us, but it would remove a recurring bit of manual work and is something a couple of competing tools already offer." + } + ] +} diff --git a/src/datacustomcode/templates/script/examples/external_callout/README.md b/src/datacustomcode/templates/script/examples/external_callout/README.md new file mode 100644 index 0000000..feec946 --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/README.md @@ -0,0 +1,169 @@ +# Transform with a Gemini Named Credential Callout + +The **transform** calls Google's **Gemini** `generateContent` API to summarize text, and +write the result back to a DLO. Gemini is reached through a **Named Credential** +(`callout:gemini`), so this code never handles the endpoint URL or the API key. + +It shows **both** callout paths against the same Named Credential. + +## Shared request template + +```python +from datacustomcode.client import Client, named_credential_request_col + +# URL, method and headers apply to every callout on this template. +request = ( + HTTPRequestBuilder() + .set_url("callout:gemini") # callout:[/] + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json"}) + .set_response_timeout_seconds(60) # optional: per-callout timeout (seconds) + .build() +) +``` + +## Driver path — one-shot on the driver + +`Client.named_credential_request` runs the callout **once** on the driver and +returns an `HTTPResponse` (`.status_code`, `.body`, `.headers`, `.is_success`). +Use it for a lookup or a shared value you reuse across the job. + +```python +response = client.named_credential_request(request, body=json.dumps(payload)) +if response.is_success: + envelope = json.loads(response.body) + text = envelope["candidates"][0]["content"]["parts"][0]["text"] +``` + +## Per-row path — fan out across the DataFrame + +`named_credential_request_col` dispatches one callout per row. The body Column +varies per row, and so can the optional `path` Column. + +```python +# One callout per row; body is a Column built from the row's data. +df = df.withColumn("_callout", named_credential_request_col(request, body=body_col)) +``` + +### Optional per-row `path` + +Pass `path` to call a different endpoint of the **same** Named Credential per +row. It is appended to the template's URL, e.g. to pick a model per row when the +Named Credential's URL stops at `.../v1beta/models`: + +```python +from pyspark.sql.functions import col, concat, lit, url_encode, when + +# callout:gemini/gemini-2.5-pro:generateContent or .../gemini-2.5-flash:generateContent +path = when(col("priority__c") == "high", lit("/gemini-2.5-pro:generateContent")).otherwise( + lit("/gemini-2.5-flash:generateContent") +) +df = df.withColumn( + "_callout", named_credential_request_col(request, body=body_col, path=path) +) + +# Query-string lookups work too: callout:geo?address=1+Market+St +geo_path = concat(lit("?address="), url_encode(col("address__c"))) +``` + +- A path starting with `/` or `?` is appended as is; any other path is joined + with `/` (`2/v1` → `callout:gemini/2/v1`), so a row can never switch to a + different Named Credential. +- A null or empty `path` uses the template's URL unchanged; omitting `path` + behaves exactly as before. +- Keep the template URL free of a query string when using `path` — put + per-row query parameters in `path` instead. + +It returns a struct Column: + +``` +{status, response: {status_code, body, headers}, error_code, error_message} +``` + +- A non-2xx response is still `status = "SUCCESS"` with the HTTP code in + `response.status_code` — extracting the model text just yields null for that row. +- A transport failure sets `status = "ERROR"`; the row survives, the job does not + abort. Pull the model text out of `response.body` with `get_json_object(...)`. + +Both paths resolve the same Named Credential and read auth from the same local +`external_callout_config.json`. + +The `gemini` Named Credential's URL already includes the full +`/v1beta/models/:generateContent` path, so the callout is just +`callout:gemini` with **no path suffix** (anything after the name is appended to +the credential's URL). + +## Configure the Named Credential + +1. Create an **External Credential** (e.g. `google_api_key`) that injects your + Gemini API key as the `X-goog-api-key` header. +2. Create a **Named Credential** named `gemini`: + - **URL**: `https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent` + - **Enabled for Callouts** + **Generate Authorization Header**: on + - **External Credential**: `google_api_key` + +## Test locally + +Copy `entrypoint.py` into your `payload/` folder (or point the run at it), then: + +```bash +DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG=/abs/path/to/external_callout_config.json \ + sf data-code-extension script run --entrypoint entrypoint.py --target-org +``` + +With `----target-org` the SDK fetches only the **URL** from the org's Named +Credential; **auth is always taken from `external_callout_config.json`** locally +(the org's External Credential is used only in the Data Cloud runtime). So the +`X-goog-api-key` must be in the local config for a local test. Omit +`--target-org` to run fully offline using `target_url`. + +```json +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { "X-goog-api-key": "YOUR_GEMINI_API_KEY" }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} +``` + +Place `external_callout_config.json` in the **parent of your payload folder** (or +point `DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG` at it). It is never packaged into +the deployment zip. Get a key from [Google AI Studio](https://aistudio.google.com/apikey); +**do not commit it.** + +## What it reads / writes + +`config.json` declares the DLO permissions for deployment: + +| | DLO | Notes | +| ------ | ------------------- | --------------------------------------------- | +| read | `Account_std__dll` | source rows; `description__c` is summarized | +| write | `Account_std_copy__dll` | adds `summary__c`, `callout_status__c`, `callout_http_code__c` | + +Adjust `_TEXT_COLUMN`, `_SOURCE_DLO` and `_TARGET_DLO` in `entrypoint.py` (and the +matching entries in `config.json`) to point at your own DLOs. + +## Auth types + +`auth_type` selects how auth is injected for local testing. It should mirror the +External Credential your Named Credential uses in the org, so local and deployed +runs behave the same. This example uses `Custom` (Gemini's `X-goog-api-key`); +all supported types: + +| `auth_type` | Fields read | Header sent | +| ----------- | ------------------------------- | --------------------------------------- | +| `Basic` | `username`, `password` | `Authorization: Basic ` | +| `Custom` | `custom_headers` (sent verbatim)| the headers you list | +| `OAuth` | `access_token` or `token` | `Authorization: Bearer ` | +| `Jwt` | `access_token` or `token` | `Authorization: Bearer ` | +| `AwsSv4` | `aws_access_key_id`, `aws_secret_access_key`, `aws_region`, `aws_service`, optional `aws_session_token` | `Authorization: AWS4-HMAC-SHA256 ...` plus `x-amz-date` / `x-amz-content-sha256` (and `x-amz-security-token` when a session token is set) | + +`OAuth`/`Jwt` take a token you supply for the local run — the SDK does not fetch +or refresh it. In the Data Cloud runtime the Named Credential handles token +acquisition; this local config only stands in for that during testing. + +`AwsSv4` signs the request with AWS Signature Version 4 using the keys you +supply, mirroring an AWS Signature Version 4 External Credential in the org. diff --git a/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py b/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py new file mode 100644 index 0000000..547bcef --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/entrypoint.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 + +""" +Data Transform with a Gemini Named Credential Callout + +This is a batch transform: read a DLO, enrich it via Google's Gemini ``generateContent`` +API, and write the result back to a DLO. + +It shows **both** callout paths against the same Named Credential +(``callout:gemini``); the request template (URL, method, headers) is shared and +only the body differs: + +- **Driver path** — :meth:`Client.named_credential_request` runs **once** on the + driver and returns an :class:`HTTPResponse`. Use it for a single job-level + callout whose result you reuse across the job. +- **Per-row path** — :func:`datacustomcode.client.named_credential_request_col` + fans the callout out across the DataFrame, one call per row, returning a struct + Column ``{status, response, error_code, error_message}`` where ``response`` is + itself ``{status_code, body, headers}``. A non-2xx response is still + ``SUCCESS`` with its code in ``response.status_code``; a transport failure sets + ``status`` to ``ERROR`` without aborting the whole Spark job. +""" + +import json +import logging + +from pyspark.sql import Column +from pyspark.sql.functions import ( + array, + col, + concat, + get_json_object, + lit, + struct, + to_json, +) + +from datacustomcode.client import Client, named_credential_request_col +from datacustomcode.io.writer.base import WriteMode +from datacustomcode.named_credential.types.http_method import HTTPMethod +from datacustomcode.named_credential.types.http_request_builder import ( + HTTPRequestBuilder, +) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +CALLOUT_URL = "callout:gemini" + +_TEXT_COLUMN = "Description__c" +_SOURCE_DLO = "Account_std__dll" +_TARGET_DLO = "Account_std_copy__dll" + +_PROMPT = ( + "Summarize the following text in one sentence. Respond with the summary " + "only, no preamble.\n\nText:\n" +) + +_REQUEST = ( + HTTPRequestBuilder() + .set_url(CALLOUT_URL) + .set_method(HTTPMethod.POST) + .set_headers({"Content-Type": "application/json", "Accept": "application/json"}) + .set_response_timeout_seconds(60) + .build() +) + + +def _gemini_body_col(text_col: Column) -> Column: + """Build a per-row Gemini ``generateContent`` request body as a JSON string. + + Using ``to_json(struct(...))`` keeps the row text properly escaped inside the + JSON payload rather than string-concatenating it. + """ + prompt = concat(lit(_PROMPT), text_col) + contents = array(struct(array(struct(prompt.alias("text"))).alias("parts"))) + return to_json(struct(contents.alias("contents"))) + + +def _gemini_body(text: str) -> str: + """Build a Gemini ``generateContent`` request body as a JSON string (driver).""" + return json.dumps({"contents": [{"parts": [{"text": _PROMPT + text}]}]}) + + +def _summarize_on_driver(client: Client, text: str) -> str: + """One-shot driver callout: summarize a single string once, not per row. + + The scalar counterpart to the per-row column path — same request template, + but dispatched once on the driver and returning an ``HTTPResponse``. + """ + response = client.named_credential_request(_REQUEST, body=_gemini_body(text)) + + # Don't raise: a failed driver callout shouldn't abort the whole job. + if not response.is_success: + logger.error(f"Driver Gemini callout failed: HTTP {response.status_code}") + return "" + + envelope = json.loads(response.body) if response.body else {} + try: + return str(envelope["candidates"][0]["content"]["parts"][0]["text"]) + except (KeyError, IndexError, TypeError): + return "" + + +def main(): + client = Client() + + df = client.read_dlo(_SOURCE_DLO) + + # Driver path: one callout on the driver over a single representative row. + sample = df.select(_TEXT_COLUMN).first() + if sample and sample[0]: + driver_summary = _summarize_on_driver(client, sample[0]) + logger.info(f"Driver-path sample summary: {driver_summary}") + + # Per-row path: one Gemini callout per row; the result struct is a column. + callout = named_credential_request_col( + _REQUEST, body=_gemini_body_col(col(_TEXT_COLUMN)) + ) + df = df.withColumn("_callout", callout) + + # Pull the model's text out of the response body. A row whose callout failed + # (non-2xx or transport error) yields null here rather than failing the job. + summary = get_json_object( + col("_callout")["response"]["body"], + "$.candidates[0].content.parts[0].text", + ) + + df = df.select( + col("id__c").alias("id__c"), + col("Description__c").alias("description__c"), + col("kq_id__c").alias("kq_id__c"), + summary.alias("summary__c"), + col("_callout")["status"].alias("callout_status__c"), + col("_callout")["response"]["status_code"].alias("callout_http_code__c"), + ) + + client.write_to_dlo(_TARGET_DLO, df, write_mode=WriteMode.APPEND) + + +if __name__ == "__main__": + main() diff --git a/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json b/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json new file mode 100644 index 0000000..2f979ea --- /dev/null +++ b/src/datacustomcode/templates/script/examples/external_callout/external_callout_config.json @@ -0,0 +1,11 @@ +{ + "credentials": { + "callout:gemini": { + "auth_type": "Custom", + "custom_headers": { + "X-goog-api-key": "YOUR_API_KEY" + }, + "target_url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-flash-latest:generateContent" + } + } +} diff --git a/tests/io/writer/test_print.py b/tests/io/writer/test_print.py index a10f2a3..258a765 100644 --- a/tests/io/writer/test_print.py +++ b/tests/io/writer/test_print.py @@ -65,6 +65,12 @@ def test_write_dlo_deltas_not_supported_locally(self, print_writer, mock_datafra print_writer.write_dlo_deltas("test_dll", mock_dataframe) assert "write_dlo_deltas" in str(exc_info.value) + def test_write_dmo_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_dmo_deltas("test_dlm", mock_dataframe) + assert "write_dmo_deltas" in str(exc_info.value) + def test_config_name(self): """Test that the CONFIG_NAME class variable is set correctly.""" assert PrintDataCloudWriter.CONFIG_NAME == "PrintDataCloudWriter" diff --git a/tests/test_client.py b/tests/test_client.py index 1293ea1..748c0fc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -440,6 +440,40 @@ def test_write_dlo_deltas_after_dmo_read_raises_exception( assert "test_dmo" in str(exc_info.value) writer.write_dlo_deltas.assert_not_called() + def test_write_dmo_deltas(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + mock_query = MagicMock() + writer.write_dmo_deltas.return_value = mock_query + + client = StreamingClient(reader=reader, writer=writer) + client._record_dmo_access("some_dmo") + + result = client.write_dmo_deltas("test_dmo", mock_df, extra_param=True) + + writer.write_dmo_deltas.assert_called_once_with( + "test_dmo", mock_df, extra_param=True + ) + assert result is mock_query + + def test_write_dmo_deltas_after_dlo_read_raises_exception( + self, reset_client, mock_spark + ): + """Streaming DMO 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_dlo_access("test_dlo") + + with pytest.raises(DataCloudAccessLayerException) as exc_info: + client.write_dmo_deltas("test_dmo", mock_df) + + assert "test_dlo" in str(exc_info.value) + writer.write_dmo_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) @@ -458,6 +492,24 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): writer.write_dlo_deltas.assert_called_once_with("target_dll", stream_df) assert "source_dll" in client._data_layer_history[DataCloudObjectType.DLO] + def test_streaming_dmo_read_write_flow(self, reset_client, mock_spark): + """A read_dmo_deltas → write_dmo_deltas flow stays within the DMO layer.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + stream_df = MagicMock(spec=DataFrame) + reader.read_dmo_deltas.return_value = stream_df + + client = StreamingClient(reader=reader, writer=writer) + + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = "source_dlm" + df = client.read_dmo_deltas() + client.write_dmo_deltas("target_dlm", df) + + reader.read_dmo_deltas.assert_called_once_with() + writer.write_dmo_deltas.assert_called_once_with("target_dlm", stream_df) + assert "source_dlm" in client._data_layer_history[DataCloudObjectType.DMO] + def test_auto_write_to_dlo(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) diff --git a/tests/test_deploy.py b/tests/test_deploy.py index c592bba..1e33f3c 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -1645,9 +1645,7 @@ def test_is_streaming_false_when_source_absent(self): ) 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)""" + def test_streaming_allows_dmo_read_dmo_write(self): config = DataTransformConfig( sdkVersion="1.0.0", entryPoint="entrypoint.py", @@ -1655,24 +1653,24 @@ def test_streaming_allows_dmo_read_dlo_write(self): streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), permissions=Permissions( read=DmoPermission(dmo=["input_dmo__dlm"]), - write=DloPermission(dlo=["output_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), ), ) assert config.is_streaming is True assert isinstance(config.permissions.read, DmoPermission) - assert isinstance(config.permissions.write, DloPermission) + assert isinstance(config.permissions.write, DmoPermission) - def test_streaming_rejects_dmo_write(self): - """Streaming writes must target a DLO.""" - with pytest.raises(ValueError, match="must write to a DLO"): + def test_streaming_rejects_dmo_read_dlo_write(self): + """Streaming follows the same same-layer rule as batch.""" + with pytest.raises(ValueError, match="both reference"): DataTransformConfig( sdkVersion="1.0.0", entryPoint="entrypoint.py", dataspace="default", - streamingSource=StreamingSource(type="dlo", name="input_dlo"), + streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), permissions=Permissions( - read=DloPermission(dlo=["input_dlo"]), - write=DmoPermission(dmo=["output_dmo__dlm"]), + read=DmoPermission(dmo=["input_dmo__dlm"]), + write=DloPermission(dlo=["output_dlo"]), ), ) diff --git a/tests/test_scan.py b/tests/test_scan.py index 716c233..71ead13 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -735,7 +735,7 @@ def my_function(event, context): client = StreamingClient() deltas = client.read_dmo_deltas() - query = client.write_dlo_deltas("Account_copy__dll", deltas) + query = client.write_dmo_deltas("Account_copy__dlm", deltas) query.awaitTermination() """ ) @@ -751,6 +751,13 @@ def test_file_is_streaming_true_for_delta_methods(self): finally: os.unlink(temp_path) + def test_file_is_streaming_true_for_dmo_delta_methods(self): + temp_path = create_test_script(STREAMING_DMO_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( """ @@ -778,13 +785,72 @@ def test_scan_file_streaming_dlo(self): finally: os.unlink(temp_path) - def test_scan_file_streaming_dmo_read(self): + def test_scan_file_streaming_dmo(self): temp_path = create_test_script(STREAMING_DMO_ENTRYPOINT) try: result = scan_file_streaming(temp_path) + assert result.read_dlo_deltas is False assert result.read_dmo_deltas is True assert result.read_layer == "dmo" - assert "Account_copy__dll" in result.write_dlo_deltas + assert "Account_copy__dlm" in result.write_dmo_deltas + assert result.write_layer == "dmo" + assert not result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_both_write_methods(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + client.write_dlo_deltas("Account_copy__dll", deltas) + client.write_dmo_deltas("Account_copy__dlm", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="Cannot write DLO and DMO deltas"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_dlo_read_dmo_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + client.write_dmo_deltas("Account_copy__dlm", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises( + ValueError, match=r"same layer \(got read=dlo, write=dmo\)" + ): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_dmo_read_dlo_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dmo_deltas() + client.write_dlo_deltas("Account_copy__dll", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises( + ValueError, match=r"same layer \(got read=dmo, write=dlo\)" + ): + scan_file_streaming(temp_path) finally: os.unlink(temp_path) @@ -875,6 +941,7 @@ def test_reasserts_layer_from_code(self): assert updated["streamingSource"]["type"] == "dmo" assert updated["streamingSource"]["name"] == "Account__dlm" assert updated["permissions"]["read"]["dmo"] == ["Account__dlm"] + assert updated["permissions"]["write"] == {"dmo": ["Account_copy__dlm"]} def test_batch_entrypoint_drops_stale_streaming_source(self): """Switching a streaming entrypoint to batch clears streamingSource."""