dbos 2.1.0a2__py3-none-any.whl → 2.4.0a7__py3-none-any.whl

dbos/__init__.py

@@ -12,6 +12,7 @@ from ._dbos_config import DBOSConfig
 from ._debouncer import Debouncer, DebouncerClient
 from ._kafka_message import KafkaMessage
 from ._queue import Queue
+from ._serialization import Serializer
 from ._sys_db import GetWorkflowsInput, WorkflowStatus, WorkflowStatusString
 
 __all__ = [
@@ -35,4 +36,5 @@ __all__ = [
     "Queue",
     "Debouncer",
     "DebouncerClient",
+    "Serializer",
 ]

dbos/_admin_server.py

@@ -338,6 +338,7 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             end_time=filters.get("end_time"),
             status=filters.get("status"),
             app_version=filters.get("application_version"),
+            forked_from=filters.get("forked_from"),
             name=filters.get("workflow_name"),
             limit=filters.get("limit"),
             offset=filters.get("offset"),
@@ -364,6 +365,7 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             start_time=filters.get("start_time"),
             end_time=filters.get("end_time"),
             status=filters.get("status"),
+            forked_from=filters.get("forked_from"),
             name=filters.get("workflow_name"),
             limit=filters.get("limit"),
             offset=filters.get("offset"),

dbos/_app_db.py

@@ -8,8 +8,8 @@ from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
 from dbos._migration import get_sqlite_timestamp_expr
+from dbos._serialization import Serializer
 
-from . import _serialization
 from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
 from ._logger import dbos_logger
 from ._schemas.application_database import ApplicationSchema
@@ -34,17 +34,52 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase(ABC):
 
+    @staticmethod
+    def create(
+        database_url: str,
+        engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
+        serializer: Serializer,
+        debug_mode: bool = False,
+    ) -> "ApplicationDatabase":
+        """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
+        if database_url.startswith("sqlite"):
+            return SQLiteApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+        else:
+            # Default to PostgreSQL for postgresql://, postgres://, or other URLs
+            return PostgresApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+
     def __init__(
         self,
         *,
         database_url: str,
         engine_kwargs: Dict[str, Any],
+        serializer: Serializer,
+        schema: Optional[str],
         debug_mode: bool = False,
     ):
+        if database_url.startswith("sqlite"):
+            self.schema = None
+        else:
+            self.schema = schema if schema else "dbos"
+        ApplicationSchema.transaction_outputs.schema = schema
         self.engine = self._create_engine(database_url, engine_kwargs)
         self._engine_kwargs = engine_kwargs
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode
+        self.serializer = serializer
 
     @abstractmethod
     def _create_engine(
@@ -156,14 +191,18 @@ class ApplicationDatabase(ABC):
                 function_id=row[0],
                 function_name=row[1],
                 output=(
-                    _serialization.deserialize(row[2]) if row[2] is not None else row[2]
+                    self.serializer.deserialize(row[2])
+                    if row[2] is not None
+                    else row[2]
                 ),
                 error=(
-                    _serialization.deserialize_exception(row[3])
+                    self.serializer.deserialize(row[3])
                     if row[3] is not None
                     else row[3]
                 ),
                 child_workflow_id=None,
+                started_at_epoch_ms=None,
+                completed_at_epoch_ms=None,
             )
             for row in rows
         ]
@@ -237,52 +276,10 @@ class ApplicationDatabase(ABC):
         """Check if the error is a serialization/concurrency error."""
         pass
 
-    @staticmethod
-    def create(
-        database_url: str,
-        engine_kwargs: Dict[str, Any],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ) -> "ApplicationDatabase":
-        """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
-        if database_url.startswith("sqlite"):
-            return SQLiteApplicationDatabase(
-                database_url=database_url,
-                engine_kwargs=engine_kwargs,
-                debug_mode=debug_mode,
-            )
-        else:
-            # Default to PostgreSQL for postgresql://, postgres://, or other URLs
-            return PostgresApplicationDatabase(
-                database_url=database_url,
-                engine_kwargs=engine_kwargs,
-                debug_mode=debug_mode,
-                schema=schema,
-            )
-
 
 class PostgresApplicationDatabase(ApplicationDatabase):
     """PostgreSQL-specific implementation of ApplicationDatabase."""
 
-    def __init__(
-        self,
-        *,
-        database_url: str,
-        engine_kwargs: Dict[str, Any],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ):
-        super().__init__(
-            database_url=database_url,
-            engine_kwargs=engine_kwargs,
-            debug_mode=debug_mode,
-        )
-        if schema is None:
-            self.schema = "dbos"
-        else:
-            self.schema = schema
-        ApplicationSchema.transaction_outputs.schema = schema
-
     def _create_engine(
         self, database_url: str, engine_kwargs: Dict[str, Any]
     ) -> sa.Engine:

dbos/_client.py

@@ -1,6 +1,6 @@
 import asyncio
+import json
 import time
-import uuid
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -16,10 +16,10 @@ from typing import (
 
 import sqlalchemy as sa
 
-from dbos import _serialization
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
 from dbos._sys_db import SystemDatabase
+from dbos._utils import generate_uuid
 
 if TYPE_CHECKING:
     from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
@@ -27,7 +27,7 @@ if TYPE_CHECKING:
 from dbos._dbos_config import get_system_database_url, is_valid_database_url
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
-from dbos._serialization import WorkflowInputs
+from dbos._serialization import DefaultSerializer, Serializer, WorkflowInputs
 from dbos._sys_db import (
     EnqueueOptionsInternal,
     StepInfo,
@@ -63,6 +63,9 @@ class EnqueueOptions(_EnqueueOptionsRequired, total=False):
     deduplication_id: str
     priority: int
     max_recovery_attempts: int
+    queue_partition_key: str
+    authenticated_user: str
+    authenticated_roles: list[str]
 
 
 def validate_enqueue_options(options: EnqueueOptions) -> None:
@@ -127,7 +130,9 @@ class DBOSClient:
         system_database_engine: Optional[sa.Engine] = None,
         application_database_url: Optional[str] = None,
         dbos_system_schema: Optional[str] = "dbos",
+        serializer: Serializer = DefaultSerializer(),
     ):
+        self._serializer = serializer
         application_database_url = (
             database_url if database_url else application_database_url
         )
@@ -144,23 +149,30 @@ class DBOSClient:
         self._sys_db = SystemDatabase.create(
             system_database_url=system_database_url,
             engine_kwargs={
+                "connect_args": {"application_name": "dbos_transact_client"},
                 "pool_timeout": 30,
                 "max_overflow": 0,
                 "pool_size": 2,
+                "pool_pre_ping": True,
             },
             engine=system_database_engine,
             schema=dbos_system_schema,
+            serializer=serializer,
+            executor_id=None,
         )
         self._sys_db.check_connection()
         if application_database_url:
             self._app_db = ApplicationDatabase.create(
                 database_url=application_database_url,
                 engine_kwargs={
+                    "connect_args": {"application_name": "dbos_transact_client"},
                     "pool_timeout": 30,
                     "max_overflow": 0,
                     "pool_size": 2,
+                    "pool_pre_ping": True,
                 },
                 schema=dbos_system_schema,
+                serializer=serializer,
             )
 
     def destroy(self) -> None:
@@ -176,14 +188,22 @@ class DBOSClient:
             max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
         workflow_id = options.get("workflow_id")
         if workflow_id is None:
-            workflow_id = str(uuid.uuid4())
+            workflow_id = generate_uuid()
         workflow_timeout = options.get("workflow_timeout", None)
         enqueue_options_internal: EnqueueOptionsInternal = {
             "deduplication_id": options.get("deduplication_id"),
             "priority": options.get("priority"),
             "app_version": options.get("app_version"),
+            "queue_partition_key": options.get("queue_partition_key"),
         }
 
+        authenticated_user = options.get("authenticated_user")
+        authenticated_roles = (
+            json.dumps(options.get("authenticated_roles"))
+            if options.get("authenticated_roles")
+            else None
+        )
+
         inputs: WorkflowInputs = {
             "args": args,
             "kwargs": kwargs,
@@ -197,9 +217,9 @@ class DBOSClient:
             "queue_name": queue_name,
             "app_version": enqueue_options_internal["app_version"],
             "config_name": None,
-            "authenticated_user": None,
+            "authenticated_user": authenticated_user,
             "assumed_role": None,
-            "authenticated_roles": None,
+            "authenticated_roles": authenticated_roles,
             "output": None,
             "error": None,
             "created_at": None,
@@ -217,7 +237,9 @@ class DBOSClient:
                 if enqueue_options_internal["priority"] is not None
                 else 0
             ),
-            "inputs": _serialization.serialize_args(inputs),
+            "inputs": self._serializer.serialize(inputs),
+            "queue_partition_key": enqueue_options_internal["queue_partition_key"],
+            "forked_from": None,
         }
 
         self._sys_db.init_workflow(
@@ -259,7 +281,7 @@ class DBOSClient:
         topic: Optional[str] = None,
         idempotency_key: Optional[str] = None,
     ) -> None:
-        idempotency_key = idempotency_key if idempotency_key else str(uuid.uuid4())
+        idempotency_key = idempotency_key if idempotency_key else generate_uuid()
         status: WorkflowStatusInternal = {
             "workflow_uuid": f"{destination_id}-{idempotency_key}",
             "status": WorkflowStatusString.SUCCESS.value,
@@ -282,7 +304,9 @@ class DBOSClient:
             "workflow_deadline_epoch_ms": None,
             "deduplication_id": None,
             "priority": 0,
-            "inputs": _serialization.serialize_args({"args": (), "kwargs": {}}),
+            "inputs": self._serializer.serialize({"args": (), "kwargs": {}}),
+            "queue_partition_key": None,
+            "forked_from": None,
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db._insert_workflow_status(

dbos/_conductor/conductor.py

@@ -2,7 +2,6 @@ import socket
 import threading
 import time
 import traceback
-import uuid
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Optional
 
@@ -11,7 +10,7 @@ from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
 from dbos._context import SetWorkflowID
-from dbos._utils import GlobalParams
+from dbos._utils import GlobalParams, generate_uuid
 from dbos._workflow_commands import (
     garbage_collect,
     get_workflow,
@@ -192,7 +191,7 @@ class ConductorWebsocket(threading.Thread):
                             fork_message = p.ForkWorkflowRequest.from_json(message)
                             new_workflow_id = fork_message.body["new_workflow_id"]
                             if new_workflow_id is None:
-                                new_workflow_id = str(uuid.uuid4())
+                                new_workflow_id = generate_uuid()
                             workflow_id = fork_message.body["workflow_id"]
                             start_step = fork_message.body["start_step"]
                             app_version = fork_message.body["application_version"]
@@ -223,22 +222,21 @@ class ConductorWebsocket(threading.Thread):
                             body = list_workflows_message.body
                             infos = []
                             try:
-                                load_input = body.get("load_input", False)
-                                load_output = body.get("load_output", False)
                                 infos = list_workflows(
                                     self.dbos._sys_db,
-                                    workflow_ids=body["workflow_uuids"],
-                                    user=body["authenticated_user"],
-                                    start_time=body["start_time"],
-                                    end_time=body["end_time"],
-                                    status=body["status"],
-                                    app_version=body["application_version"],
-                                    name=body["workflow_name"],
-                                    limit=body["limit"],
-                                    offset=body["offset"],
-                                    sort_desc=body["sort_desc"],
-                                    load_input=load_input,
-                                    load_output=load_output,
+                                    workflow_ids=body.get("workflow_uuids", None),
+                                    user=body.get("authenticated_user", None),
+                                    start_time=body.get("start_time", None),
+                                    end_time=body.get("end_time", None),
+                                    status=body.get("status", None),
+                                    app_version=body.get("application_version", None),
+                                    forked_from=body.get("forked_from", None),
+                                    name=body.get("workflow_name", None),
+                                    limit=body.get("limit", None),
+                                    offset=body.get("offset", None),
+                                    sort_desc=body.get("sort_desc", False),
+                                    load_input=body.get("load_input", False),
+                                    load_output=body.get("load_output", False),
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when listing workflows: {traceback.format_exc()}"
@@ -261,18 +259,18 @@ class ConductorWebsocket(threading.Thread):
                             q_body = list_queued_workflows_message.body
                             infos = []
                             try:
-                                q_load_input = q_body.get("load_input", False)
                                 infos = list_queued_workflows(
                                     self.dbos._sys_db,
-                                    start_time=q_body["start_time"],
-                                    end_time=q_body["end_time"],
-                                    status=q_body["status"],
-                                    name=q_body["workflow_name"],
-                                    limit=q_body["limit"],
-                                    offset=q_body["offset"],
-                                    queue_name=q_body["queue_name"],
-                                    sort_desc=q_body["sort_desc"],
-                                    load_input=q_load_input,
+                                    start_time=q_body.get("start_time", None),
+                                    end_time=q_body.get("end_time", None),
+                                    status=q_body.get("status", None),
+                                    forked_from=q_body.get("forked_from", None),
+                                    name=q_body.get("workflow_name", None),
+                                    limit=q_body.get("limit", None),
+                                    offset=q_body.get("offset", None),
+                                    queue_name=q_body.get("queue_name", None),
+                                    sort_desc=q_body.get("sort_desc", False),
+                                    load_input=q_body.get("load_input", False),
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when listing queued workflows: {traceback.format_exc()}"

dbos/_conductor/protocol.py

@@ -118,6 +118,7 @@ class ListWorkflowsBody(TypedDict, total=False):
     end_time: Optional[str]
     status: Optional[str]
     application_version: Optional[str]
+    forked_from: Optional[str]
     limit: Optional[int]
     offset: Optional[int]
     sort_desc: bool
@@ -143,6 +144,12 @@ class WorkflowsOutput:
     QueueName: Optional[str]
     ApplicationVersion: Optional[str]
     ExecutorID: Optional[str]
+    WorkflowTimeoutMS: Optional[str]
+    WorkflowDeadlineEpochMS: Optional[str]
+    DeduplicationID: Optional[str]
+    Priority: Optional[str]
+    QueuePartitionKey: Optional[str]
+    ForkedFrom: Optional[str]
 
     @classmethod
     def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
@@ -152,12 +159,22 @@ class WorkflowsOutput:
         inputs_str = str(info.input) if info.input is not None else None
         outputs_str = str(info.output) if info.output is not None else None
         error_str = str(info.error) if info.error is not None else None
-        request_str = None
         roles_str = (
             str(info.authenticated_roles)
             if info.authenticated_roles is not None
             else None
         )
+        workflow_timeout_ms_str = (
+            str(info.workflow_timeout_ms)
+            if info.workflow_timeout_ms is not None
+            else None
+        )
+        workflow_deadline_epoch_ms_str = (
+            str(info.workflow_deadline_epoch_ms)
+            if info.workflow_deadline_epoch_ms is not None
+            else None
+        )
+        priority_str = str(info.priority) if info.priority is not None else None
 
         return cls(
             WorkflowUUID=info.workflow_id,
@@ -176,6 +193,12 @@ class WorkflowsOutput:
             QueueName=info.queue_name,
             ApplicationVersion=info.app_version,
             ExecutorID=info.executor_id,
+            WorkflowTimeoutMS=workflow_timeout_ms_str,
+            WorkflowDeadlineEpochMS=workflow_deadline_epoch_ms_str,
+            DeduplicationID=info.deduplication_id,
+            Priority=priority_str,
+            QueuePartitionKey=info.queue_partition_key,
+            ForkedFrom=info.forked_from,
         )
 
 
@@ -186,14 +209,28 @@ class WorkflowSteps:
     output: Optional[str]
     error: Optional[str]
     child_workflow_id: Optional[str]
+    started_at_epoch_ms: Optional[str]
+    completed_at_epoch_ms: Optional[str]
 
     @classmethod
     def from_step_info(cls, info: StepInfo) -> "WorkflowSteps":
         output_str = str(info["output"]) if info["output"] is not None else None
         error_str = str(info["error"]) if info["error"] is not None else None
+        started_at_str = (
+            str(info["started_at_epoch_ms"])
+            if info["started_at_epoch_ms"] is not None
+            else None
+        )
+        completed_at_str = (
+            str(info["completed_at_epoch_ms"])
+            if info["completed_at_epoch_ms"] is not None
+            else None
+        )
         return cls(
             function_id=info["function_id"],
             function_name=info["function_name"],
+            started_at_epoch_ms=started_at_str,
+            completed_at_epoch_ms=completed_at_str,
             output=output_str,
             error=error_str,
             child_workflow_id=info["child_workflow_id"],
@@ -216,6 +253,7 @@ class ListQueuedWorkflowsBody(TypedDict, total=False):
     start_time: Optional[str]
     end_time: Optional[str]
     status: Optional[str]
+    forked_from: Optional[str]
     queue_name: Optional[str]
     limit: Optional[int]
     offset: Optional[int]

dbos/_context.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import json
 import os
-import uuid
 from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -15,7 +14,7 @@ if TYPE_CHECKING:
 
 from sqlalchemy.orm import Session
 
-from dbos._utils import GlobalParams
+from dbos._utils import GlobalParams, generate_uuid
 
 from ._logger import dbos_logger
 from ._tracer import dbos_tracer
@@ -120,6 +119,8 @@ class DBOSContext:
         self.deduplication_id: Optional[str] = None
         # A user-specified priority for the enqueuing workflow.
         self.priority: Optional[int] = None
+        # If the workflow is enqueued on a partitioned queue, its partition key
+        self.queue_partition_key: Optional[str] = None
 
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
@@ -149,7 +150,7 @@ class DBOSContext:
                 self.logger.warning(
                     f"Multiple workflows started in the same SetWorkflowID block. Only the first workflow is assigned the specified workflow ID; subsequent workflows will use a generated workflow ID."
                 )
-            wfid = str(uuid.uuid4())
+            wfid = generate_uuid()
         return wfid
 
     def start_workflow(
@@ -479,6 +480,7 @@ class SetEnqueueOptions:
         deduplication_id: Optional[str] = None,
         priority: Optional[int] = None,
         app_version: Optional[str] = None,
+        queue_partition_key: Optional[str] = None,
     ) -> None:
         self.created_ctx = False
         self.deduplication_id: Optional[str] = deduplication_id
@@ -491,6 +493,8 @@ class SetEnqueueOptions:
         self.saved_priority: Optional[int] = None
         self.app_version: Optional[str] = app_version
         self.saved_app_version: Optional[str] = None
+        self.queue_partition_key = queue_partition_key
+        self.saved_queue_partition_key: Optional[str] = None
 
     def __enter__(self) -> SetEnqueueOptions:
         # Code to create a basic context
@@ -505,6 +509,8 @@ class SetEnqueueOptions:
         ctx.priority = self.priority
         self.saved_app_version = ctx.app_version
         ctx.app_version = self.app_version
+        self.saved_queue_partition_key = ctx.queue_partition_key
+        ctx.queue_partition_key = self.queue_partition_key
         return self
 
     def __exit__(
@@ -517,6 +523,7 @@ class SetEnqueueOptions:
         curr_ctx.deduplication_id = self.saved_deduplication_id
         curr_ctx.priority = self.saved_priority
         curr_ctx.app_version = self.saved_app_version
+        curr_ctx.queue_partition_key = self.saved_queue_partition_key
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()