diff --git a/AGENTS.md b/AGENTS.md index 6c174db60..497183132 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,6 +24,7 @@ dapr/ # Core SDK package (single PyPI dist: `pip install ├── serializers/ # JSON and pluggable serializers ├── version/ # Version metadata └── ext/ # Extensions, installable as extras to the base package + ├── databricks/ # Databricks Lakeflow sink ← see dapr/ext/databricks/AGENTS.md (`pip install dapr[databricks]`) ├── fastapi/ # FastAPI integration ← see dapr/ext/fastapi/AGENTS.md (`pip install dapr[fastapi]`) ├── flask/ # Flask integration ← see dapr/ext/flask/AGENTS.md (`pip install dapr[flask]`) ├── grpc/ # gRPC App extension ← see dapr/ext/grpc/AGENTS.md (`pip install dapr[grpc]`) @@ -62,6 +63,7 @@ Extensions are bundled into the core `dapr` wheel and exposed as installable ext | `dapr[flask]` | `dapr.ext.flask` | Flask integration for pub/sub and actors (legacy `flask_dapr` import path is a deprecated shim) | Low | | `dapr[langgraph]` | `dapr.ext.langgraph` | LangGraph checkpoint persistence to Dapr state store | Moderate | | `dapr[strands]` | `dapr.ext.strands` | Strands agent session management via Dapr state store | New | +| `dapr[databricks]` | `dapr.ext.databricks` | Databricks Lakeflow streaming sink -> Dapr Workflow | New | The previously-separate distributions (`dapr-ext-*`, `flask-dapr`) are no longer published. `dapr/__init__.py` emits a `FutureWarning` if it detects a legacy install at import time; see `RELEASE.md` for the migration recipe. @@ -109,6 +111,7 @@ uv run python -m unittest discover -v ./tests/ext/grpc uv run python -m unittest discover -v ./tests/ext/fastapi uv run python -m unittest discover -v ./tests/ext/langgraph uv run python -m unittest discover -v ./tests/ext/strands +uv run python -m unittest discover -v ./tests/ext/databricks # pytest-style suites: uv run pytest -m "not e2e" ./tests/ext/workflow/durabletask/ diff --git a/README.md b/README.md index 0b1e13892..24f094d85 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ This includes the following packages: And the following extensions, installed as extras on the core `dapr` package: +* [dapr.ext.databricks](./dapr/ext/databricks): Databricks Lakeflow streaming sink for Dapr Workflow — `pip install "dapr[databricks]"` * [dapr.ext.fastapi](./dapr/ext/fastapi): FastAPI integration (actors, pub/sub) — `pip install "dapr[fastapi]"` * [dapr.ext.flask](./dapr/ext/flask): Flask integration (actors, pub/sub) — `pip install "dapr[flask]"` * [dapr.ext.grpc](./dapr/ext/grpc): gRPC AppCallback server — `pip install "dapr[grpc]"` diff --git a/dapr/ext/databricks/AGENTS.md b/dapr/ext/databricks/AGENTS.md new file mode 100644 index 000000000..8b5880214 --- /dev/null +++ b/dapr/ext/databricks/AGENTS.md @@ -0,0 +1,294 @@ +# AGENTS.md — dapr.ext.databricks + +Turns records from a Databricks Lakeflow streaming pipeline into durable Dapr Workflow +executions. Databricks keeps doing what it's good at (ingestion, transformation, streaming +checkpoints, Lakeflow execution); Dapr Workflow takes over the business process a record +triggers (durable execution, retries, timers, human interaction, external systems). + +## Source layout + +``` +dapr/ext/databricks/ +├── __init__.py # Public API exports +├── sink.py # register_workflow_sink() + lazy pyspark.pipelines import +├── batch_handler.py # DaprWorkflowBatchHandler — the reusable low-level API +├── scheduling.py # ensure_workflow_scheduled() — the dedup/retry-safety core +├── identity.py # Deterministic instance-ID derivation + sanitization +├── mapping.py # default_row_mapper() — Row -> JSON-safe dict +├── config.py # WorkflowSinkConfig dataclass + validation +├── exceptions.py # DaprDatabricksError hierarchy +├── _typing.py # RowLike / BatchDataFrameLike Protocols (no pyspark import) +└── py.typed + +tests/ext/databricks/ +├── _fakes.py # Shared test doubles (FakeRow, FakeDataFrame, FakeWorkflowClient, ...) +├── test_identity.py # Deterministic IDs, sanitization, Unicode, hashing +├── test_mapping.py # Default + custom row mappers +├── test_scheduling.py # ensure_workflow_scheduled: new/existing/race/lost-response +├── test_batch_handler.py # End-to-end batch processing, retry, concurrency, metadata +├── test_config.py # WorkflowSinkConfig validation +└── test_sink.py # register_workflow_sink + pyspark.pipelines wiring + +tests/integration/test_databricks_sink.py # Runs against a real Dapr sidecar (no Databricks needed) +examples/databricks/ # Fraud-remediation example (see its own README) +``` + +Installed via the `databricks` extra: `pip install "dapr[databricks]"`. Like `workflow`, this +extra has **no third-party runtime dependencies** — see "Dependency design" below. + +## Architecture + +``` +┌───────────────────────────────────────────────────────────┐ +│ Databricks Lakeflow pipeline │ +│ @dp.append_flow(target=) -> readStream │ +└──────────────────────────┬──────────────────────────────────┘ + │ micro-batch (DataFrame, batch_id) +┌──────────────────────────▼──────────────────────────────────┐ +│ register_workflow_sink() │ +│ - lazily imports pyspark.pipelines │ +│ - registers @dp.foreach_batch_sink(name=...) wrapping │ +│ DaprWorkflowBatchHandler.process │ +└──────────────────────────┬────────────────────────────────────┘ + │ +┌──────────────────────────▼────────────────────────────────────┐ +│ DaprWorkflowBatchHandler.process(df, batch_id) │ +│ - df.toLocalIterator(prefetchPartitions=True) — never .collect() │ +│ - bounded concurrency via ThreadPoolExecutor(max_in_flight) │ +│ - per row: identity.py -> scheduling.py -> optional metadata wrap │ +│ - any unresolved row -> raises, failing the micro-batch │ +└──────────────────────────┬────────────────────────────────────────────┘ + │ +┌──────────────────────────▼────────────────────────────────────────────┐ +│ scheduling.ensure_workflow_scheduled(client, workflow, instance_id, ...) │ +│ 1. get_workflow_state(instance_id) -> exists? already accepted, done │ +│ 2. else schedule_new_workflow(...); ALREADY_EXISTS response -> also done │ +│ 3. any other error propagates (auth, unavailable, timeout, throttling, ...) │ +└──────────────────────────┬──────────────────────────────────────────────────┘ + │ + dapr.ext.workflow.DaprWorkflowClient (existing, unmodified) + │ + Dapr sidecar / endpoint (gRPC) +``` + +No new workflow engine, no new transport: this extension is entirely a scheduling and +identity layer on top of the existing `dapr.ext.workflow.DaprWorkflowClient`. + +## Public API + +```python +from dapr.ext.databricks import ( + register_workflow_sink, # high-level: registers a foreach_batch_sink + DaprWorkflowBatchHandler, # low-level: reusable batch-processing handler + WorkflowSinkConfig, # typed config, if constructing a handler directly + default_row_mapper, # the default Row -> dict mapper + DaprDatabricksError, # base exception + SinkConfigurationError, # bad register_workflow_sink() config + MissingBusinessKeyError, # configured id_field/id_fields absent or null on a row + DaprDatabricksSinkError, # a micro-batch could not be durably handed off +) +``` + +`register_workflow_sink(name, workflow, *, id_field=None, id_fields=None, namespace='default', +generation='v1', input_mapper=None, instance_id_factory=None, metadata=True, max_in_flight=8, +max_records_per_batch=None, host=None, port=None)` is the primary entry point. See its +docstring in `sink.py` for the full parameter reference — kept there rather than duplicated here +so it can't drift out of sync. + +## Delivery semantics (read this before changing `identity.py` or `scheduling.py`) + +This is the part of the extension that must stay correct under review; the rest is +straightforward plumbing. + +**The property this extension provides:** retry-safe handoff while Dapr workflow +instance/history is retained. Concretely: + +``` +deterministic instance ID (namespace-sink-generation-business_key) + + +check-then-schedule (scheduling.ensure_workflow_scheduled) + + +Dapr rejecting a duplicate non-terminal instance ID + = +a Lakeflow retry of the same micro-batch recognizes work it already +handed off, instead of scheduling a second execution +``` + +**What it is not:** strict exactly-once delivery to arbitrary downstream systems. It is +exactly-once-Dapr-workflow-*scheduling* per `(namespace, sink, generation, business_key)`, +for as long as that instance's history has not been purged. Never describe this as +exactly-once in docs/comments/log messages without that qualification. + +**The three scenarios this must handle** (each has a dedicated test in +`test_scheduling.py` and/or `test_batch_handler.py` — do not remove test coverage for any +of them): + +1. **Micro-batch retry after partial failure.** Batch has records A and B; A schedules, + B fails. The batch raises (see below), Lakeflow retries. On retry, A's + `get_workflow_state` finds it already exists (skipped, not re-scheduled); B schedules + for the first time. Net result: exactly one workflow each. +2. **Lost response.** `schedule_new_workflow` durably succeeds on Dapr's side, but the + client never observes success (timeout, connection drop). The attempt still raises + (we cannot prove acceptance), the batch fails, Lakeflow retries. On retry, + `get_workflow_state` finds the instance Dapr already accepted and skips scheduling. +3. **Concurrent duplicate race.** Two callers derive the same instance ID (e.g. two + in-flight retries) and both see "absent" before either schedules. Both call + `schedule_new_workflow`; Dapr accepts exactly one and rejects the other with a + duplicate-instance error. `scheduling.is_duplicate_instance_error` recognizes that + rejection and both callers report success. + +**Why the "exists" check ignores workflow status.** `ensure_workflow_scheduled` treats +*any* existing instance (including a long-completed one) as "already handled" — it does +not re-schedule just because the prior run reached a terminal state. Dapr itself would +allow reusing a terminal instance ID; this extension deliberately does not, because the +business key already produced one execution in this generation and re-running it would +be exactly the accidental-duplicate-side-effect this extension exists to prevent. This is +also why `generation` exists at all (see below) — it is the only sanctioned way to get a +fresh identity space on purpose. + +**Duplicate-instance detection is dual-signaled** (`scheduling.is_duplicate_instance_error`): +`grpc.StatusCode.ALREADY_EXISTS` first, then a case-insensitive `'already exists'` substring +match on the error details as a fallback. The Dapr Workflow HTTP API reference documents a +409 response with that phrase for this condition; matching on text in addition to the status +code follows the existing precedent in `DaprWorkflowClient.get_workflow_state`, which already +relies on message-text matching (`'no such instance exists'`) rather than a status code alone. + +**The batch-level failure rule:** if any record's outcome cannot be established as either +newly-scheduled or already-durably-present, `DaprWorkflowBatchHandler.process` raises +`DaprDatabricksSinkError`, chaining the first underlying error. This is deliberate — a +`foreach_batch_sink` that swallows a partial failure would let Lakeflow believe the batch +succeeded while some records were silently never scheduled. + +## Full-refresh semantics (`generation`) + +Lakeflow's `batch_id` is **not** a globally unique event identity: `batch_id` restarts at 0 +both for a fresh stream and after a full pipeline refresh (Databricks' own docs: "A `batch_id` +of `0` represents the start of a stream, or the beginning of a full refresh"). If instance IDs +were derived from `batch_id` alone, a full refresh would silently replay every business action. + +`generation` (default `'v1'`) is part of the deterministic instance-ID template specifically to +make this safe by default: normal retries keep `generation` unchanged and dedupe correctly +against prior runs in the same generation; a user who *intentionally* wants to replay business +actions after a full refresh changes `generation` (e.g. to `'v2'`), which — combined with +`namespace` and the sink name — produces an entirely fresh identity space. Nothing about a full +refresh changes `generation` automatically; that is a deliberate, safe-by-default choice: an +accidental full refresh must not replay business actions just because Databricks reset its +checkpoint. + +## ID derivation and sanitization (`identity.py`) + +Template: `---`, or +`----` when no business-key strategy is +configured (weaker guarantee — see its docstring caveat about relying on stable row ordering +across retries). + +Business key source, in priority order (mutually exclusive; `WorkflowSinkConfig` rejects +configuring more than one): `instance_id_factory(row, batch_id)` > `id_field` > `id_fields` +(joined with `_`). All three still get namespaced by `namespace`/sink/`generation` and run +through the same sanitizer — `instance_id_factory` is an escape hatch for *how the business key +is computed*, not a bypass of the identity/generation safety net. + +Sanitization (`sanitize_segment`): a segment that is already short (<= 80 chars) and matches +`[A-Za-z0-9_-]+` (Dapr's documented allowed instance-ID characters) passes through unchanged, +so ordinary business keys stay human-readable. Anything else — Unicode, punctuation, empty +strings, over-length values — is replaced by a SHA-256 hex digest of the *whole* segment. +Hashing the whole segment (rather than stripping/replacing individual bad characters) is +deliberate: a character-by-character transliteration can collide distinct inputs (ASCII-folding +would collapse many distinct non-Latin business keys to the same, near-empty output; stripping +would collide `"a/b"` and `"a-b"`). The final composed ID is also capped at 128 characters +(this extension's own conservative bound — Dapr does not publish a maximum) via the same +hash-the-whole-thing strategy, careful to never truncate away the hash suffix itself (which +would collide every record in a sink with a pathologically long namespace/sink/generation +prefix onto the same ID — see the comment in `derive_instance_id`). + +## Spark execution model + +`DaprWorkflowBatchHandler.process` iterates `df.toLocalIterator(prefetchPartitions=True)`, +never `df.collect()` — a micro-batch is never required to fit entirely in driver memory at +once. Concurrency is bounded by `max_in_flight` (a plain `ThreadPoolExecutor`, sized to that +many workers — no unbounded fan-out of scheduling calls). `max_records_per_batch`, when set, is +a hard fail-fast cap, not a silent truncation: exceeding it raises immediately (already-submitted +records finish first) so an operator notices a business-action sink got pointed at a bulk/ +analytical stream, rather than quietly scheduling millions of workflows. + +Everything in `process()` runs on the Spark **driver** (this is inherent to +`foreach_batch_sink`/`toLocalIterator`, not a choice this extension makes) — appropriate for +business-action streams (the intended use case), not bulk data replication. + +**`toLocalIterator()` is not universally available.** Confirmed via a live end-to-end run +against real Databricks serverless Lakeflow compute: `df.toLocalIterator(prefetchPartitions=True)` +itself raised `Exception: toLocalIterator() is not supported when using file-based collect` +(a Spark Connect-backed-compute limitation, not anything this extension controls) — +synchronously, before any row was read. `DaprWorkflowBatchHandler._iter_rows` catches exactly +that message and falls back to a `collect()`, bounded by `max_records_per_batch` (via +`df.limit(max_records_per_batch + 1)`) when configured, and logs a warning either way — this +is a degraded-safety fallback, not a silent one. Any other exception from `toLocalIterator` +(a corrupt source table, an expired storage token, etc.) is not this fallback's concern and +propagates normally. **Set `max_records_per_batch` when running on compute where +`toLocalIterator` is unavailable** — without it, the fallback has no bound and behaves like a +plain `collect()`. + +## Authentication and secrets + +No new auth/TLS mechanism: `DaprWorkflowBatchHandler` constructs (or accepts) a +`dapr.ext.workflow.DaprWorkflowClient`, which resolves `DAPR_GRPC_ENDPOINT` / +`DAPR_RUNTIME_HOST` / `DAPR_GRPC_PORT` / `DAPR_API_TOKEN` exactly as every other workflow +client in this SDK does. Nothing in this extension logs a payload's business data by default — +structured log lines (see `batch_handler._process_row`) carry only sink/workflow/instance-id/ +batch-id/namespace/generation/outcome/latency, never the row or workflow input itself. + +## Dependency design + +`pyspark.pipelines` is provided by the Databricks Lakeflow (Spark Declarative Pipelines) +runtime — it is not part of a plain `pip install pyspark`, and even where a `pyspark` package +is present, `pyspark.pipelines` may not be. So: + +- The `databricks` extra in `pyproject.toml` is **empty**, like `workflow` — no pyspark + dependency, on purpose. +- `sink.py` imports `pyspark.pipelines` lazily, inside `register_workflow_sink`, not at module + scope. `dapr.ext.databricks` (and `DaprWorkflowBatchHandler`) import cleanly with zero pyspark + installed; only calling `register_workflow_sink` outside a Lakeflow pipeline fails, with the + message: *"Databricks Lakeflow support requires pyspark.pipelines.foreach_batch_sink. Run + this integration inside a supported Databricks Lakeflow pipeline."* +- `_typing.py` defines `RowLike`/`BatchDataFrameLike` as structural `Protocol`s instead of + importing `pyspark.sql.Row`/`DataFrame`, so signatures stay strongly typed without a pyspark + import anywhere, including under `TYPE_CHECKING`. `pyproject.toml` still needs + `[[tool.mypy.overrides]] module = ["pyspark.*"] ignore_missing_imports = true` for the one + `from pyspark import pipelines` line inside `sink.py`'s lazy-import function. + +## Testing + +```bash +uv run pytest tests/ext/databricks/ # unit tests, no pyspark/Databricks needed +uv run pytest tests/integration/test_databricks_sink.py # against a real Dapr sidecar +uv run pytest tests/examples/test_databricks.py # fraud-remediation example +``` + +All unit tests mock/fake three things (`tests/ext/databricks/_fakes.py`): `Row`/`DataFrame` +(`FakeRow`/`FakeDataFrame` — plain Python, no pyspark), the Dapr workflow client +(`FakeWorkflowClient` — an in-memory instance-exists/schedule simulator that models Dapr's +actual duplicate-instance rejection, including under real thread concurrency via a lock), and +`pyspark.pipelines` itself (`test_sink.py` injects a fake module via `sys.modules`, since +pyspark is never installed in this environment). + +`BarrierSyncedWorkflowClient` (in `_fakes.py`) exists specifically to make the "concurrent +duplicate race" test deterministic — it holds every caller at the start of +`schedule_new_workflow` until all parties have arrived, so the race is guaranteed to occur +rather than depending on incidental thread-scheduling timing. + +## Known limitations + +- Not exactly-once to arbitrary downstream systems — see "Delivery semantics" above. +- The batch/record-index fallback identity (no business key configured) depends on Lakeflow + redelivering the same rows in the same order on retry; this is generally true for a + straightforward read-then-sink pipeline but is a materially weaker guarantee than a business + key, and is documented as such rather than presented as equivalent. +- `max_records_per_batch` protects against runaway fan-out but does not itself rate-limit the + Lakeflow source; pair it with upstream rate limiting (e.g. `maxFilesPerTrigger`/ + `maxBytesPerTrigger` on the source `readStream`) for real protection. +- Purging a workflow instance's history removes Dapr's record of "already handled" for that + instance ID; retention/purge policy is the operator's responsibility, same as for any other + Dapr Workflow usage. +- On compute where `toLocalIterator()` is unavailable (observed on Databricks serverless), the + handler falls back to a `collect()` bounded only by `max_records_per_batch` — set that + explicitly on such compute; see "Spark execution model" above. diff --git a/dapr/ext/databricks/README.md b/dapr/ext/databricks/README.md new file mode 100644 index 000000000..05895afc7 --- /dev/null +++ b/dapr/ext/databricks/README.md @@ -0,0 +1,180 @@ +# dapr.ext.databricks + +Turns records emitted by a Databricks Lakeflow streaming pipeline into durable Dapr Workflow +executions. + +```text +Databricks Lakeflow + │ + │ streaming records + ▼ +Dapr Databricks Sink + │ + ▼ +Dapr Workflow + │ + ┌────┼──────────────┐ + ▼ ▼ ▼ + APIs SaaS Humans + systems +``` + +```sh +pip install "dapr[databricks]" +``` + +## Why this exists + +External side effects (freezing a card, calling a partner API, opening a case for a human to +review) have fundamentally different reliability requirements than transformations inside a +data pipeline. Instead of writing that business logic directly inside `foreach_batch_sink`, +hand the record to a Dapr Workflow and let it continue independently — with retries, timers, +human-in-the-loop waits, and crash recovery, all outside the streaming query's lifetime. + +## Quick start + +```python +from pyspark import pipelines as dp +from dapr.ext.databricks import register_workflow_sink + +register_workflow_sink( + name='order_actions', + workflow='process_order', + id_field='order_id', + namespace='orders', +) + +@dp.append_flow(target='order_actions', name='order_actions_flow') +def order_actions_flow(): + return spark.readStream.table('validated_orders') +``` + +`register_workflow_sink` registers the `foreach_batch_sink` for you — you do not write one by +hand for the standard case. See `examples/databricks/` in this repository for a complete, +runnable (no Databricks needed) fraud-remediation walkthrough. + +### Custom input mapping + +```python +register_workflow_sink( + name='customer_actions', + workflow='process_customer', + id_field='customer_id', + input_mapper=lambda row: {'customer': row['customer_id'], 'status': row['status']}, +) +``` + +Without `input_mapper`, each row becomes a JSON-safe dict via `Row.asDict(recursive=True)` +(datetimes/dates -> ISO 8601 strings, `Decimal` -> string to avoid precision loss, bytes -> +base64). + +### Composite keys and the escape hatch + +```python +register_workflow_sink( + name='transfers', + workflow='process_transfer', + id_fields=['account_id', 'transaction_id'], +) + +# Or, for full control over how the business key is computed: +register_workflow_sink( + name='transfers', + workflow='process_transfer', + instance_id_factory=lambda row, batch_id: f"{row['account_id']}:{row['transaction_id']}", +) +``` + +`instance_id_factory` computes the business-key *component* — the result is still namespaced by +`namespace`/`name`/`generation` and sanitized like any other key, so it can't accidentally +disable the full-refresh safety net described below. + +### Optional lower-level API + +```python +from dapr.ext.databricks import DaprWorkflowBatchHandler, WorkflowSinkConfig + +handler = DaprWorkflowBatchHandler( + WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') +) + +@dp.foreach_batch_sink(name='orders') +def orders_handler(df, batch_id): + handler.process(df, batch_id) +``` + +Use this if you need to fold the handoff into a hand-written `foreach_batch_sink` (custom +pre/post-processing around the same batch). The high-level `register_workflow_sink` remains the +recommended path for everything else. + +## Delivery semantics + +```text +Lakeflow delivery + + +deterministic workflow identity + + +Dapr workflow persistence + = +retry-safe handoff while workflow history is retained +``` + +Every workflow instance ID is derived deterministically — +`---` — **never** a random UUID. Before scheduling, +the sink checks whether that instance already exists; if so, the record is treated as already +handled. If not, it schedules, and treats Dapr's "instance already exists" rejection (a +concurrent scheduler winning a race, or this exact instance from a previous attempt) the same +way. If neither check can be completed (Dapr unavailable, a timeout, an auth failure, ...), the +whole micro-batch fails so Lakeflow retries it — the sink never reports success while some +record's fate is unknown. + +This is **retry-safe, not exactly-once**: it depends on Dapr retaining that instance's +history/state. Once a workflow instance is purged, Dapr (and therefore this extension) can no +longer tell that its business key was already handled; scheduling it again would start a new +execution. Plan your workflow-history retention accordingly if you need this guarantee to hold +indefinitely. + +### Full refresh + +Lakeflow's `batch_id` restarts at `0` both for a brand-new stream and after a full pipeline +refresh — it is never, by itself, a safe uniqueness key. `generation` (default `'v1'`) is part +of the deterministic instance ID for exactly this reason: normal retries keep `generation` +unchanged and continue deduplicating correctly; if you deliberately want to replay business +actions after a full refresh, bump `generation` (e.g. to `'v2'`) to get a fresh identity space +on purpose. Nothing changes `generation` for you — an accidental full refresh must not silently +replay business actions. + +## Observability + +Each scheduling attempt logs one structured line via the standard `logging` module (logger +`dapr.ext.databricks.batch_handler`) with sink name, workflow name, instance ID, Lakeflow batch +ID, namespace, generation, whether the instance was newly scheduled or already existed, and +scheduling latency. Business record contents are never logged. + +## Configuration reference + +```python +register_workflow_sink( + name='orders', # sink name (foreach_batch_sink name) + workflow='process_order', # registered Dapr Workflow name + + id_field='order_id', # business key column (mutually exclusive with the two below) + id_fields=None, # composite business key columns + instance_id_factory=None, # (row, batch_id) -> business-key component + + namespace='orders', # identity partition; part of the instance ID + generation='v1', # identity epoch; bump to replay after a full refresh + + input_mapper=None, # Row -> dict; defaults to a JSON-safe asDict() + metadata=True, # wrap input as {'data': ..., 'metadata': {...}} + + max_in_flight=8, # bounded concurrent schedule_new_workflow calls + max_records_per_batch=None, # optional hard cap; fails the batch, never truncates silently + + host=None, port=None, # Dapr endpoint; defaults to the standard SDK env/settings +) +``` + +See `dapr/ext/databricks/AGENTS.md` in this repository for the full architecture, the exact +retry/race scenarios this extension handles, and the reasoning behind the ID-sanitization +scheme. diff --git a/dapr/ext/databricks/__init__.py b/dapr/ext/databricks/__init__.py new file mode 100644 index 000000000..e080eac3a --- /dev/null +++ b/dapr/ext/databricks/__init__.py @@ -0,0 +1,42 @@ +""" +Copyright 2026 The Dapr Authors +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. + +Turns Databricks Lakeflow streaming records into durable Dapr Workflow +executions. See ``AGENTS.md`` (architecture, delivery semantics) and +``README.md`` (usage) alongside this file for the full picture. + +Importing this package never requires pyspark: only calling +``register_workflow_sink`` (or otherwise reaching into ``pyspark.pipelines``) +does, and it fails with a clear error outside a Databricks Lakeflow pipeline. +""" + +from dapr.ext.databricks.batch_handler import DaprWorkflowBatchHandler +from dapr.ext.databricks.config import WorkflowSinkConfig +from dapr.ext.databricks.exceptions import ( + DaprDatabricksError, + DaprDatabricksSinkError, + MissingBusinessKeyError, + SinkConfigurationError, +) +from dapr.ext.databricks.mapping import default_row_mapper +from dapr.ext.databricks.sink import register_workflow_sink + +__all__ = [ + 'register_workflow_sink', + 'DaprWorkflowBatchHandler', + 'WorkflowSinkConfig', + 'default_row_mapper', + 'DaprDatabricksError', + 'SinkConfigurationError', + 'MissingBusinessKeyError', + 'DaprDatabricksSinkError', +] diff --git a/dapr/ext/databricks/_typing.py b/dapr/ext/databricks/_typing.py new file mode 100644 index 000000000..35792cfb4 --- /dev/null +++ b/dapr/ext/databricks/_typing.py @@ -0,0 +1,57 @@ +""" +Copyright 2026 The Dapr Authors +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 Any, Callable, Dict, Iterator, List, Protocol, runtime_checkable + +# This extension must import cleanly without pyspark installed (see AGENTS.md). +# Structural protocols describe exactly the pyspark.sql.Row / DataFrame surface +# we call, so callers get real static typing without a hard runtime dependency +# on pyspark, and tests can supply plain fakes instead of real Spark objects. + + +@runtime_checkable +class RowLike(Protocol): + """The subset of ``pyspark.sql.Row`` this extension relies on. + + Deliberately excludes ``__contains__``: ``Row`` subclasses ``tuple``, so + ``'x' in row`` tests membership among *values*, not field names — a + well-known footgun. Field presence is always checked via ``asDict()``. + """ + + def asDict(self, recursive: bool = ...) -> Dict[str, Any]: ... + + def __getitem__(self, key: str) -> Any: ... + + +@runtime_checkable +class BatchDataFrameLike(Protocol): + """The subset of ``pyspark.sql.DataFrame`` this extension relies on. + + ``limit``/``collect`` back the ``toLocalIterator`` fallback (see + ``batch_handler._iter_rows``): some compute (observed on Databricks + serverless / Spark Connect) raises on ``toLocalIterator`` itself + ("not supported when using file-based collect"), so a bounded + ``collect()`` is the fallback path, not an alternative primary path. + """ + + def toLocalIterator(self, prefetchPartitions: bool = ...) -> Iterator[RowLike]: ... + + def limit(self, num: int) -> 'BatchDataFrameLike': ... + + def collect(self) -> List[RowLike]: ... + + +RowMapper = Callable[[RowLike], Dict[str, Any]] +InstanceIdFactory = Callable[[RowLike, int], str] diff --git a/dapr/ext/databricks/batch_handler.py b/dapr/ext/databricks/batch_handler.py new file mode 100644 index 000000000..ba3704ac4 --- /dev/null +++ b/dapr/ext/databricks/batch_handler.py @@ -0,0 +1,253 @@ +""" +Copyright 2026 The Dapr Authors +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 logging +import time +from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from typing import Any, Dict, Iterable, List, Optional + +from dapr.ext.databricks._typing import BatchDataFrameLike, RowLike +from dapr.ext.databricks.config import WorkflowSinkConfig +from dapr.ext.databricks.exceptions import DaprDatabricksSinkError +from dapr.ext.databricks.identity import derive_instance_id, extract_business_key +from dapr.ext.databricks.scheduling import ensure_workflow_scheduled +from dapr.ext.workflow.dapr_workflow_client import DaprWorkflowClient + +_logger = logging.getLogger(__name__) + +# Observed on Databricks serverless / Spark Connect-backed compute: +# `df.toLocalIterator()` itself raises this, synchronously, before any row is +# read. Matching on this substring (rather than any exception at all) keeps +# genuinely unrelated failures — a corrupt source table, an expired ADLS +# token — propagating normally instead of being swallowed by the fallback. +_TO_LOCAL_ITERATOR_UNSUPPORTED_MARKER = 'toLocalIterator() is not supported' + + +class DaprWorkflowBatchHandler: + """Schedules one Dapr Workflow execution per row of a Lakeflow micro-batch. + + This is the reusable, lower-level building block behind + ``register_workflow_sink``. Most users should call that function + instead; use this class directly only if you need to fold the handoff + into a hand-written ``foreach_batch_sink`` (e.g. custom pre/post-processing + around the same batch): + + ```python + from pyspark import pipelines as dp + from dapr.ext.databricks import DaprWorkflowBatchHandler, WorkflowSinkConfig + + handler = DaprWorkflowBatchHandler( + WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') + ) + + @dp.foreach_batch_sink(name='orders') + def orders_handler(df, batch_id): + handler.process(df, batch_id) + ``` + + Only waits for Dapr to durably accept (or already hold) each workflow + instance — never for the workflow itself to finish — so a micro-batch + stays fast regardless of how long the triggered business process runs. + """ + + def __init__( + self, + config: WorkflowSinkConfig, + *, + workflow_client: Optional[DaprWorkflowClient] = None, + ) -> None: + """Creates a batch handler for one sink configuration. + + Args: + config: The sink's configuration. + workflow_client: An existing client to reuse — e.g. a test fake, + or one client shared across multiple sinks. When omitted, a + new ``DaprWorkflowClient`` is created from ``config.host``/ + ``config.port``, falling back to the standard Dapr SDK + environment/settings (``DAPR_GRPC_ENDPOINT``, ``DAPR_RUNTIME_HOST``, + ``DAPR_GRPC_PORT``, ``DAPR_API_TOKEN``). This is the same + resolution used everywhere else in ``dapr.ext.workflow``, so + it works whether Dapr is a local sidecar or a + network-accessible endpoint (e.g. a managed Dapr deployment) + — Databricks compute is not expected to have a sidecar + running in-process. + """ + self._config = config + self._owns_client = workflow_client is None + self._client = workflow_client or DaprWorkflowClient(host=config.host, port=config.port) + + def process(self, df: BatchDataFrameLike, batch_id: int) -> None: + """Schedules one Dapr Workflow execution per row in this micro-batch. + + Bounded by ``max_in_flight`` concurrent scheduling calls, iterating + ``df`` via ``toLocalIterator()`` rather than ``collect()`` so a + micro-batch never has to fit entirely in driver memory at once. + + If any record's outcome cannot be established as either newly + scheduled or already durably present, this raises so the + ``foreach_batch_sink`` call fails and Lakeflow retries the whole + micro-batch. Records already durably accepted in this same attempt + are unaffected by that retry: their deterministic instance IDs make + them idempotent to re-process. + + Args: + df: The micro-batch DataFrame Lakeflow passes to + ``foreach_batch_sink``. + batch_id: The Lakeflow micro-batch ID. Note that this resets to 0 + on a full pipeline refresh — see ``config.generation`` for + how this extension avoids treating that as a fresh identity + space by accident. + + Raises: + DaprDatabricksSinkError: One or more records could not be + durably scheduled, or the batch exceeded + ``max_records_per_batch``. + """ + limit = self._config.max_records_per_batch + record_count = 0 + failures: List[Exception] = [] + + with ThreadPoolExecutor(max_workers=self._config.max_in_flight) as pool: + futures: Dict[Future, int] = {} + for record_index, row in enumerate(self._iter_rows(df)): + if limit is not None and record_index >= limit: + raise DaprDatabricksSinkError( + f"batch {batch_id} for sink '{self._config.name}' has more than " + f'max_records_per_batch={limit} records. dapr.ext.databricks is ' + 'designed for business-action streams, not bulk data replication; ' + 'add upstream rate limiting (e.g. maxFilesPerTrigger/' + 'maxBytesPerTrigger on the source readStream), or raise ' + 'max_records_per_batch if this volume is intentional.' + ) + record_count += 1 + future = pool.submit(self._process_row, row, batch_id, record_index) + futures[future] = record_index + + for future in as_completed(futures): + try: + future.result() + except Exception as error: + failures.append(error) + + if failures: + raise DaprDatabricksSinkError( + f'{len(failures)}/{record_count} record(s) in batch {batch_id} for sink ' + f"'{self._config.name}' could not be durably scheduled; failing the " + 'micro-batch so Lakeflow retries it.' + ) from failures[0] + + def _iter_rows(self, df: BatchDataFrameLike) -> Iterable[RowLike]: + """Iterates ``df`` memory-safely, falling back to a bounded ``collect()`` + where ``toLocalIterator`` itself is unavailable. + + Some compute (observed on Databricks serverless / Spark Connect-backed + execution) raises directly from the ``toLocalIterator()`` call itself + — before any row is read — with "not supported when using file-based + collect". Nothing has been processed yet at that point, so falling + back is safe. The fallback still respects ``max_records_per_batch`` + as a hard cap via ``limit()`` when configured; without that cap it + logs a warning and collects the whole batch, since there is no other + memory-safe primitive to fall back to on such compute. + """ + try: + return df.toLocalIterator(prefetchPartitions=True) + except Exception as error: + if _TO_LOCAL_ITERATOR_UNSUPPORTED_MARKER not in str(error): + raise + + limit = self._config.max_records_per_batch + if limit is None: + _logger.warning( + 'dapr.ext.databricks: toLocalIterator() is unavailable in this Spark ' + 'environment; falling back to collect() with no record limit. Set ' + 'max_records_per_batch to bound driver memory use here.' + ) + bounded_df = df + else: + _logger.warning( + 'dapr.ext.databricks: toLocalIterator() is unavailable in this Spark ' + 'environment; falling back to a collect() bounded by ' + 'max_records_per_batch=%s.', + limit, + ) + bounded_df = df.limit(limit + 1) + return iter(bounded_df.collect()) + + def close(self) -> None: + """Closes the underlying Dapr Workflow client, if this handler created it.""" + if self._owns_client: + self._client.close() + + def _process_row(self, row: RowLike, batch_id: int, record_index: int) -> None: + start = time.monotonic() + try: + business_key = extract_business_key( + row, + batch_id, + id_field=self._config.id_field, + id_fields=self._config.id_fields, + instance_id_factory=self._config.instance_id_factory, + ) + instance_id = derive_instance_id( + namespace=self._config.namespace, + sink_name=self._config.name, + generation=self._config.generation, + business_key=business_key, + batch_id=batch_id, + record_index=record_index, + ) + payload = self._build_payload(row, batch_id, instance_id) + outcome = ensure_workflow_scheduled( + self._client, self._config.workflow, instance_id, payload + ) + except Exception: + _logger.exception( + 'dapr.ext.databricks: sink=%s workflow=%s batch_id=%s record_index=%s ' + 'failed to durably schedule workflow', + self._config.name, + self._config.workflow, + batch_id, + record_index, + ) + raise + + latency_ms = (time.monotonic() - start) * 1000 + _logger.info( + 'dapr.ext.databricks: sink=%s workflow=%s instance_id=%s batch_id=%s ' + 'namespace=%s generation=%s outcome=%s latency_ms=%.1f', + self._config.name, + self._config.workflow, + instance_id, + batch_id, + self._config.namespace, + self._config.generation, + 'newly_scheduled' if outcome.newly_scheduled else 'already_existed', + latency_ms, + ) + + def _build_payload(self, row: RowLike, batch_id: int, instance_id: str) -> Any: + data = self._config.row_mapper(row) + if not self._config.metadata: + return data + return { + 'data': data, + 'metadata': { + 'sink': self._config.name, + 'workflow': self._config.workflow, + 'batch_id': batch_id, + 'namespace': self._config.namespace, + 'generation': self._config.generation, + }, + } diff --git a/dapr/ext/databricks/config.py b/dapr/ext/databricks/config.py new file mode 100644 index 000000000..968a9b0ee --- /dev/null +++ b/dapr/ext/databricks/config.py @@ -0,0 +1,76 @@ +""" +Copyright 2026 The Dapr Authors +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 dataclasses import dataclass +from typing import Optional, Sequence + +from dapr.ext.databricks._typing import InstanceIdFactory, RowMapper +from dapr.ext.databricks.exceptions import SinkConfigurationError +from dapr.ext.databricks.mapping import default_row_mapper + + +@dataclass(frozen=True) +class WorkflowSinkConfig: + """Configuration for a single Lakeflow-to-Dapr-Workflow sink. + + See ``register_workflow_sink`` for the parameter descriptions this + mirrors, and ``dapr/ext/databricks/AGENTS.md`` for the delivery- and + full-refresh-semantics this configuration controls. + """ + + name: str + workflow: str + id_field: Optional[str] = None + id_fields: Optional[Sequence[str]] = None + namespace: str = 'default' + generation: str = 'v1' + input_mapper: Optional[RowMapper] = None + instance_id_factory: Optional[InstanceIdFactory] = None + metadata: bool = True + max_in_flight: int = 8 + max_records_per_batch: Optional[int] = None + host: Optional[str] = None + port: Optional[str] = None + + def __post_init__(self) -> None: + if not self.name: + raise SinkConfigurationError('name must be a non-empty string') + if not self.workflow: + raise SinkConfigurationError('workflow must be a non-empty string') + if not self.namespace: + raise SinkConfigurationError('namespace must be a non-empty string') + if not self.generation: + raise SinkConfigurationError('generation must be a non-empty string') + + key_strategies_given = sum( + strategy is not None + for strategy in (self.id_field, self.id_fields, self.instance_id_factory) + ) + if key_strategies_given > 1: + raise SinkConfigurationError( + 'specify at most one of id_field, id_fields, or instance_id_factory' + ) + if self.id_fields is not None and len(self.id_fields) == 0: + raise SinkConfigurationError('id_fields must not be empty') + + if self.max_in_flight < 1: + raise SinkConfigurationError('max_in_flight must be >= 1') + if self.max_records_per_batch is not None and self.max_records_per_batch < 1: + raise SinkConfigurationError('max_records_per_batch must be >= 1 when set') + + @property + def row_mapper(self) -> RowMapper: + """The effective row mapper: the configured ``input_mapper``, or the default.""" + return self.input_mapper or default_row_mapper diff --git a/dapr/ext/databricks/exceptions.py b/dapr/ext/databricks/exceptions.py new file mode 100644 index 000000000..04786c2b3 --- /dev/null +++ b/dapr/ext/databricks/exceptions.py @@ -0,0 +1,48 @@ +""" +Copyright 2026 The Dapr Authors +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 + + +class DaprDatabricksError(Exception): + """Base class for all errors raised by ``dapr.ext.databricks``.""" + + +class SinkConfigurationError(DaprDatabricksError, ValueError): + """Raised when ``register_workflow_sink`` (or ``WorkflowSinkConfig``) is misconfigured. + + This is a caller/programming error detected before any Spark row is + processed — e.g. specifying more than one of ``id_field``, ``id_fields``, + and ``instance_id_factory``. + """ + + +class MissingBusinessKeyError(DaprDatabricksError, ValueError): + """Raised when a configured ``id_field``/``id_fields`` is missing or null on a row. + + Treated as a malformed-record failure: it propagates out of batch + processing so the micro-batch fails and Lakeflow retries, rather than + silently falling back to a weaker, batch-position-based identity for a + record the caller explicitly said should be keyed by business ID. + """ + + +class DaprDatabricksSinkError(DaprDatabricksError, RuntimeError): + """Raised when a Lakeflow micro-batch cannot be durably handed off to Dapr Workflow. + + Any record for which durable acceptance (newly scheduled, or already + present in Dapr's workflow store) could not be established causes this to + be raised, which fails the ``foreach_batch_sink`` call and lets Lakeflow + retry the whole micro-batch. The original failure is chained via + ``__cause__``. + """ diff --git a/dapr/ext/databricks/identity.py b/dapr/ext/databricks/identity.py new file mode 100644 index 000000000..ee30607e8 --- /dev/null +++ b/dapr/ext/databricks/identity.py @@ -0,0 +1,187 @@ +""" +Copyright 2026 The Dapr Authors +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. + +Deterministic Dapr Workflow instance-ID derivation for the Databricks sink. + +This module is the crux of the sink's retry-safety: the same logical record +must always sanitize to the same instance ID, across processes and across +Lakeflow retries of the same micro-batch, so ``scheduling.ensure_workflow_scheduled`` +can recognize "already handled" via a plain lookup instead of guessing. +""" + +from __future__ import annotations + +import hashlib +import re +from typing import Optional, Sequence + +from dapr.ext.databricks._typing import InstanceIdFactory, RowLike +from dapr.ext.databricks.exceptions import MissingBusinessKeyError + +# Dapr workflow instance IDs may only contain alphanumeric characters, +# underscores, and dashes (see the Dapr Workflow "manage workflows" docs). +_CLEAN_SEGMENT_RE = re.compile(r'[A-Za-z0-9_-]+') + +# Dapr does not publish a maximum instance-ID length. These bounds are this +# extension's own conservative safety margin, keeping generated IDs short +# enough to stay friendly to every backing actor/state store Dapr might use. +_MAX_CLEAN_SEGMENT_LENGTH = 80 +_MAX_INSTANCE_ID_LENGTH = 128 +_HASH_LENGTH = 32 + + +def sanitize_segment(raw: str) -> str: + """Returns a Dapr-instance-ID-safe segment derived deterministically from ``raw``. + + A segment that is already short and composed entirely of + ``[A-Za-z0-9_-]`` is returned unchanged, so common business keys (order + numbers, UUIDs, account IDs) stay human-readable in the resulting + instance ID. Anything else — Unicode, punctuation, empty strings, or + values over ``_MAX_CLEAN_SEGMENT_LENGTH`` chars — is replaced by a SHA-256 + hex digest of its UTF-8 encoding. Hashing the *whole* segment (rather than + stripping individual bad characters) avoids collisions that a + character-by-character replacement could introduce, e.g. between "a/b" + and "a-b" once both map to the same cleaned output. + + Args: + raw: The unsanitized segment value. + + Returns: + A non-empty string containing only alphanumeric characters, + underscores, and dashes. + """ + is_clean = ( + len(raw) <= _MAX_CLEAN_SEGMENT_LENGTH and _CLEAN_SEGMENT_RE.fullmatch(raw) is not None + ) + if is_clean: + return raw + return hashlib.sha256(raw.encode('utf-8')).hexdigest()[:_HASH_LENGTH] + + +def _extract_single_field(row: RowLike, field: str) -> str: + """Reads one required, non-null field from a row as a string business key.""" + row_dict = row.asDict() + if field not in row_dict: + raise MissingBusinessKeyError(f"id field '{field}' is not present on this row") + value = row_dict[field] + if value is None: + raise MissingBusinessKeyError(f"id field '{field}' is null on this row") + return str(value) + + +def extract_business_key( + row: RowLike, + batch_id: int, + *, + id_field: Optional[str], + id_fields: Optional[Sequence[str]], + instance_id_factory: Optional[InstanceIdFactory], +) -> Optional[str]: + """Derives the business-key component of a workflow instance ID from a row. + + Exactly one of ``id_field``, ``id_fields``, or ``instance_id_factory`` is + expected to be set (``WorkflowSinkConfig`` enforces this at configuration + time); if none are set, ``None`` is returned and the caller falls back to + batch/record-position identity. + + Args: + row: The Spark row for a single record. + batch_id: The Lakeflow micro-batch ID, passed through to a custom + ``instance_id_factory``. + id_field: Single business-key column name, if configured. + id_fields: Composite business-key column names, if configured. + instance_id_factory: Escape hatch that computes the business-key + component directly from the row; still namespaced, sanitized, and + composed with ``generation`` like any other business key. + + Returns: + The raw (not yet sanitized) business-key string, or ``None`` if no + business-key strategy is configured. + + Raises: + MissingBusinessKeyError: A configured ``id_field``/``id_fields`` + column is absent or null on this row. + """ + if instance_id_factory is not None: + return str(instance_id_factory(row, batch_id)) + if id_field is not None: + return _extract_single_field(row, id_field) + if id_fields: + key_components = [_extract_single_field(row, field) for field in id_fields] + return '_'.join(key_components) + return None + + +def derive_instance_id( + *, + namespace: str, + sink_name: str, + generation: str, + business_key: Optional[str], + batch_id: int, + record_index: int, +) -> str: + """Composes the deterministic Dapr Workflow instance ID for one record. + + The template is ``---`` when a + business key is available, or ``----`` + otherwise. ``generation`` is always part of the identity: this is what + lets a full pipeline refresh (which restarts ``batch_id`` from 0 and can + replay already-handled business keys) be given a fresh, non-colliding + identity space simply by bumping ``generation`` — see the module README + for the full explanation of full-refresh semantics. + + Args: + namespace: Logical partition for this sink's workflow identities + (e.g. a business domain like ``"orders"``). + sink_name: The registered sink name. + generation: Identity epoch. Bump this to intentionally replay + business actions after a full refresh; keep it stable for normal + retries to dedupe correctly. + business_key: Pre-extracted business key, or ``None`` to fall back to + batch/record-position identity. + batch_id: The Lakeflow micro-batch ID. + record_index: The record's 0-based position within the micro-batch, + used only when ``business_key`` is ``None``. + + Returns: + A stable, Dapr-instance-ID-safe string: the same logical record + always produces the same output. + """ + namespace_segment = sanitize_segment(namespace) + sink_segment = sanitize_segment(sink_name) + generation_segment = sanitize_segment(generation) + + if business_key is not None: + key_segment = sanitize_segment(business_key) + instance_id = f'{namespace_segment}-{sink_segment}-{generation_segment}-{key_segment}' + else: + instance_id = ( + f'{namespace_segment}-{sink_segment}-{generation_segment}-{batch_id}-{record_index}' + ) + + if len(instance_id) <= _MAX_INSTANCE_ID_LENGTH: + return instance_id + + # Individually-clean segments can still add up to an over-long ID (e.g. a + # long but valid namespace/sink pair). Collapse deterministically instead + # of naively truncating the composed string, which could chop off the + # record-specific suffix entirely and collide every record in the sink + # onto the same truncated prefix. The digest (derived from the full, + # per-record instance_id) is always kept intact; only the human-readable + # prefix is shortened to make room. + digest = hashlib.sha256(instance_id.encode('utf-8')).hexdigest()[:_HASH_LENGTH] + prefix = f'{namespace_segment}-{sink_segment}-{generation_segment}' + max_prefix_length = _MAX_INSTANCE_ID_LENGTH - len(digest) - 1 # 1 for the joining '-' + if max_prefix_length <= 0: + return digest + return f'{prefix[:max_prefix_length]}-{digest}' diff --git a/dapr/ext/databricks/mapping.py b/dapr/ext/databricks/mapping.py new file mode 100644 index 000000000..25b374c22 --- /dev/null +++ b/dapr/ext/databricks/mapping.py @@ -0,0 +1,60 @@ +""" +Copyright 2026 The Dapr Authors +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 +from datetime import date, datetime +from decimal import Decimal +from typing import Any, Dict + +from dapr.ext.databricks._typing import RowLike + + +def default_row_mapper(row: RowLike) -> Dict[str, Any]: + """Default ``Row`` -> JSON-compatible ``dict`` mapping for workflow input. + + Delegates field extraction to ``Row.asDict(recursive=True)`` (so nested + structs/arrays/maps become plain dicts/lists), then normalizes the value + types that ``asDict`` leaves in a non-JSON-serializable shape. + + Args: + row: The Spark row for a single record. + + Returns: + A JSON-serializable dict suitable for use as workflow input. + """ + return _json_safe(row.asDict(recursive=True)) + + +def _json_safe(value: Any) -> Any: + """Recursively converts Spark/Python values that ``json.dumps`` cannot handle. + + - ``datetime``/``date`` -> ISO 8601 string. + - ``Decimal`` -> string, to avoid silently losing precision on monetary + fields by round-tripping through ``float``. + - ``bytes``/``bytearray`` -> base64-encoded string. + - ``dict``/``list``/``tuple`` -> recursed into. + - Anything else is passed through unchanged. + """ + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, (datetime, date)): + return value.isoformat() + if isinstance(value, Decimal): + return str(value) + if isinstance(value, (bytes, bytearray)): + return base64.b64encode(bytes(value)).decode('ascii') + return value diff --git a/dapr/ext/databricks/py.typed b/dapr/ext/databricks/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/dapr/ext/databricks/scheduling.py b/dapr/ext/databricks/scheduling.py new file mode 100644 index 000000000..e6980877c --- /dev/null +++ b/dapr/ext/databricks/scheduling.py @@ -0,0 +1,117 @@ +""" +Copyright 2026 The Dapr Authors +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. + +Retry-safe "ensure scheduled" logic — the primary recovery property this +extension provides. See ``dapr/ext/databricks/AGENTS.md`` for the full +delivery-semantics writeup this implements. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import grpc + +from dapr.ext.workflow.dapr_workflow_client import DaprWorkflowClient + +# Dapr's workflow "start instance" call rejects an instance ID that is +# already in use with an HTTP 409 / gRPC ALREADY_EXISTS-shaped error (see the +# Dapr Workflow API reference: "A workflow with the given instance ID already +# exists and is not yet reusable"). Match on the status code first; fall back +# to the message text since this SDK's own client already relies on message +# matching elsewhere (DaprWorkflowClient.get_workflow_state, for "no such +# instance exists") rather than assuming every backend/transport surfaces the +# same status code for this condition. +_DUPLICATE_INSTANCE_MESSAGE_MARKER = 'already exists' + + +@dataclass(frozen=True) +class ScheduleOutcome: + """Result of ensuring a single workflow instance is durably scheduled.""" + + instance_id: str + newly_scheduled: bool + + +def is_duplicate_instance_error(error: grpc.RpcError) -> bool: + """Reports whether ``error`` represents a duplicate-instance-ID conflict. + + Args: + error: The gRPC error raised by ``schedule_new_workflow``. + + Returns: + True if this error means "an instance with this ID already exists" + rather than a genuine transport/auth/validation failure. + """ + if error.code() == grpc.StatusCode.ALREADY_EXISTS: + return True + details = error.details() or '' + return _DUPLICATE_INSTANCE_MESSAGE_MARKER in details.lower() + + +def ensure_workflow_scheduled( + client: DaprWorkflowClient, + workflow: str, + instance_id: str, + payload: Any, +) -> ScheduleOutcome: + """Schedules ``workflow`` at ``instance_id`` unless it is already durably present. + + This is a check-then-act sequence, not a transaction — Dapr Workflow has + no compare-and-swap "schedule if absent" primitive — so it relies on + ``instance_id`` being deterministic for the same logical record and on + Dapr rejecting a second ``schedule_new_workflow`` for an ID that is + already in use: + + 1. ``get_workflow_state`` — if the instance already exists (in any + status), the record was already handled by a prior attempt; treat it + as durably accepted and do nothing further. + 2. Otherwise, call ``schedule_new_workflow``. If Dapr reports the ID + already exists (``is_duplicate_instance_error``), a concurrent caller + won the race between step 1 and step 2; treat it the same as a normal + existing-instance discovery. Any other error (unavailable, timeout, + auth failure, throttling, malformed input, invalid workflow name) + propagates so the caller can fail the micro-batch. + + A failure of ``schedule_new_workflow`` itself (e.g. a network timeout + *after* Dapr already durably accepted the instance) also propagates here + — this function cannot distinguish that from a genuine failure on this + attempt. Recovery happens on the next Lakeflow retry of the same record, + whose ``get_workflow_state`` call in step 1 will find the instance Dapr + already accepted. + + Args: + client: A connected ``DaprWorkflowClient``. + workflow: Registered workflow name to schedule. + instance_id: Deterministic instance ID for this record. + payload: JSON-serializable workflow input. + + Returns: + A ``ScheduleOutcome`` recording whether this call newly scheduled the + instance or found it already durably present. + + Raises: + grpc.RpcError: Any non-duplicate-instance scheduling failure. + """ + existing_state = client.get_workflow_state(instance_id, fetch_payloads=False) + if existing_state is not None: + return ScheduleOutcome(instance_id=instance_id, newly_scheduled=False) + + try: + client.schedule_new_workflow(workflow, input=payload, instance_id=instance_id) + except grpc.RpcError as error: + if is_duplicate_instance_error(error): + return ScheduleOutcome(instance_id=instance_id, newly_scheduled=False) + raise + + return ScheduleOutcome(instance_id=instance_id, newly_scheduled=True) diff --git a/dapr/ext/databricks/sink.py b/dapr/ext/databricks/sink.py new file mode 100644 index 000000000..b467545e4 --- /dev/null +++ b/dapr/ext/databricks/sink.py @@ -0,0 +1,171 @@ +""" +Copyright 2026 The Dapr Authors +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 Optional, Sequence + +from dapr.ext.databricks._typing import InstanceIdFactory, RowMapper +from dapr.ext.databricks.batch_handler import DaprWorkflowBatchHandler +from dapr.ext.databricks.config import WorkflowSinkConfig + + +def _pipelines_module(): + """Lazily imports ``pyspark.pipelines``, with a clear error outside Lakeflow. + + ``pyspark.pipelines`` is provided by the Databricks Lakeflow (Spark + Declarative Pipelines) runtime, not by a plain ``pip install pyspark``. + Importing it at module scope would either force every ``dapr`` user to + have a compatible pyspark on their path, or fail import of this whole + extension outside a Lakeflow pipeline — so the import happens here, + on first use, instead. + """ + try: + from pyspark import pipelines + except ImportError as error: + raise ImportError( + 'Databricks Lakeflow support requires pyspark.pipelines.foreach_batch_sink. ' + 'Run this integration inside a supported Databricks Lakeflow pipeline.' + ) from error + return pipelines + + +def register_workflow_sink( + name: str, + workflow: str, + *, + id_field: Optional[str] = None, + id_fields: Optional[Sequence[str]] = None, + namespace: str = 'default', + generation: str = 'v1', + input_mapper: Optional[RowMapper] = None, + instance_id_factory: Optional[InstanceIdFactory] = None, + metadata: bool = True, + max_in_flight: int = 8, + max_records_per_batch: Optional[int] = None, + host: Optional[str] = None, + port: Optional[str] = None, +) -> DaprWorkflowBatchHandler: + """Registers a Lakeflow sink that schedules a Dapr Workflow for each streaming record. + + Internally registers a ``pyspark.pipelines.foreach_batch_sink`` named + ``name``; reference it as a flow ``target`` the same way you would any + other sink: + + ```python + from pyspark import pipelines as dp + from dapr.ext.databricks import register_workflow_sink + + register_workflow_sink( + name='order_actions', + workflow='process_order', + id_field='order_id', + namespace='orders', + ) + + @dp.append_flow(target='order_actions', name='order_actions_flow') + def order_actions_flow(): + return spark.readStream.table('validated_orders') + ``` + + This only waits for Dapr to durably accept each workflow instance, not + for the workflow to finish — the whole point is to decouple fast + streaming ingestion from potentially long-running business processes. + See ``dapr/ext/databricks/AGENTS.md`` (or the extension README) for the + full delivery-semantics and full-refresh writeup; the summary: + + - Instance IDs are derived deterministically from ``namespace``, `name`, + ``generation``, and the record's business key — never randomly — so a + Lakeflow retry of a micro-batch recognizes work it already handed off + instead of scheduling a duplicate execution. + - This is retry-safe, not exactly-once: it depends on Dapr's workflow + instance/history retention. Once a workflow instance is purged, its ID + is no longer recognized as "already handled". + - ``batch_id`` resets to 0 on a full pipeline refresh, so it is never + used alone as identity. Bump ``generation`` to intentionally replay + business actions after a full refresh; leave it unchanged for normal + retries to keep deduplicating correctly. + + Args: + name: Unique sink name within the pipeline (passed to + ``foreach_batch_sink``). + workflow: Registered Dapr Workflow name to schedule for each record. + id_field: Row column to use as the business key. Mutually exclusive + with ``id_fields`` and ``instance_id_factory``. + id_fields: Row columns to combine into a composite business key. + Mutually exclusive with ``id_field`` and ``instance_id_factory``. + namespace: Logical partition for this sink's workflow identities + (e.g. a business domain). Part of the deterministic instance ID. + generation: Identity epoch, also part of the deterministic instance + ID. Keep stable across normal retries; change it to deliberately + replay business actions (e.g. after a full pipeline refresh). + Defaults to ``'v1'`` so normal operation never accidentally + replays already-handled records. + input_mapper: Optional ``Row -> dict`` mapper for workflow input. + Defaults to a JSON-safe ``Row.asDict(recursive=True)``. + instance_id_factory: Escape hatch: ``(row, batch_id) -> str`` + computing the business-key component directly. Still namespaced + by ``namespace``/``name``/``generation`` and sanitized like any + other business key. Mutually exclusive with ``id_field`` and + ``id_fields``. + metadata: When ``True`` (default), wraps workflow input as + ``{'data': , 'metadata': {...sink/workflow/batch/namespace/generation...}}``. + When ``False``, the workflow input is exactly the mapped row, + with no wrapper. + max_in_flight: Maximum concurrent ``schedule_new_workflow`` calls per + micro-batch. + max_records_per_batch: Optional hard cap on records processed per + micro-batch. Exceeding it fails the batch outright (see + ``DaprWorkflowBatchHandler.process``) rather than silently + dropping records — this integration targets business-action + streams, not bulk data replication. + host: Dapr sidecar/endpoint gRPC host. Defaults to the standard Dapr + SDK environment/settings resolution (``DAPR_GRPC_ENDPOINT``, + ``DAPR_RUNTIME_HOST``). Must be network-reachable from the + Databricks compute running the pipeline. + port: Dapr sidecar/endpoint gRPC port. Defaults to + ``DAPR_GRPC_PORT``. + + Returns: + The ``DaprWorkflowBatchHandler`` backing the registered sink, for + advanced use (e.g. sharing it, or calling ``.close()`` explicitly). + + Raises: + SinkConfigurationError: Invalid configuration (e.g. more than one of + ``id_field``/``id_fields``/``instance_id_factory`` given). + ImportError: Called outside a Databricks Lakeflow pipeline. + """ + pipelines = _pipelines_module() + + config = WorkflowSinkConfig( + name=name, + workflow=workflow, + id_field=id_field, + id_fields=id_fields, + namespace=namespace, + generation=generation, + input_mapper=input_mapper, + instance_id_factory=instance_id_factory, + metadata=metadata, + max_in_flight=max_in_flight, + max_records_per_batch=max_records_per_batch, + host=host, + port=port, + ) + handler = DaprWorkflowBatchHandler(config) + + @pipelines.foreach_batch_sink(name=name) + def _dapr_workflow_sink(df, batch_id): + handler.process(df, batch_id) + + return handler diff --git a/examples/AGENTS.md b/examples/AGENTS.md index f098a2b10..6c0839bfc 100644 --- a/examples/AGENTS.md +++ b/examples/AGENTS.md @@ -126,6 +126,17 @@ The `workflow` example includes: `simple.py`, `task_chaining.py`, `fan_out_fan_i | `conversation` | Standalone client | `dapr` (base, uses sidecar) | No (uses config/) | | `langgraph-checkpointer` | Standalone gRPC server | `dapr[langgraph]`, LangGraph, LangChain | Yes | +### Data pipelines +| Example | Pattern | SDK packages | Has components | +|---------|---------|-------------|----------------| +| `databricks` | Workflow app runnable locally + illustrative Lakeflow pipeline file | `dapr[workflow,databricks]` | No | + +`databricks/fraud_remediation_workflow.py` is the runnable half (workflow + activities, +simulates a Lakeflow micro-batch in-process so it needs no pyspark/Databricks). Its counterpart +`databricks/fraud_remediation_pipeline.py` shows the real `pyspark.pipelines` wiring but is +illustrative only — it is not executed by `tests/examples/` since it requires an actual +Databricks Lakeflow runtime. + ## Adding a new example 1. Create a directory under `examples/` with a descriptive kebab-case name diff --git a/examples/databricks/README.md b/examples/databricks/README.md new file mode 100644 index 000000000..cea68d913 --- /dev/null +++ b/examples/databricks/README.md @@ -0,0 +1,87 @@ +# Databricks Lakeflow -> Dapr Workflow: fraud remediation + +Demonstrates `dapr.ext.databricks`: turning records from a Databricks Lakeflow +streaming pipeline into durable Dapr Workflow executions. + +```text +Databricks Lakeflow identifies a suspicious transaction + | + dapr.ext.databricks sink + | + FraudRemediationWorkflow + | + freeze card -> notify customer -> open investigation + | + wait for analyst decision + | + resolve case +``` + +## Files + +- **`fraud_remediation_workflow.py`** -- the Dapr Workflow side: the + `fraud_remediation` workflow and its activities (freeze card, notify + customer, open investigation, wait for an analyst decision, resolve the + case). Runnable locally via `dapr run`, no Databricks or pyspark required: + its `__main__` block simulates a small Lakeflow micro-batch in-process + using `DaprWorkflowBatchHandler` directly -- the same reusable building + block `register_workflow_sink` uses internally -- and then retries the + identical batch to demonstrate that no duplicate workflow executions are + created. +- **`fraud_remediation_pipeline.py`** -- the Databricks Lakeflow side: + `register_workflow_sink` plus the `@dp.append_flow` that feeds it. This + file is illustrative only. It is not executed by this repository's test + suite and cannot run locally -- it depends on `pyspark.pipelines` (provided + by the Databricks Lakeflow runtime, not `pip install pyspark`), the + implicit `spark` session Lakeflow injects into pipeline source files, and a + Unity Catalog table named `detected_fraud`. Copy its pattern into an actual + Databricks Lakeflow pipeline to wire this integration up for real. + +See [`dapr/ext/databricks/README.md`](../../dapr/ext/databricks/README.md) +for the full public API, delivery-semantics, and full-refresh writeup. + +## Prerequisites + +- [Dapr CLI and initialized environment](https://docs.dapr.io/getting-started) +- [Install Python 3.10+](https://www.python.org/downloads/) +- `pip install "dapr[workflow,databricks]"` (this repo's dev environment + already has both via `uv sync --all-packages --group dev`) + +## Run the example + +```sh +dapr run --app-id fraud-remediation-demo -- python3 fraud_remediation_workflow.py +``` + +Expected output (interleaved with Dapr/durabletask log lines): + +```text +*** Lakeflow micro-batch 1: scheduling fraud remediation workflows +*** Micro-batch 1 handed off; each workflow now runs independently +*** Simulating a Lakeflow retry of the same micro-batch +*** Retry complete: no duplicate workflow executions were created +*** Triggered by sink 'fraud_actions', batch 1 +*** Freezing card for transaction T-1001 (customer C-1) +*** Notifying customer C-1 about the frozen card +*** Opened investigation CASE-T-1001 +*** Resolved CASE-T-1001: CONFIRMED_FRAUD +*** Workflow fraud-fraud_actions-v1-T-1001 completed: {"case_id": "CASE-T-1001", "decision": "CONFIRMED_FRAUD"} +... (and the same for T-1002) +``` + +Look for the structured `dapr.ext.databricks: sink=... outcome=...` log +lines: the first micro-batch reports `outcome=newly_scheduled` for both +transactions, and the simulated retry reports `outcome=already_existed` for +both -- the same transaction never starts a second remediation workflow. + +The example purges its two demo workflow instances on startup so repeated +local runs behave like a fresh run; a real Lakeflow pipeline would never do +this; see the "Delivery semantics" section of the extension README for why +reusing an instance ID on purpose depends on Dapr's workflow history +retention. + +## Cleanup + +```sh +dapr stop --app-id fraud-remediation-demo +``` diff --git a/examples/databricks/fraud_remediation_pipeline.py b/examples/databricks/fraud_remediation_pipeline.py new file mode 100644 index 000000000..fc3845c3e --- /dev/null +++ b/examples/databricks/fraud_remediation_pipeline.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- +# Copyright 2026 The Dapr Authors +# 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. + +"""Fraud remediation: the Databricks Lakeflow side of the dapr.ext.databricks example. + +Illustrative only -- this file is meant to be a source file inside a real +Databricks Lakeflow pipeline. It is NOT executed by this repository's test +suite and cannot run locally: it needs `pyspark.pipelines` (provided by the +Databricks Lakeflow runtime, not by `pip install pyspark`), the implicit +`spark` session Lakeflow injects into pipeline source files, and a Unity +Catalog table named `detected_fraud`. + +See `fraud_remediation_workflow.py` in this same directory for the Dapr +Workflow side (freeze card / notify customer / open investigation / wait for +analyst / resolve case) -- that file *is* runnable locally via `dapr run` and +is covered by `tests/examples/test_databricks.py`. +""" + +from pyspark import pipelines as dp + +from dapr.ext.databricks import register_workflow_sink + +# Deterministic instance IDs (---) +# mean a Lakeflow retry of a micro-batch recognizes transactions it already +# handed off instead of starting a second remediation. See this repository's +# dapr/ext/databricks/README.md and AGENTS.md for the full delivery-semantics +# and full-refresh (generation) writeup -- in particular, why `generation` +# defaults to a fixed value instead of changing on every full pipeline +# refresh, and what to do if you deliberately want to replay fraud actions +# after one. +register_workflow_sink( + name='fraud_actions', + workflow='fraud_remediation', + id_field='transaction_id', + namespace='fraud', +) + + +@dp.append_flow( + target='fraud_actions', + name='fraud_action_flow', +) +def fraud_action_flow(): + # `spark` is injected into pipeline source files by the Lakeflow runtime; + # it is not something this file imports or defines. + return spark.readStream.table('detected_fraud') # noqa: F821 diff --git a/examples/databricks/fraud_remediation_workflow.py b/examples/databricks/fraud_remediation_workflow.py new file mode 100644 index 000000000..5e1dc791d --- /dev/null +++ b/examples/databricks/fraud_remediation_workflow.py @@ -0,0 +1,211 @@ +# -*- coding: utf-8 -*- +# Copyright 2026 The Dapr Authors +# 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. + +"""Fraud remediation: the Dapr Workflow side of the dapr.ext.databricks example. + + Databricks Lakeflow identifies a suspicious transaction + | + dapr.ext.databricks sink + | + FraudRemediationWorkflow <-- this file + | + freeze card -> notify customer -> open investigation + | + wait for analyst decision + | + resolve case + +This file is runnable on its own via `dapr run` (see the README in this +directory) and does not require Databricks or pyspark: it simulates a small +Lakeflow micro-batch in-process using `DaprWorkflowBatchHandler` directly, +the same reusable building block `register_workflow_sink` uses internally. +`fraud_remediation_pipeline.py`, alongside this file, shows the real +`pyspark.pipelines` wiring you'd use inside an actual Databricks Lakeflow +pipeline instead of this simulation. +""" + +import logging +import threading +import time +from datetime import timedelta +from typing import Any, Dict, Iterator + +import dapr.ext.workflow as wf +from dapr.ext.databricks import DaprWorkflowBatchHandler, WorkflowSinkConfig + +wfr = wf.WorkflowRuntime() + +ANALYST_DECISION_EVENT = 'analyst_decision' +ANALYST_TIMEOUT = timedelta(seconds=20) + + +@wfr.workflow(name='fraud_remediation') +def fraud_remediation_workflow(ctx: wf.DaprWorkflowContext, envelope: Dict[str, Any]): + # register_workflow_sink's default `metadata=True` wraps the mapped row as + # {'data': ..., 'metadata': {'sink', 'workflow', 'batch_id', 'namespace', + # 'generation'}}, so the workflow can see which Lakeflow sink/batch/ + # generation triggered it without that bookkeeping polluting `data`, which + # stays exactly what `input_mapper` (or the default row mapper) produced. + transaction = envelope['data'] + if not ctx.is_replaying: + print( + f"*** Triggered by sink '{envelope['metadata']['sink']}', " + f'batch {envelope["metadata"]["batch_id"]}' + ) + + yield ctx.call_activity(freeze_card, input=transaction) + yield ctx.call_activity(notify_customer, input=transaction) + case_id = yield ctx.call_activity(open_investigation, input=transaction) + + decision_event = ctx.wait_for_external_event(ANALYST_DECISION_EVENT) + timeout_event = ctx.create_timer(ANALYST_TIMEOUT) + winner = yield wf.when_any([decision_event, timeout_event]) + decision = decision_event.get_result() if winner == decision_event else 'ESCALATED_NO_RESPONSE' + + yield ctx.call_activity(resolve_case, input={'case_id': case_id, 'decision': decision}) + return {'case_id': case_id, 'decision': decision} + + +@wfr.activity(name='freeze_card') +def freeze_card(_, transaction: Dict[str, Any]) -> None: + print( + f'*** Freezing card for transaction {transaction["transaction_id"]} ' + f'(customer {transaction["customer_id"]})' + ) + + +@wfr.activity(name='notify_customer') +def notify_customer(_, transaction: Dict[str, Any]) -> None: + print(f'*** Notifying customer {transaction["customer_id"]} about the frozen card') + + +@wfr.activity(name='open_investigation') +def open_investigation(_, transaction: Dict[str, Any]) -> str: + case_id = f'CASE-{transaction["transaction_id"]}' + print(f'*** Opened investigation {case_id}') + return case_id + + +@wfr.activity(name='resolve_case') +def resolve_case(_, resolution: Dict[str, Any]) -> None: + print(f'*** Resolved {resolution["case_id"]}: {resolution["decision"]}') + + +class _SimulatedLakeflowRow: + """Stands in for a `pyspark.sql.Row` so this example runs without pyspark. + + A real Lakeflow micro-batch delivers actual `pyspark.sql.Row` objects; + `DaprWorkflowBatchHandler` only ever calls `asDict()`/`__getitem__` on + them (see `dapr.ext.databricks._typing.RowLike`), so this minimal stand-in + is enough to demonstrate the real code path end-to-end. + """ + + def __init__(self, **fields: Any) -> None: + self._fields = fields + + def asDict(self, recursive: bool = False) -> Dict[str, Any]: + return dict(self._fields) + + def __getitem__(self, key: str) -> Any: + return self._fields[key] + + +class _SimulatedLakeflowBatch: + """Stands in for a `pyspark.sql.DataFrame` micro-batch; see the class above.""" + + def __init__(self, rows) -> None: + self._rows = rows + + def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Any]: + return iter(self._rows) + + +def _instance_id_for(transaction_id: str) -> str: + """The same deterministic template `dapr.ext.databricks` derives internally: + `---`. Recomputing it here (rather + than importing the internal `identity` module) is exactly what the + documented, stable ID scheme is for: callers can predict an instance ID + without having scheduled it themselves. + """ + return f'fraud-fraud_actions-v1-{transaction_id}' + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO, format='%(message)s') + + wfr.start() + wfr.wait_for_worker_ready() + + handler = DaprWorkflowBatchHandler( + WorkflowSinkConfig( + name='fraud_actions', + workflow='fraud_remediation', + id_field='transaction_id', + namespace='fraud', + ) + ) + + detected_fraud_batch = [ + _SimulatedLakeflowRow(transaction_id='T-1001', customer_id='C-1', amount=4200.00), + _SimulatedLakeflowRow(transaction_id='T-1002', customer_id='C-2', amount=999.50), + ] + + # Best-effort cleanup so re-running this demo script against the same + # Dapr instance behaves like a fresh run instead of finding yesterday's + # instances "already existed". A real Lakeflow pipeline would never do + # this -- reusing an instance ID on purpose is exactly the durable-history + # dependency this extension documents (see the README/AGENTS.md). + wf_client = wf.DaprWorkflowClient() + for row in detected_fraud_batch: + try: + wf_client.purge_workflow(_instance_id_for(row['transaction_id'])) + except Exception: + pass + + # This call only waits for Dapr to durably accept each workflow instance + # -- not for fraud_remediation_workflow to finish running. It returns as + # soon as scheduling is confirmed, which is what keeps a real Lakeflow + # micro-batch fast regardless of how long remediation takes. + print('*** Lakeflow micro-batch 1: scheduling fraud remediation workflows') + handler.process(_SimulatedLakeflowBatch(detected_fraud_batch), batch_id=1) + print('*** Micro-batch 1 handed off; each workflow now runs independently') + + # Simulates Lakeflow retrying the exact same micro-batch (e.g. after a + # worker restart). Deterministic instance IDs mean this must not start a + # second remediation for either transaction. + print('*** Simulating a Lakeflow retry of the same micro-batch') + handler.process(_SimulatedLakeflowBatch(detected_fraud_batch), batch_id=1) + print('*** Retry complete: no duplicate workflow executions were created') + + def _resolve_after_delay() -> None: + time.sleep(2) + for row in detected_fraud_batch: + wf_client.raise_workflow_event( + _instance_id_for(row['transaction_id']), + ANALYST_DECISION_EVENT, + data='CONFIRMED_FRAUD', + ) + + threading.Thread(target=_resolve_after_delay, daemon=True).start() + + for row in detected_fraud_batch: + instance_id = _instance_id_for(row['transaction_id']) + state = wf_client.wait_for_workflow_completion(instance_id, timeout_in_seconds=30) + if state and state.runtime_status.name == 'COMPLETED': + print(f'*** Workflow {instance_id} completed: {state.serialized_output}') + else: + status = state.runtime_status.name if state else 'NOT_FOUND' + print(f'*** Workflow {instance_id} ended with status: {status}') + + wf_client.close() + handler.close() + wfr.shutdown() diff --git a/pyproject.toml b/pyproject.toml index 68264315d..77fed5fa0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,15 @@ strands = [ # exists so users get a consistent `pip install dapr[]` shape and so the # extra is reserved if workflow ever grows external deps. workflow = [] -all = ["dapr[fastapi,flask,grpc,langgraph,strands,workflow]"] +# databricks is also empty on purpose: it schedules Dapr Workflow executions +# (via dapr.ext.workflow, already in core) from inside a Databricks Lakeflow +# pipeline. pyspark.pipelines is provided by that Databricks runtime itself — +# depending on a pip-installable "pyspark" here would be both unnecessary for +# real usage and insufficient for local imports, since pyspark.pipelines is +# not part of the open-source PyPI pyspark package. dapr.ext.databricks +# imports it lazily and raises a clear error if it's unavailable. +databricks = [] +all = ["dapr[fastapi,flask,grpc,langgraph,strands,workflow,databricks]"] [project.urls] Documentation = "https://github.com/dapr/docs" @@ -189,6 +197,14 @@ ignore_missing_imports = true module = ["langgraph.*", "langchain.*", "strands.*", "strands_agents.*"] ignore_missing_imports = true +# pyspark (and pyspark.pipelines specifically) is provided by the Databricks +# Lakeflow runtime, not a project dependency — see the `databricks` extra +# above. dapr.ext.databricks imports it lazily at call time; this override +# only lets mypy check the TYPE_CHECKING-only references to it. +[[tool.mypy.overrides]] +module = ["pyspark.*"] +ignore_missing_imports = true + # Bundling the extensions into the core wheel brings their source under mypy's # namespace walk for the first time. These modules carry pre-existing type # errors that were previously hidden by the separate-workspace-package layout. diff --git a/tests/examples/test_databricks.py b/tests/examples/test_databricks.py new file mode 100644 index 000000000..4dc6ab0d1 --- /dev/null +++ b/tests/examples/test_databricks.py @@ -0,0 +1,37 @@ +import pytest + +EXPECTED_FRAUD_REMEDIATION = [ + '*** Lakeflow micro-batch 1: scheduling fraud remediation workflows', + 'dapr.ext.databricks: sink=fraud_actions workflow=fraud_remediation instance_id=fraud-fraud_actions-v1-T-1001 batch_id=1 namespace=fraud generation=v1 outcome=newly_scheduled', + 'dapr.ext.databricks: sink=fraud_actions workflow=fraud_remediation instance_id=fraud-fraud_actions-v1-T-1002 batch_id=1 namespace=fraud generation=v1 outcome=newly_scheduled', + '*** Micro-batch 1 handed off; each workflow now runs independently', + '*** Simulating a Lakeflow retry of the same micro-batch', + 'dapr.ext.databricks: sink=fraud_actions workflow=fraud_remediation instance_id=fraud-fraud_actions-v1-T-1001 batch_id=1 namespace=fraud generation=v1 outcome=already_existed', + 'dapr.ext.databricks: sink=fraud_actions workflow=fraud_remediation instance_id=fraud-fraud_actions-v1-T-1002 batch_id=1 namespace=fraud generation=v1 outcome=already_existed', + '*** Retry complete: no duplicate workflow executions were created', + '*** Freezing card for transaction T-1001 (customer C-1)', + '*** Freezing card for transaction T-1002 (customer C-2)', + '*** Notifying customer C-1 about the frozen card', + '*** Notifying customer C-2 about the frozen card', + '*** Opened investigation CASE-T-1001', + '*** Opened investigation CASE-T-1002', + '*** Resolved CASE-T-1001: CONFIRMED_FRAUD', + '*** Resolved CASE-T-1002: CONFIRMED_FRAUD', + '*** Workflow fraud-fraud_actions-v1-T-1001 completed: {"case_id": "CASE-T-1001", "decision": "CONFIRMED_FRAUD"}', + '*** Workflow fraud-fraud_actions-v1-T-1002 completed: {"case_id": "CASE-T-1002", "decision": "CONFIRMED_FRAUD"}', +] + + +@pytest.mark.example_dir('databricks') +def test_fraud_remediation_workflow(dapr): + output = dapr.run( + '--app-id fraud-remediation-demo -- python3 fraud_remediation_workflow.py', + timeout=60, + ) + for line in EXPECTED_FRAUD_REMEDIATION: + assert line in output, f'Missing in output: {line}' + + # The whole point of the sink is to schedule-and-move-on: scheduling must + # never be reported as newly_scheduled a second time for the retried batch. + assert output.count('outcome=newly_scheduled') == 2 + assert output.count('outcome=already_existed') == 2 diff --git a/tests/ext/databricks/__init__.py b/tests/ext/databricks/__init__.py new file mode 100644 index 000000000..8fb63c69c --- /dev/null +++ b/tests/ext/databricks/__init__.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +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/tests/ext/databricks/_fakes.py b/tests/ext/databricks/_fakes.py new file mode 100644 index 000000000..95b26b406 --- /dev/null +++ b/tests/ext/databricks/_fakes.py @@ -0,0 +1,178 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +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. + +Shared test doubles for dapr.ext.databricks tests. Not a test module itself +(no test_ prefix), so unittest/pytest discovery skips it. +""" + +from __future__ import annotations + +import threading +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple + +import grpc + + +class SimulatedRpcError(grpc.RpcError): + """A constructible ``grpc.RpcError`` for tests, mirroring the pattern used in + ``tests/ext/workflow/test_workflow_client.py``.""" + + def __init__(self, code: Any, details: str) -> None: + self._code = code + self._details = details + + def code(self) -> Any: + return self._code + + def details(self) -> str: + return self._details + + +class FakeRow: + """A minimal stand-in for ``pyspark.sql.Row``.""" + + def __init__(self, **fields: Any) -> None: + self._fields = dict(fields) + + def asDict(self, recursive: bool = False) -> Dict[str, Any]: + if not recursive: + return dict(self._fields) + return {key: self._recurse(value) for key, value in self._fields.items()} + + @staticmethod + def _recurse(value: Any) -> Any: + if isinstance(value, FakeRow): + return value.asDict(recursive=True) + if isinstance(value, (list, tuple)): + return [FakeRow._recurse(item) for item in value] + return value + + def __getitem__(self, key: str) -> Any: + return self._fields[key] + + def __repr__(self) -> str: + return f'FakeRow({self._fields!r})' + + +class FakeDataFrame: + """A minimal stand-in for ``pyspark.sql.DataFrame``, backed by a plain list of rows. + + ``to_local_iterator_error``, when set, is raised by ``toLocalIterator`` + instead of iterating — used to simulate the real "toLocalIterator() is + not supported when using file-based collect" failure observed on + Databricks serverless / Spark Connect-backed compute, so the + ``collect()``-based fallback path (see ``DaprWorkflowBatchHandler._iter_rows``) + has real test coverage instead of only a local/classic-compute code path. + """ + + def __init__( + self, + rows: Iterable[FakeRow], + *, + to_local_iterator_error: Optional[BaseException] = None, + ) -> None: + self._rows = list(rows) + self._to_local_iterator_error = to_local_iterator_error + + def toLocalIterator(self, prefetchPartitions: bool = False): + if self._to_local_iterator_error is not None: + raise self._to_local_iterator_error + return iter(self._rows) + + def limit(self, num: int) -> 'FakeDataFrame': + return FakeDataFrame(self._rows[:num]) + + def collect(self) -> List[FakeRow]: + return list(self._rows) + + +class FakeWorkflowClient: + """An in-memory stand-in for ``DaprWorkflowClient`` that emulates the one + invariant this extension depends on: at most one non-purged instance per + ID, with a second ``schedule_new_workflow`` for the same ID rejected as a + duplicate — the same behavior the real Dapr Workflow instance-start API + documents (HTTP 409 / "already exists and is not yet reusable"). + + Individual instance IDs can be configured to fail in specific ways via + ``raise_on_get_state`` / ``raise_on_schedule``, to simulate transport, + auth, and validation failures independently of the exists/doesn't-exist + bookkeeping. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self.existing: Set[str] = set() + self.scheduled: List[Tuple[str, str, Any]] = [] + self.get_state_calls: List[str] = [] + self.schedule_calls: List[str] = [] + self.raise_on_get_state: Dict[str, BaseException] = {} + self.raise_on_schedule: Dict[str, BaseException] = {} + self.lost_response_for: Set[str] = set() + self.closed = False + + def get_workflow_state(self, instance_id: str, *, fetch_payloads: bool = True) -> Any: + self.get_state_calls.append(instance_id) + if instance_id in self.raise_on_get_state: + raise self.raise_on_get_state.pop(instance_id) + return object() if instance_id in self.existing else None + + def schedule_new_workflow( + self, + workflow: str, + *, + input: Optional[Any] = None, + instance_id: Optional[str] = None, + start_at: Optional[Any] = None, + reuse_id_policy: Optional[Any] = None, + ) -> str: + assert instance_id is not None + self.schedule_calls.append(instance_id) + with self._lock: + if instance_id in self.existing: + raise SimulatedRpcError( + grpc.StatusCode.ALREADY_EXISTS, + f"an active workflow with ID '{instance_id}' already exists", + ) + if instance_id in self.raise_on_schedule: + error = self.raise_on_schedule.pop(instance_id) + if instance_id in self.lost_response_for: + # Dapr durably accepted the instance server-side, but the + # caller experiences this attempt as a failure (e.g. the + # response never arrived) — the "lost response" scenario. + self.existing.add(instance_id) + raise error + self.existing.add(instance_id) + self.scheduled.append((workflow, instance_id, input)) + return instance_id + + def close(self) -> None: + self.closed = True + + +class BarrierSyncedWorkflowClient(FakeWorkflowClient): + """A ``FakeWorkflowClient`` that pauses every caller at the start of + ``schedule_new_workflow`` until ``party_count`` callers have all arrived. + + Used to deterministically reproduce the "two callers race to schedule the + same deterministic instance ID" scenario instead of relying on incidental + thread-scheduling timing. + """ + + def __init__(self, party_count: int) -> None: + super().__init__() + self._barrier = threading.Barrier(party_count) + + def schedule_new_workflow(self, workflow: str, **kwargs: Any) -> str: + self._barrier.wait(timeout=5.0) + return super().schedule_new_workflow(workflow, **kwargs) diff --git a/tests/ext/databricks/test_batch_handler.py b/tests/ext/databricks/test_batch_handler.py new file mode 100644 index 000000000..f884b5c2f --- /dev/null +++ b/tests/ext/databricks/test_batch_handler.py @@ -0,0 +1,328 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import threading +import time +import unittest + +import grpc + +from dapr.ext.databricks.batch_handler import DaprWorkflowBatchHandler +from dapr.ext.databricks.config import WorkflowSinkConfig +from dapr.ext.databricks.exceptions import DaprDatabricksSinkError +from tests.ext.databricks._fakes import ( + FakeDataFrame, + FakeRow, + FakeWorkflowClient, + SimulatedRpcError, +) + + +def _order_row(order_id, customer_id=1, status='READY'): + return FakeRow(order_id=order_id, customer_id=customer_id, status=status) + + +class BasicSchedulingTests(unittest.TestCase): + def test_one_row_schedules_one_workflow(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', workflow='process_order', id_field='order_id', namespace='orders' + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) + + self.assertEqual(len(client.scheduled), 1) + workflow, instance_id, payload = client.scheduled[0] + self.assertEqual(workflow, 'process_order') + self.assertEqual(instance_id, 'orders-order_actions-v1-123') + self.assertEqual(payload['data'], {'order_id': 123, 'customer_id': 1, 'status': 'READY'}) + + def test_reprocessing_same_row_does_not_duplicate(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', workflow='process_order', id_field='order_id', namespace='orders' + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) + + self.assertEqual(len(client.scheduled), 1) # only ever scheduled once + self.assertEqual(client.schedule_calls.count('orders-order_actions-v1-123'), 1) + + +class MicroBatchRetryTests(unittest.TestCase): + """The most important scenario: batch 42 has records A and B, B fails on + the first attempt, and the batch retry must produce exactly one workflow + for A and one for B — no duplicate for A, no missing execution for B.""" + + def test_partial_failure_then_retry_yields_exactly_one_workflow_each(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', + workflow='process_order', + id_field='order_id', + namespace='orders', + max_in_flight=1, # deterministic ordering for this test + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + batch = FakeDataFrame([_order_row('A'), _order_row('B')]) + + a_instance_id = 'orders-order_actions-v1-A' + b_instance_id = 'orders-order_actions-v1-B' + client.raise_on_schedule[b_instance_id] = SimulatedRpcError( + grpc.StatusCode.UNAVAILABLE, 'dapr sidecar unavailable' + ) + + with self.assertRaises(DaprDatabricksSinkError): + handler.process(batch, batch_id=42) + + self.assertIn(a_instance_id, client.existing) + self.assertNotIn(b_instance_id, client.existing) + self.assertEqual(len(client.scheduled), 1) + + # Lakeflow retries the same micro-batch (same rows, same batch_id). + handler.process(FakeDataFrame([_order_row('A'), _order_row('B')]), batch_id=42) + + # A is found already-existing on retry, so it is never re-attempted. + # B genuinely failed the first attempt, so it is legitimately + # attempted twice — the invariant is that it *succeeds* exactly once. + self.assertEqual(client.schedule_calls.count(a_instance_id), 1) + self.assertEqual(client.schedule_calls.count(b_instance_id), 2) + scheduled_ids = [instance_id for _, instance_id, _ in client.scheduled] + self.assertEqual(sorted(scheduled_ids), sorted([a_instance_id, b_instance_id])) + self.assertEqual(len(client.scheduled), 2) # exactly one workflow for A, one for B + + def test_lost_response_is_recovered_on_retry_without_duplicate(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', workflow='process_order', id_field='order_id', namespace='orders' + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + instance_id = 'orders-order_actions-v1-123' + + # Dapr accepts the schedule call durably, but the client never + # observes success (e.g. the response is lost to a network blip). + client.raise_on_schedule[instance_id] = SimulatedRpcError( + grpc.StatusCode.DEADLINE_EXCEEDED, 'deadline exceeded' + ) + client.lost_response_for.add(instance_id) + + with self.assertRaises(DaprDatabricksSinkError): + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) + + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) # Lakeflow retries + + self.assertEqual(client.schedule_calls, [instance_id]) # exactly one schedule attempt ever + self.assertEqual(len(client.scheduled), 0) # that one attempt never observed success + self.assertEqual(client.get_state_calls.count(instance_id), 2) # found on the 2nd check + + +class MaxRecordsPerBatchTests(unittest.TestCase): + def test_exceeding_limit_fails_the_batch_without_silent_truncation(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', + workflow='process_order', + id_field='order_id', + max_in_flight=1, + max_records_per_batch=2, + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + batch = FakeDataFrame([_order_row(1), _order_row(2), _order_row(3)]) + + with self.assertRaises(DaprDatabricksSinkError) as ctx: + handler.process(batch, batch_id=1) + + self.assertIn('max_records_per_batch', str(ctx.exception)) + # Already-submitted records before the cap was hit are still durably + # scheduled; we fail loudly rather than silently dropping record 3. + self.assertEqual(len(client.scheduled), 2) + + +class MetadataWrappingTests(unittest.TestCase): + def test_metadata_enabled_wraps_data_and_metadata(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', + workflow='process_order', + id_field='order_id', + namespace='orders', + generation='v2', + metadata=True, + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + + handler.process(FakeDataFrame([_order_row(123)]), batch_id=7) + + _, _, payload = client.scheduled[0] + self.assertEqual(set(payload.keys()), {'data', 'metadata'}) + self.assertEqual( + payload['metadata'], + { + 'sink': 'order_actions', + 'workflow': 'process_order', + 'batch_id': 7, + 'namespace': 'orders', + 'generation': 'v2', + }, + ) + + def test_metadata_disabled_sends_raw_mapped_row(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', workflow='process_order', id_field='order_id', metadata=False + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + + handler.process(FakeDataFrame([_order_row(123)]), batch_id=1) + + _, _, payload = client.scheduled[0] + self.assertEqual(payload, {'order_id': 123, 'customer_id': 1, 'status': 'READY'}) + + +class InputMapperTests(unittest.TestCase): + def test_custom_input_mapper_is_used(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='customer_actions', + workflow='process_customer', + id_field='customer_id', + input_mapper=lambda row: {'customer': row['customer_id'], 'status': row['status']}, + metadata=False, + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + row = FakeRow(customer_id=456, status='ACTIVE', internal_field='secret') + + handler.process(FakeDataFrame([row]), batch_id=1) + + _, _, payload = client.scheduled[0] + self.assertEqual(payload, {'customer': 456, 'status': 'ACTIVE'}) + + +class MalformedInputTests(unittest.TestCase): + def test_missing_business_key_field_fails_the_batch(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + row = FakeRow(customer_id=1) # no order_id + + with self.assertRaises(DaprDatabricksSinkError): + handler.process(FakeDataFrame([row]), batch_id=1) + + self.assertEqual(client.scheduled, []) + + +class BoundedConcurrencyTests(unittest.TestCase): + def test_never_exceeds_max_in_flight_concurrent_schedule_calls(self): + max_in_flight = 3 + lock = threading.Lock() + state = {'current': 0, 'peak': 0} + + class SlowWorkflowClient(FakeWorkflowClient): + def schedule_new_workflow(self, workflow, **kwargs): + with lock: + state['current'] += 1 + state['peak'] = max(state['peak'], state['current']) + time.sleep(0.05) + try: + return super().schedule_new_workflow(workflow, **kwargs) + finally: + with lock: + state['current'] -= 1 + + client = SlowWorkflowClient() + config = WorkflowSinkConfig( + name='order_actions', + workflow='process_order', + id_field='order_id', + max_in_flight=max_in_flight, + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + rows = [_order_row(i) for i in range(12)] + + handler.process(FakeDataFrame(rows), batch_id=1) + + self.assertEqual(len(client.scheduled), 12) + self.assertLessEqual(state['peak'], max_in_flight) + self.assertGreater( + state['peak'], 1 + ) # actually exercised concurrency, not accidentally serial + + +# The exact text a real Databricks serverless / Spark Connect-backed compute +# raised from `df.toLocalIterator()` itself during live end-to-end testing — +# see dapr/ext/databricks/AGENTS.md for the full story. +_TO_LOCAL_ITERATOR_UNSUPPORTED = Exception( + 'toLocalIterator() is not supported when using file-based collect' +) + + +class ToLocalIteratorFallbackTests(unittest.TestCase): + """Covers DaprWorkflowBatchHandler._iter_rows falling back to collect() + on compute where toLocalIterator() itself is unavailable.""" + + def test_falls_back_to_collect_and_still_schedules_every_row(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + batch = FakeDataFrame( + [_order_row(1), _order_row(2)], + to_local_iterator_error=_TO_LOCAL_ITERATOR_UNSUPPORTED, + ) + + handler.process(batch, batch_id=1) + + self.assertEqual(len(client.scheduled), 2) + + def test_fallback_still_enforces_max_records_per_batch(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig( + name='orders', + workflow='process_order', + id_field='order_id', + max_in_flight=1, + max_records_per_batch=2, + ) + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + batch = FakeDataFrame( + [_order_row(1), _order_row(2), _order_row(3)], + to_local_iterator_error=_TO_LOCAL_ITERATOR_UNSUPPORTED, + ) + + with self.assertRaises(DaprDatabricksSinkError) as ctx: + handler.process(batch, batch_id=1) + + self.assertIn('max_records_per_batch', str(ctx.exception)) + self.assertEqual(len(client.scheduled), 2) # bounded, not all 3 + + def test_unrelated_to_local_iterator_failure_is_not_swallowed(self): + client = FakeWorkflowClient() + config = WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') + handler = DaprWorkflowBatchHandler(config, workflow_client=client) + batch = FakeDataFrame( + [_order_row(1)], + to_local_iterator_error=RuntimeError('source table permission denied'), + ) + + with self.assertRaises(RuntimeError): + handler.process(batch, batch_id=1) + + self.assertEqual(client.scheduled, []) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/ext/databricks/test_config.py b/tests/ext/databricks/test_config.py new file mode 100644 index 000000000..885c10547 --- /dev/null +++ b/tests/ext/databricks/test_config.py @@ -0,0 +1,98 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +from dapr.ext.databricks.config import WorkflowSinkConfig +from dapr.ext.databricks.exceptions import SinkConfigurationError +from dapr.ext.databricks.mapping import default_row_mapper + + +class WorkflowSinkConfigTests(unittest.TestCase): + def test_minimal_valid_config(self): + config = WorkflowSinkConfig(name='orders', workflow='process_order', id_field='order_id') + self.assertEqual(config.namespace, 'default') + self.assertEqual(config.generation, 'v1') + self.assertTrue(config.metadata) + self.assertEqual(config.max_in_flight, 8) + self.assertIsNone(config.max_records_per_batch) + + def test_empty_name_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='', workflow='process_order') + + def test_empty_workflow_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='') + + def test_empty_namespace_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='process_order', namespace='') + + def test_empty_generation_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='process_order', generation='') + + def test_id_field_and_id_fields_together_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig( + name='orders', + workflow='process_order', + id_field='order_id', + id_fields=['account_id', 'transaction_id'], + ) + + def test_id_field_and_instance_id_factory_together_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig( + name='orders', + workflow='process_order', + id_field='order_id', + instance_id_factory=lambda row, batch_id: 'x', + ) + + def test_empty_id_fields_rejected(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='process_order', id_fields=[]) + + def test_max_in_flight_must_be_positive(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='process_order', max_in_flight=0) + + def test_max_records_per_batch_must_be_positive_when_set(self): + with self.assertRaises(SinkConfigurationError): + WorkflowSinkConfig(name='orders', workflow='process_order', max_records_per_batch=0) + + def test_row_mapper_defaults_to_default_row_mapper(self): + config = WorkflowSinkConfig(name='orders', workflow='process_order') + self.assertIs(config.row_mapper, default_row_mapper) + + def test_row_mapper_uses_configured_input_mapper(self): + def custom_mapper(row): + return {'x': 1} + + config = WorkflowSinkConfig( + name='orders', workflow='process_order', input_mapper=custom_mapper + ) + self.assertIs(config.row_mapper, custom_mapper) + + def test_config_is_immutable(self): + config = WorkflowSinkConfig(name='orders', workflow='process_order') + with self.assertRaises(Exception): + config.name = 'other' # type: ignore[misc] + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/ext/databricks/test_identity.py b/tests/ext/databricks/test_identity.py new file mode 100644 index 000000000..53aa986a8 --- /dev/null +++ b/tests/ext/databricks/test_identity.py @@ -0,0 +1,267 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest + +from dapr.ext.databricks.exceptions import MissingBusinessKeyError +from dapr.ext.databricks.identity import ( + _MAX_INSTANCE_ID_LENGTH, + derive_instance_id, + extract_business_key, + sanitize_segment, +) +from tests.ext.databricks._fakes import FakeRow + +_ALLOWED_CHARS = set('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-') + + +class SanitizeSegmentTests(unittest.TestCase): + def test_clean_short_segment_is_unchanged(self): + self.assertEqual(sanitize_segment('order-123_abc'), 'order-123_abc') + + def test_invalid_characters_are_hashed_deterministically(self): + first = sanitize_segment('order/123:456') + second = sanitize_segment('order/123:456') + + self.assertEqual(first, second) + self.assertTrue(set(first) <= _ALLOWED_CHARS) + self.assertNotEqual(first, 'order/123:456') + + def test_different_dirty_inputs_hash_differently(self): + self.assertNotEqual(sanitize_segment('a/b'), sanitize_segment('a/c')) + + def test_unicode_input_is_hashed_and_stable(self): + first = sanitize_segment('café-北京-🎉') + second = sanitize_segment('café-北京-🎉') + + self.assertEqual(first, second) + self.assertTrue(set(first) <= _ALLOWED_CHARS) + + def test_distinct_unicode_inputs_do_not_collide(self): + # A naive ASCII-fold/transliterate approach could collapse distinct + # non-Latin business keys onto the same (near-empty) output; hashing + # the whole segment instead must keep them distinct. + self.assertNotEqual(sanitize_segment('北京'), sanitize_segment('上海')) + + def test_empty_string_is_hashed_not_dropped(self): + result = sanitize_segment('') + self.assertTrue(result) + self.assertTrue(set(result) <= _ALLOWED_CHARS) + + def test_overlong_clean_segment_is_hashed(self): + long_clean = 'a' * 500 + result = sanitize_segment(long_clean) + self.assertNotEqual(result, long_clean) + self.assertLessEqual(len(result), 64) + + +class ExtractBusinessKeyTests(unittest.TestCase): + def test_single_id_field(self): + row = FakeRow(order_id=123, status='READY') + key = extract_business_key( + row, batch_id=1, id_field='order_id', id_fields=None, instance_id_factory=None + ) + self.assertEqual(key, '123') + + def test_composite_id_fields_join_in_order(self): + row = FakeRow(account_id='A1', transaction_id='T9') + key = extract_business_key( + row, + batch_id=1, + id_field=None, + id_fields=['account_id', 'transaction_id'], + instance_id_factory=None, + ) + self.assertEqual(key, 'A1_T9') + + def test_instance_id_factory_takes_priority_and_receives_batch_id(self): + seen_args = [] + + def factory(row, batch_id): + seen_args.append((row['order_id'], batch_id)) + return f'custom-{row["order_id"]}' + + row = FakeRow(order_id=42) + key = extract_business_key( + row, + batch_id=7, + id_field='order_id', # would be ignored: factory takes priority + id_fields=None, + instance_id_factory=factory, + ) + self.assertEqual(key, 'custom-42') + self.assertEqual(seen_args, [(42, 7)]) + + def test_no_strategy_configured_returns_none(self): + row = FakeRow(order_id=42) + key = extract_business_key( + row, batch_id=1, id_field=None, id_fields=None, instance_id_factory=None + ) + self.assertIsNone(key) + + def test_missing_id_field_raises(self): + row = FakeRow(customer_id=1) + with self.assertRaises(MissingBusinessKeyError): + extract_business_key( + row, batch_id=1, id_field='order_id', id_fields=None, instance_id_factory=None + ) + + def test_null_id_field_raises(self): + row = FakeRow(order_id=None) + with self.assertRaises(MissingBusinessKeyError): + extract_business_key( + row, batch_id=1, id_field='order_id', id_fields=None, instance_id_factory=None + ) + + def test_missing_field_within_composite_key_raises(self): + row = FakeRow(account_id='A1') + with self.assertRaises(MissingBusinessKeyError): + extract_business_key( + row, + batch_id=1, + id_field=None, + id_fields=['account_id', 'transaction_id'], + instance_id_factory=None, + ) + + +class DeriveInstanceIdTests(unittest.TestCase): + def test_deterministic_for_same_business_key(self): + first = derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key='123', + batch_id=42, + record_index=0, + ) + second = derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key='123', + batch_id=999, # different batch/record: must not matter when a business key exists + record_index=5, + ) + self.assertEqual(first, second) + self.assertEqual(first, 'orders-order_actions-v1-123') + + def test_different_business_keys_produce_different_ids(self): + def make_id(key): + return derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key=key, + batch_id=1, + record_index=0, + ) + + self.assertNotEqual(make_id('123'), make_id('124')) + + def test_generation_changes_identity(self): + def make_id(generation): + return derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation=generation, + business_key='123', + batch_id=1, + record_index=0, + ) + + self.assertNotEqual(make_id('v1'), make_id('v2')) + + def test_namespace_changes_identity(self): + def make_id(namespace): + return derive_instance_id( + namespace=namespace, + sink_name='order_actions', + generation='v1', + business_key='123', + batch_id=1, + record_index=0, + ) + + self.assertNotEqual(make_id('orders'), make_id('fraud')) + + def test_fallback_identity_uses_batch_and_record_index_when_no_business_key(self): + instance_id = derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key=None, + batch_id=42, + record_index=3, + ) + self.assertEqual(instance_id, 'orders-order_actions-v1-42-3') + + def test_fallback_identity_differs_by_record_index(self): + def make_id(idx): + return derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key=None, + batch_id=42, + record_index=idx, + ) + + self.assertNotEqual(make_id(0), make_id(1)) + + def test_result_always_within_length_bound(self): + instance_id = derive_instance_id( + namespace='n' * 100, + sink_name='s' * 100, + generation='g' * 100, + business_key='k' * 500, + batch_id=1, + record_index=0, + ) + self.assertLessEqual(len(instance_id), _MAX_INSTANCE_ID_LENGTH) + + def test_pathologically_long_clean_prefix_still_distinguishes_records(self): + # namespace/sink/generation alone can exceed the instance-id budget + # even though each segment is individually "clean". The final digest + # must never be truncated away, or every record in the sink would + # collide onto the same instance ID. + def make_id(key): + return derive_instance_id( + namespace='n' * 80, + sink_name='s' * 80, + generation='g' * 80, + business_key=key, + batch_id=1, + record_index=0, + ) + + first, second = make_id('key-one'), make_id('key-two') + self.assertLessEqual(len(first), _MAX_INSTANCE_ID_LENGTH) + self.assertNotEqual(first, second) + + def test_only_valid_instance_id_characters_are_produced(self): + instance_id = derive_instance_id( + namespace='orders', + sink_name='order_actions', + generation='v1', + business_key='id with spaces/slashes:colons', + batch_id=1, + record_index=0, + ) + self.assertTrue(set(instance_id) <= _ALLOWED_CHARS) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/ext/databricks/test_mapping.py b/tests/ext/databricks/test_mapping.py new file mode 100644 index 000000000..0bc1d4b72 --- /dev/null +++ b/tests/ext/databricks/test_mapping.py @@ -0,0 +1,82 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import json +import unittest +from datetime import date, datetime, timezone +from decimal import Decimal + +from dapr.ext.databricks.mapping import default_row_mapper +from tests.ext.databricks._fakes import FakeRow + + +class DefaultRowMapperTests(unittest.TestCase): + def test_simple_fields_pass_through(self): + row = FakeRow(order_id=123, customer_id=456, status='READY') + self.assertEqual( + default_row_mapper(row), + {'order_id': 123, 'customer_id': 456, 'status': 'READY'}, + ) + + def test_result_is_json_serializable(self): + row = FakeRow(order_id=123, status='READY') + json.dumps(default_row_mapper(row)) # must not raise + + def test_datetime_becomes_isoformat_string(self): + when = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc) + row = FakeRow(created_at=when) + self.assertEqual(default_row_mapper(row)['created_at'], when.isoformat()) + + def test_date_becomes_isoformat_string(self): + day = date(2026, 1, 2) + row = FakeRow(order_date=day) + self.assertEqual(default_row_mapper(row)['order_date'], day.isoformat()) + + def test_decimal_becomes_string_to_avoid_precision_loss(self): + row = FakeRow(amount=Decimal('19.999999999999999999')) + mapped = default_row_mapper(row) + self.assertEqual(mapped['amount'], '19.999999999999999999') + self.assertIsInstance(mapped['amount'], str) + + def test_bytes_become_base64_string(self): + row = FakeRow(payload=b'\x00\x01\xff') + mapped = default_row_mapper(row) + self.assertEqual(mapped['payload'], 'AAH/') + + def test_nested_row_is_recursively_converted(self): + row = FakeRow(order_id=1, customer=FakeRow(id=9, name='Ada')) + mapped = default_row_mapper(row) + self.assertEqual(mapped['customer'], {'id': 9, 'name': 'Ada'}) + + def test_list_of_rows_is_recursively_converted(self): + row = FakeRow(items=[FakeRow(sku='A'), FakeRow(sku='B')]) + mapped = default_row_mapper(row) + self.assertEqual(mapped['items'], [{'sku': 'A'}, {'sku': 'B'}]) + + +class CustomInputMapperTests(unittest.TestCase): + def test_custom_mapper_overrides_default_shape(self): + row = FakeRow(customer_id=456, status='READY', internal_score=0.87) + + def custom_mapper(r): + return {'customer': r['customer_id'], 'status': r['status']} + + mapped = custom_mapper(row) + self.assertEqual(mapped, {'customer': 456, 'status': 'READY'}) + self.assertNotIn('internal_score', mapped) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/ext/databricks/test_scheduling.py b/tests/ext/databricks/test_scheduling.py new file mode 100644 index 000000000..c5f519477 --- /dev/null +++ b/tests/ext/databricks/test_scheduling.py @@ -0,0 +1,133 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import unittest +from concurrent.futures import ThreadPoolExecutor + +import grpc + +from dapr.ext.databricks.scheduling import ensure_workflow_scheduled, is_duplicate_instance_error +from tests.ext.databricks._fakes import ( + BarrierSyncedWorkflowClient, + FakeWorkflowClient, + SimulatedRpcError, +) + + +class IsDuplicateInstanceErrorTests(unittest.TestCase): + def test_true_for_already_exists_status_code(self): + error = SimulatedRpcError(grpc.StatusCode.ALREADY_EXISTS, 'irrelevant message') + self.assertTrue(is_duplicate_instance_error(error)) + + def test_true_for_message_text_regardless_of_code(self): + error = SimulatedRpcError( + grpc.StatusCode.FAILED_PRECONDITION, + 'a workflow with the given instance ID already exists and is not yet reusable', + ) + self.assertTrue(is_duplicate_instance_error(error)) + + def test_false_for_unrelated_error(self): + error = SimulatedRpcError(grpc.StatusCode.UNAVAILABLE, 'connection refused') + self.assertFalse(is_duplicate_instance_error(error)) + + +class EnsureWorkflowScheduledTests(unittest.TestCase): + def test_new_instance_is_scheduled(self): + client = FakeWorkflowClient() + + outcome = ensure_workflow_scheduled(client, 'process_order', 'orders-1', {'order_id': 1}) + + self.assertTrue(outcome.newly_scheduled) + self.assertEqual(outcome.instance_id, 'orders-1') + self.assertEqual(client.scheduled, [('process_order', 'orders-1', {'order_id': 1})]) + self.assertEqual(client.get_state_calls, ['orders-1']) + + def test_existing_instance_is_not_rescheduled(self): + client = FakeWorkflowClient() + client.existing.add('orders-1') + + outcome = ensure_workflow_scheduled(client, 'process_order', 'orders-1', {'order_id': 1}) + + self.assertFalse(outcome.newly_scheduled) + self.assertEqual(client.scheduled, []) + self.assertEqual(client.schedule_calls, []) + + def test_lost_response_then_retry_finds_existing_instance(self): + """schedule succeeds durably server-side, but this attempt sees an error; + a later retry's existence check must find it and must not reschedule.""" + client = FakeWorkflowClient() + client.raise_on_schedule['orders-1'] = SimulatedRpcError( + grpc.StatusCode.DEADLINE_EXCEEDED, 'deadline exceeded' + ) + client.lost_response_for.add('orders-1') + + with self.assertRaises(grpc.RpcError): + ensure_workflow_scheduled(client, 'process_order', 'orders-1', {'order_id': 1}) + + # Dapr durably accepted it despite the client-visible failure above. + self.assertIn('orders-1', client.existing) + self.assertEqual(client.scheduled, []) # this attempt never observed success + + retry_outcome = ensure_workflow_scheduled( + client, 'process_order', 'orders-1', {'order_id': 1} + ) + + self.assertFalse(retry_outcome.newly_scheduled) + self.assertEqual(client.schedule_calls, ['orders-1']) # no second schedule call + + def test_non_duplicate_schedule_error_propagates(self): + client = FakeWorkflowClient() + client.raise_on_schedule['orders-1'] = SimulatedRpcError( + grpc.StatusCode.UNAUTHENTICATED, 'invalid api token' + ) + + with self.assertRaises(grpc.RpcError): + ensure_workflow_scheduled(client, 'process_order', 'orders-1', {'order_id': 1}) + + def test_get_workflow_state_error_propagates_without_scheduling(self): + client = FakeWorkflowClient() + client.raise_on_get_state['orders-1'] = SimulatedRpcError( + grpc.StatusCode.UNAVAILABLE, 'dapr sidecar unavailable' + ) + + with self.assertRaises(grpc.RpcError): + ensure_workflow_scheduled(client, 'process_order', 'orders-1', {'order_id': 1}) + + self.assertEqual(client.scheduled, []) + self.assertEqual(client.schedule_calls, []) + + def test_concurrent_race_resolves_to_exactly_one_new_and_rest_existing(self): + party_count = 2 + client = BarrierSyncedWorkflowClient(party_count=party_count) + + with ThreadPoolExecutor(max_workers=party_count) as pool: + futures = [ + pool.submit( + ensure_workflow_scheduled, client, 'process_order', 'orders-1', {'order_id': 1} + ) + for _ in range(party_count) + ] + outcomes = [future.result() for future in futures] + + newly_scheduled_count = sum(1 for outcome in outcomes if outcome.newly_scheduled) + self.assertEqual(newly_scheduled_count, 1) + self.assertEqual(len(outcomes) - newly_scheduled_count, party_count - 1) + for outcome in outcomes: + self.assertEqual(outcome.instance_id, 'orders-1') + self.assertEqual(client.scheduled, [('process_order', 'orders-1', {'order_id': 1})]) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/ext/databricks/test_sink.py b/tests/ext/databricks/test_sink.py new file mode 100644 index 000000000..2967df075 --- /dev/null +++ b/tests/ext/databricks/test_sink.py @@ -0,0 +1,128 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +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. + +Tests for register_workflow_sink's pyspark.pipelines wiring. pyspark is +never installed for these tests (see AGENTS.md): a fake module is injected +into sys.modules for the "inside Lakeflow" tests, and removed/absent for the +"outside Lakeflow" test, matching how the extension is actually used. +""" + +import sys +import types +import unittest +from unittest import mock + +from dapr.ext.databricks.batch_handler import DaprWorkflowBatchHandler +from tests.ext.databricks._fakes import FakeDataFrame, FakeRow, FakeWorkflowClient + + +class FakePipelinesModule(types.ModuleType): + """Fakes just enough of ``pyspark.pipelines`` for ``register_workflow_sink``.""" + + def __init__(self): + super().__init__('pyspark.pipelines') + self.registered_sinks = {} + + def foreach_batch_sink(self, name=None, **_kwargs): + def decorator(fn): + self.registered_sinks[name or fn.__name__] = fn + return fn + + return decorator + + +def _install_fake_pyspark(): + fake_pipelines = FakePipelinesModule() + fake_pyspark = types.ModuleType('pyspark') + fake_pyspark.pipelines = fake_pipelines + return fake_pyspark, fake_pipelines + + +class RegisterWorkflowSinkOutsideLakeflowTests(unittest.TestCase): + def test_clear_error_when_pyspark_is_unavailable(self): + # No pyspark.* entries are patched into sys.modules here, and the + # module is not actually installed in this environment (by design: + # see the `databricks` extra in pyproject.toml). + with mock.patch.dict(sys.modules): + sys.modules.pop('pyspark', None) + sys.modules.pop('pyspark.pipelines', None) + + from dapr.ext.databricks.sink import register_workflow_sink + + with self.assertRaises(ImportError) as ctx: + register_workflow_sink(name='orders', workflow='process_order') + + self.assertIn('pyspark.pipelines.foreach_batch_sink', str(ctx.exception)) + self.assertIn('Databricks Lakeflow', str(ctx.exception)) + + +class RegisterWorkflowSinkInsideLakeflowTests(unittest.TestCase): + def setUp(self): + self.fake_pyspark, self.fake_pipelines = _install_fake_pyspark() + patcher = mock.patch.dict( + sys.modules, + {'pyspark': self.fake_pyspark, 'pyspark.pipelines': self.fake_pipelines}, + ) + patcher.start() + self.addCleanup(patcher.stop) + + client_patcher = mock.patch( + 'dapr.ext.databricks.batch_handler.DaprWorkflowClient', autospec=False + ) + self.mock_client_cls = client_patcher.start() + self.addCleanup(client_patcher.stop) + self.fake_client = FakeWorkflowClient() + self.mock_client_cls.return_value = self.fake_client + + def test_registers_a_foreach_batch_sink_under_the_given_name(self): + from dapr.ext.databricks.sink import register_workflow_sink + + handler = register_workflow_sink( + name='order_actions', workflow='process_order', id_field='order_id' + ) + + self.assertIn('order_actions', self.fake_pipelines.registered_sinks) + self.assertIsInstance(handler, DaprWorkflowBatchHandler) + + def test_registered_sink_delegates_to_the_handler(self): + from dapr.ext.databricks.sink import register_workflow_sink + + register_workflow_sink( + name='order_actions', workflow='process_order', id_field='order_id', namespace='orders' + ) + sink_fn = self.fake_pipelines.registered_sinks['order_actions'] + + sink_fn(FakeDataFrame([FakeRow(order_id=123)]), 1) + + self.assertEqual(len(self.fake_client.scheduled), 1) + _, instance_id, _ = self.fake_client.scheduled[0] + self.assertEqual(instance_id, 'orders-order_actions-v1-123') + + def test_invalid_configuration_raises_before_registering_pyspark_sink(self): + from dapr.ext.databricks.exceptions import SinkConfigurationError + from dapr.ext.databricks.sink import register_workflow_sink + + with self.assertRaises(SinkConfigurationError): + register_workflow_sink( + name='order_actions', + workflow='process_order', + id_field='order_id', + id_fields=['a', 'b'], + ) + + self.assertNotIn('order_actions', self.fake_pipelines.registered_sinks) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/integration/test_databricks_sink.py b/tests/integration/test_databricks_sink.py new file mode 100644 index 000000000..aa02913dd --- /dev/null +++ b/tests/integration/test_databricks_sink.py @@ -0,0 +1,191 @@ +# -*- coding: utf-8 -*- + +""" +Copyright 2026 The Dapr Authors +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. + +Runs dapr.ext.databricks against a real Dapr sidecar (via ``dapr_env`` from +conftest.py) instead of the mocked ``DaprWorkflowClient`` the unit tests +under tests/ext/databricks use. No Databricks workspace or pyspark install +is required: Lakeflow rows/DataFrames are faked locally, the same way the +unit test suite does it (see tests/ext/databricks/_fakes.py for the +equivalent used there). +""" + +import threading + +import grpc +import pytest + +import dapr.ext.workflow as wf +from dapr.ext.databricks import DaprWorkflowBatchHandler, WorkflowSinkConfig +from dapr.ext.databricks.exceptions import DaprDatabricksSinkError +from dapr.ext.databricks.scheduling import is_duplicate_instance_error + +HOST = '127.0.0.1' +GRPC_PORT = '13501' +NAMESPACE = 'itest-orders' +SINK_NAME = 'order_actions' +WORKFLOW_NAME = 'itest_record_order' + + +class _FakeLakeflowRow: + """Minimal ``pyspark.sql.Row`` stand-in; see ``dapr.ext.databricks._typing.RowLike`` + for the exact (small) surface this needs to implement.""" + + def __init__(self, **fields): + self._fields = fields + + def asDict(self, recursive=False): + return dict(self._fields) + + def __getitem__(self, key): + return self._fields[key] + + +class _FakeLakeflowBatch: + """Minimal ``pyspark.sql.DataFrame`` micro-batch stand-in.""" + + def __init__(self, rows): + self._rows = list(rows) + + def toLocalIterator(self, prefetchPartitions=False): + return iter(self._rows) + + +def _instance_id(order_id) -> str: + return f'{NAMESPACE}-{SINK_NAME}-v1-{order_id}' + + +def _purge(order_ids): + wf_client = wf.DaprWorkflowClient(host=HOST, port=GRPC_PORT) + for order_id in order_ids: + try: + wf_client.purge_workflow(_instance_id(order_id)) + except Exception: + pass + wf_client.close() + + +@pytest.fixture(scope='module') +def sidecar(dapr_env): + return dapr_env.start_sidecar(app_id='test-databricks-sink') + + +@pytest.fixture(scope='module') +def execution_log(): + return [] + + +@pytest.fixture(scope='module') +def runtime(sidecar, execution_log): + """A tiny real Dapr Workflow app: one workflow, one activity, registered + against the sidecar started above -- this is the "start/register a small + workflow" step the sink is tested against.""" + lock = threading.Lock() + rt = wf.WorkflowRuntime(host=HOST, port=GRPC_PORT) + + @rt.activity(name='itest_record_execution') + def record_execution(ctx, order_id): + with lock: + execution_log.append(order_id) + return order_id + + @rt.workflow(name=WORKFLOW_NAME) + def record_order(ctx, envelope): + order_id = envelope['data']['order_id'] + result = yield ctx.call_activity(record_execution, input=order_id) + return result + + rt.start() + rt.wait_for_worker_ready(timeout=30) + yield rt + rt.shutdown() + + +@pytest.fixture +def handler(sidecar): + config = WorkflowSinkConfig( + name=SINK_NAME, + workflow=WORKFLOW_NAME, + id_field='order_id', + namespace=NAMESPACE, + host=HOST, + port=GRPC_PORT, + ) + h = DaprWorkflowBatchHandler(config) + yield h + h.close() + + +def test_batch_schedules_one_workflow_per_row_and_retry_creates_no_duplicates( + runtime, handler, execution_log +): + order_ids = ['ITEST-1', 'ITEST-2'] + _purge(order_ids) + execution_log.clear() + + def _batch(): + return _FakeLakeflowBatch([_FakeLakeflowRow(order_id=oid) for oid in order_ids]) + + # 1. Pass fake Lakeflow rows through the sink handler -- schedules two new + # workflow instances against the real sidecar. + handler.process(_batch(), batch_id=1) + + # 2. Verify the workflows were created and actually ran to completion. + wf_client = wf.DaprWorkflowClient(host=HOST, port=GRPC_PORT) + for order_id in order_ids: + state = wf_client.wait_for_workflow_completion( + _instance_id(order_id), timeout_in_seconds=30 + ) + assert state is not None + assert state.runtime_status.name == 'COMPLETED' + assert sorted(execution_log) == sorted(order_ids) + + # 3. Retry the identical micro-batch -- what a Lakeflow retry after a + # worker restart looks like. + handler.process(_batch(), batch_id=1) + + # 4. Verify no duplicate workflows were created: the activity never ran a + # second time for either order, proving the retry was recognized as + # already-handled rather than executing a second instance. + assert sorted(execution_log) == sorted(order_ids) + + wf_client.close() + + +def test_duplicate_schedule_is_rejected_by_real_dapr_and_recognized(runtime): + """Validates scheduling.is_duplicate_instance_error against the actual + Dapr sidecar response (not the simulated fakes tests/ext/databricks uses).""" + wf_client = wf.DaprWorkflowClient(host=HOST, port=GRPC_PORT) + instance_id = f'{NAMESPACE}-duplicate-check' + try: + wf_client.purge_workflow(instance_id) + except Exception: + pass + + wf_client.schedule_new_workflow( + WORKFLOW_NAME, input={'data': {'order_id': 'dup-check'}}, instance_id=instance_id + ) + + with pytest.raises(grpc.RpcError) as exc_info: + wf_client.schedule_new_workflow( + WORKFLOW_NAME, input={'data': {'order_id': 'dup-check'}}, instance_id=instance_id + ) + + assert is_duplicate_instance_error(exc_info.value) + wf_client.close() + + +def test_missing_business_key_fails_the_batch(handler): + bad_row = _FakeLakeflowRow(customer_id=1) # no order_id + with pytest.raises(DaprDatabricksSinkError): + handler.process(_FakeLakeflowBatch([bad_row]), batch_id=2) diff --git a/uv.lock b/uv.lock index 17ae62033..9ddc215c6 100644 --- a/uv.lock +++ b/uv.lock @@ -800,7 +800,7 @@ requires-dist = [ { name = "uvicorn", marker = "extra == 'all'", specifier = ">=0.11.6,<1.0.0" }, { name = "uvicorn", marker = "extra == 'fastapi'", specifier = ">=0.11.6,<1.0.0" }, ] -provides-extras = ["all", "fastapi", "flask", "grpc", "langgraph", "strands", "workflow"] +provides-extras = ["all", "databricks", "fastapi", "flask", "grpc", "langgraph", "strands", "workflow"] [package.metadata.requires-dev] dev = [ @@ -880,7 +880,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [