"""

duroxide - Python SDK for the Duroxide durable execution runtime.

Generator-based orchestrations: users write generator functions that yield

ScheduledTask descriptors. The Rust runtime handles DurableFutures.

"""

from duroxide._duroxide import (

PySqliteProvider,

PyPostgresProvider,

PyPostgresEntraOptions,

PyClient,

PyRuntime,

RuntimeOptions,

OrchestrationStatus,

SystemMetrics,

SystemStats,

QueueDepths,

InstanceInfo,

ExecutionInfo,

InstanceTree,

DeleteInstanceResult,

PruneOptions,

PruneResult,

InstanceFilter,

Event,

MetricsSnapshot,

activity_trace_log,

orchestration_trace_log,

orchestration_set_custom_status,

orchestration_reset_custom_status,

orchestration_get_custom_status,

activity_is_cancelled,

activity_tag,

init_tracing,

)

from duroxide.context import OrchestrationContext, ActivityContext, ScheduledTask

from duroxide.driver import create_generator, next_step, dispose_generator

# Backwards-compatibility aliases (deprecated, use clean names)

PyOrchestrationStatus = OrchestrationStatus

PySystemMetrics = SystemMetrics

PySystemStats = SystemStats

PyQueueDepths = QueueDepths

PyInstanceInfo = InstanceInfo

PyExecutionInfo = ExecutionInfo

PyInstanceTree = InstanceTree

PyDeleteInstanceResult = DeleteInstanceResult

PyPruneOptions = PruneOptions

PyPruneResult = PruneResult

PyInstanceFilter = InstanceFilter

PyEvent = Event

PyMetricsSnapshot = MetricsSnapshot

PyRuntimeOptions = RuntimeOptions

import json

# ─── Generator Driver Registry ───────────────────────────────────

_orchestration_functions: dict = {}

# ─── Result Wrapper ──────────────────────────────────────────────

class OrchestrationResult:

"""Wrapper for orchestration status with parsed output."""

def __init__(self, status, output=None, error=None, custom_status=None, custom_status_version=0):

self.status = status

self.output = output

self.error = error

self.custom_status = custom_status

self.custom_status_version = custom_status_version

def _parse_status(raw):

"""Convert an OrchestrationStatus to an OrchestrationResult with parsed output."""

output = raw.output

if output is not None:

try:

output = json.loads(output)

except (json.JSONDecodeError, TypeError):

pass

return OrchestrationResult(

status=raw.status,

output=output,

error=raw.error,

custom_status=raw.custom_status,

custom_status_version=raw.custom_status_version,

)

def parse_result(result):

"""Parse a result from join/race calls.

Handles these formats:

- List of {"ok": val} dicts (from join/all) → [val, ...]

- Single {"ok": val} dict → val

- Single {"err": val} dict → raises Exception

Useful for manually processing results from non-typed join/race calls.

"""

if isinstance(result, list):

return [

item.get("ok") if isinstance(item, dict) and "ok" in item else item

for item in result

]

if isinstance(result, dict):

if "ok" in result:

return result["ok"]

if "err" in result:

raise Exception(result["err"])

return result

# ─── Public API ───────────────────────────────────────────────────

class SqliteProvider:

"""SQLite provider for duroxide."""

def __init__(self, native):

self._native = native

@staticmethod

def open(path: str) -> "SqliteProvider":

"""Open a SQLite database file."""

return SqliteProvider(PySqliteProvider.open(path))

@staticmethod

def in_memory() -> "SqliteProvider":

"""Create an in-memory SQLite database."""

return SqliteProvider(PySqliteProvider.in_memory())

class PostgresEntraOptions:

"""Options for Entra ID (Azure AD) authentication with PostgreSQL.

All parameters are keyword-only and optional; omitting a parameter uses

the duroxide-pg default for that setting.

Parameters

----------

audience:

Token audience/scope. Override for sovereign clouds (e.g., Azure US

Government: ``https://ossrdbms-aad.database.usgovcloudapi.net/.default``).

max_connections:

Maximum pool connection count.

acquire_timeout_ms:

Pool connection acquisition timeout in milliseconds.

refresh_interval_ms:

Upper bound on time between token refresh attempts, in milliseconds.

"""

def __init__(

self,

*,

audience: "str | None" = None,

max_connections: "int | None" = None,

acquire_timeout_ms: "int | None" = None,

refresh_interval_ms: "int | None" = None,

):

self._native = PyPostgresEntraOptions(

audience=audience,

max_connections=max_connections,

acquire_timeout_ms=acquire_timeout_ms,

refresh_interval_ms=refresh_interval_ms,

)

class PostgresProvider:

"""PostgreSQL provider for duroxide."""

def __init__(self, native):

self._native = native

self._type = "postgres"

@staticmethod

def connect(database_url: str) -> "PostgresProvider":

"""Connect to a PostgreSQL database (uses 'public' schema)."""

return PostgresProvider(PyPostgresProvider.connect(database_url))

@staticmethod

def connect_with_schema(database_url: str, schema: str) -> "PostgresProvider":

"""Connect to a PostgreSQL database with a custom schema."""

return PostgresProvider(

PyPostgresProvider.connect_with_schema(database_url, schema)

)

@staticmethod

def connect_with_entra(

host: str,

port: int,

database: str,

user: str,

options: "PostgresEntraOptions | None" = None,

) -> "PostgresProvider":

"""Connect to Azure Database for PostgreSQL using Entra ID (Azure AD) auth.

The SDK fetches and refreshes the access token automatically via the

DefaultAzureCredential chain (managed identity, environment variables,

Azure CLI, etc.).

Parameters

----------

host:

PostgreSQL server hostname (e.g. ``myserver.postgres.database.azure.com``).

port:

PostgreSQL server port (usually 5432).

database:

Target database name.

user:

Entra principal name mapped to a PostgreSQL role on the server.

options:

Optional :class:`PostgresEntraOptions` for tuning. Pass ``None``

(or omit) to use defaults.

"""

native_opts = options._native if options is not None else None

return PostgresProvider(

PyPostgresProvider.connect_with_entra(host, port, database, user, native_opts)

)

@staticmethod

def connect_with_schema_and_entra(

host: str,

port: int,

database: str,

user: str,

schema: str,

options: "PostgresEntraOptions | None" = None,

) -> "PostgresProvider":

"""Same as :meth:`connect_with_entra` but uses a custom schema.

The schema will be created if it does not exist. Useful for

multi-tenant deployments where each tenant has its own schema.

"""

native_opts = options._native if options is not None else None

return PostgresProvider(

PyPostgresProvider.connect_with_schema_and_entra(

host, port, database, user, schema, native_opts

)

)

class Client:

"""Client for starting and managing orchestration instances."""

def __init__(self, provider):

if getattr(provider, "_type", None) == "postgres":

self._native = PyClient.from_postgres(provider._native)

else:

self._native = PyClient.from_sqlite(provider._native)

def start_orchestration(self, instance_id: str, name: str, input=None):

self._native.start_orchestration(

instance_id, name, json.dumps(input)

)

def start_orchestration_versioned(

self, instance_id: str, name: str, input, version: str

):

self._native.start_orchestration_versioned(

instance_id, name, json.dumps(input), version

)

def get_status(self, instance_id: str):

result = self._native.get_status(instance_id)

return _parse_status(result)

def wait_for_orchestration(self, instance_id: str, timeout_ms: int = 30000):

result = self._native.wait_for_orchestration(instance_id, timeout_ms)

return _parse_status(result)

def wait_for_status_change(

self,

instance_id: str,

last_seen_version: int = 0,

poll_interval_ms: int = 200,

timeout_ms: int = 30000,

):

"""Wait for custom status changes on an orchestration instance.

Polls until the custom_status_version changes from last_seen_version,

or the orchestration reaches a terminal state.

Returns an OrchestrationResult with custom_status and custom_status_version.

"""

result = self._native.wait_for_status_change(

instance_id, last_seen_version, poll_interval_ms, timeout_ms

)

return _parse_status(result)

def get_kv_value(self, instance_id: str, key: str) -> "Optional[str]":

"""Read a single KV entry for the given instance."""

return self._native.get_kv_value(instance_id, key)

def get_kv_value_typed(self, instance_id: str, key: str):

"""Read a single KV entry for the given instance and JSON-decode it."""

raw = self.get_kv_value(instance_id, key)

if raw is None:

return None

return json.loads(raw)

def wait_for_kv_value(self, instance_id: str, key: str, timeout_ms: int = 30000) -> str:

"""Wait for a KV value to be set on an instance."""

try:

return self._native.wait_for_kv_value(instance_id, key, timeout_ms)

except RuntimeError as exc:

if "timed out" in str(exc).lower():

raise TimeoutError(

f"Timed out waiting for value '{key}' on instance '{instance_id}'"

) from exc

raise

def wait_for_kv_value_typed(self, instance_id: str, key: str, timeout_ms: int = 30000):

"""Wait for a KV value to be set on an instance and JSON-decode it."""

return json.loads(self.wait_for_kv_value(instance_id, key, timeout_ms))

def cancel_instance(self, instance_id: str, reason: str = None):

self._native.cancel_instance(instance_id, reason)

def raise_event(self, instance_id: str, event_name: str, data=None):

self._native.raise_event(

instance_id, event_name, json.dumps(data)

)

def enqueue_event(self, instance_id: str, queue_name: str, data=None):

"""Enqueue an event into a named queue for an orchestration instance.

Uses FIFO mailbox semantics. Matched by ctx.dequeue_event() in the orchestration.

"""

self._native.enqueue_event(

instance_id, queue_name, json.dumps(data)

)

# ─── Typed convenience wrappers ────────────────────────

def start_orchestration_typed(self, instance_id: str, name: str, input=None):

"""Start an orchestration with auto JSON input serialization."""

self.start_orchestration(instance_id, name, input)

def start_orchestration_versioned_typed(

self, instance_id: str, name: str, input, version: str

):

"""Start a versioned orchestration with auto JSON input serialization."""

self.start_orchestration_versioned(instance_id, name, input, version)

def raise_event_typed(self, instance_id: str, event_name: str, data=None):

"""Raise an event with auto JSON data serialization."""

self.raise_event(instance_id, event_name, data)

def enqueue_event_typed(self, instance_id: str, queue_name: str, data=None):

"""Enqueue an event with auto JSON data serialization."""

self.enqueue_event(instance_id, queue_name, data)

def wait_for_orchestration_typed(self, instance_id: str, timeout_ms: int = 30000):

"""Wait for an orchestration and return the parsed output directly.

If the orchestration failed, raises an Exception with the error message.

"""

result = self.wait_for_orchestration(instance_id, timeout_ms)

if result.status == "Failed":

raise Exception(result.error or "Orchestration failed")

return result.output

def get_system_metrics(self):

return self._native.get_system_metrics()

def get_orchestration_stats(self, instance_id: str):

return self._native.get_orchestration_stats(instance_id)

def get_queue_depths(self):

return self._native.get_queue_depths()

# ─── Management / Admin API ─────────────────────────────

def list_all_instances(self):

return self._native.list_all_instances()

def list_instances_by_status(self, status: str):

return self._native.list_instances_by_status(status)

def get_instance_info(self, instance_id: str):

return self._native.get_instance_info(instance_id)

def get_execution_info(self, instance_id: str, execution_id: int):

return self._native.get_execution_info(instance_id, execution_id)

def list_executions(self, instance_id: str):

return self._native.list_executions(instance_id)

def read_execution_history(self, instance_id: str, execution_id: int):

return self._native.read_execution_history(instance_id, execution_id)

def get_instance_tree(self, instance_id: str):

return self._native.get_instance_tree(instance_id)

def delete_instance(self, instance_id: str, force: bool = False):

return self._native.delete_instance(instance_id, force)

def delete_instance_bulk(self, filter=None):

if filter is None:

filter = InstanceFilter()

return self._native.delete_instance_bulk(filter)

def prune_executions(self, instance_id: str, options=None):

if options is None:

options = PruneOptions()

return self._native.prune_executions(instance_id, options)

def prune_executions_bulk(self, filter=None, options=None):

if filter is None:

filter = InstanceFilter()

if options is None:

options = PruneOptions()

return self._native.prune_executions_bulk(filter, options)

class Runtime:

"""Durable execution runtime."""

def __init__(self, provider, options=None):

if getattr(provider, "_type", None) == "postgres":

self._native = PyRuntime.from_postgres(provider._native, options)

else:

self._native = PyRuntime.from_sqlite(provider._native, options)

# Wire up the generator driver functions

self._native.set_generator_driver(

create_generator, next_step, dispose_generator

)

def register_activity(self, name: str, fn=None):

"""Register an activity function. Can be used as a decorator.

Usage:

@runtime.register_activity("my_activity")

def my_activity(ctx, input):

return {"result": "done"}

# or:

runtime.register_activity("my_activity", my_fn)

"""

if fn is not None:

self._register_activity_impl(name, fn)

return fn

# Decorator usage

def decorator(func):

self._register_activity_impl(name, func)

return func

return decorator

def _register_activity_impl(self, name: str, fn):

def wrapped_fn(payload: str) -> str:

newline_idx = payload.index("\n")

ctx_info_str = payload[:newline_idx]

input_str = payload[newline_idx + 1 :]

ctx_info = json.loads(ctx_info_str)

ctx = ActivityContext(ctx_info)

try:

input_val = json.loads(input_str)

except (json.JSONDecodeError, TypeError):

input_val = input_str

result = fn(ctx, input_val)

return json.dumps(result if result is not None else None)

self._native.register_activity(name, wrapped_fn)

def register_activity_typed(self, name: str, fn=None):

"""Register a typed activity. Mirrors Rust core's register_typed.

Input is auto-parsed from JSON; output is auto-serialized.

Usage:

@runtime.register_activity_typed("my_activity")

def my_activity(ctx, input): # input is already a dict/object

return {"result": "done"} # return value auto-serialized

"""

if fn is not None:

self._register_activity_impl(name, fn)

return fn

def decorator(func):

self._register_activity_impl(name, func)

return func

return decorator

def register_orchestration(self, name: str, fn=None):

"""Register an orchestration generator function. Can be used as a decorator.

Usage:

@runtime.register_orchestration("my_orch")

def my_orch(ctx, input):

result = yield ctx.schedule_activity("work", input)

return result

# or:

runtime.register_orchestration("my_orch", my_fn)

"""

if fn is not None:

_orchestration_functions[name] = fn

self._native.register_orchestration(name)

return fn

def decorator(func):

_orchestration_functions[name] = func

self._native.register_orchestration(name)

return func

return decorator

def register_orchestration_typed(self, name: str, fn=None):

"""Register a typed orchestration. Input is auto-parsed; output is auto-serialized.

Mirrors Rust core's register_typed on OrchestrationRegistryBuilder.

"""

return self.register_orchestration(name, fn)

def register_orchestration_versioned(self, name: str, version: str, fn=None):

"""Register a versioned orchestration generator function. Can be used as a decorator."""

if fn is not None:

key = f"{name}@{version}"

_orchestration_functions[key] = fn

self._native.register_orchestration_versioned(name, version)

return fn

def decorator(func):

key = f"{name}@{version}"

_orchestration_functions[key] = func

self._native.register_orchestration_versioned(name, version)

return func

return decorator

def register_orchestration_versioned_typed(self, name: str, version: str, fn=None):

"""Register a typed versioned orchestration. Input is auto-parsed; output is auto-serialized.

Can be used as a decorator:

@runtime.register_orchestration_versioned_typed("MyOrch", "1.0.0")

def my_orch(ctx, input):

...

"""

return self.register_orchestration_versioned(name, version, fn)

def start(self):

"""Start the runtime. Blocks until shutdown is called."""

self._native.start()

def shutdown(self, timeout_ms: int = None):

"""Shutdown the runtime gracefully."""

self._native.shutdown(timeout_ms)

def metrics_snapshot(self):

"""Get a snapshot of runtime metrics."""

return self._native.metrics_snapshot()

# Tag limits

MAX_WORKER_TAGS = 5

MAX_TAG_NAME_BYTES = 256

MAX_KV_KEYS = 150

MAX_KV_VALUE_BYTES = 65536

class TagFilter:

"""Helper for constructing worker tag filter values.

Usage with RuntimeOptions:

RuntimeOptions(worker_tag_filter=TagFilter.DEFAULT_ONLY)

RuntimeOptions(worker_tag_filter=TagFilter.tags(["gpu", "cpu"]))

RuntimeOptions(worker_tag_filter=TagFilter.default_and(["gpu"]))

RuntimeOptions(worker_tag_filter=TagFilter.ANY)

RuntimeOptions(worker_tag_filter=TagFilter.NONE)

"""

DEFAULT_ONLY = "default_only"

ANY = "any"

NONE = "none"

@staticmethod

def tags(tags: list) -> str:

"""Process only activities with the specified tags (not untagged)."""

return json.dumps({"tags": list(tags)})

@staticmethod

def default_and(tags: list) -> str:

"""Process untagged activities AND activities with the specified tags."""

return json.dumps({"default_and": list(tags)})

__all__ = [

"SqliteProvider",

"PostgresProvider",

"PostgresEntraOptions",

"Client",

"Runtime",

"OrchestrationResult",

"RuntimeOptions",

"OrchestrationContext",

"ActivityContext",

"ScheduledTask",

"TagFilter",

"OrchestrationStatus",

"SystemMetrics",

"SystemStats",

"QueueDepths",

"InstanceInfo",

"ExecutionInfo",

"InstanceTree",

"DeleteInstanceResult",

"PruneOptions",

"PruneResult",

"InstanceFilter",

"Event",

"MetricsSnapshot",

"init_tracing",

"parse_result",

"MAX_WORKER_TAGS",

"MAX_TAG_NAME_BYTES",

"MAX_KV_KEYS",

"MAX_KV_VALUE_BYTES",

# Backwards-compatibility aliases

"PyRuntimeOptions",

"PyOrchestrationStatus",

"PySystemMetrics",

"PySystemStats",

"PyQueueDepths",

"PyInstanceInfo",

"PyExecutionInfo",

"PyInstanceTree",

"PyDeleteInstanceResult",

"PyPruneOptions",

"PyPruneResult",

"PyInstanceFilter",

"PyEvent",

"PyMetricsSnapshot",

]