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

dbos/_core.py

@@ -23,7 +23,6 @@ from typing import (
 from dbos._outcome import Immediate, NoResult, Outcome, Pending
 from dbos._utils import GlobalParams, retriable_postgres_exception
 
-from . import _serialization
 from ._app_db import ApplicationDatabase, TransactionResultInternal
 from ._context import (
     DBOSAssumeRole,
@@ -94,14 +93,6 @@ TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
 DEBOUNCER_WORKFLOW_NAME = "_dbos_debouncer_workflow"
 
 
-def check_is_in_coroutine() -> bool:
-    try:
-        asyncio.get_running_loop()
-        return True
-    except RuntimeError:
-        return False
-
-
 class WorkflowHandleFuture(Generic[R]):
 
     def __init__(self, workflow_id: str, future: Future[R], dbos: "DBOS"):
@@ -116,10 +107,10 @@ class WorkflowHandleFuture(Generic[R]):
         try:
             r = self.future.result()
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
@@ -143,10 +134,10 @@ class WorkflowHandlePolling(Generic[R]):
         try:
             r: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         self.dbos._sys_db.record_get_result(self.workflow_id, serialized_r, None)
         return r
 
@@ -171,7 +162,7 @@ class WorkflowHandleAsyncTask(Generic[R]):
         try:
             r = await self.task
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             await asyncio.to_thread(
                 self.dbos._sys_db.record_get_result,
                 self.workflow_id,
@@ -179,7 +170,7 @@ class WorkflowHandleAsyncTask(Generic[R]):
                 serialized_e,
             )
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         await asyncio.to_thread(
             self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
         )
@@ -207,7 +198,7 @@ class WorkflowHandleAsyncPolling(Generic[R]):
                 self.dbos._sys_db.await_workflow_result, self.workflow_id
             )
         except Exception as e:
-            serialized_e = _serialization.serialize_exception(e)
+            serialized_e = self.dbos._serializer.serialize(e)
             await asyncio.to_thread(
                 self.dbos._sys_db.record_get_result,
                 self.workflow_id,
@@ -215,7 +206,7 @@ class WorkflowHandleAsyncPolling(Generic[R]):
                 serialized_e,
             )
             raise
-        serialized_r = _serialization.serialize(r)
+        serialized_r = self.dbos._serializer.serialize(r)
         await asyncio.to_thread(
             self.dbos._sys_db.record_get_result, self.workflow_id, serialized_r, None
         )
@@ -303,7 +294,13 @@ def _init_workflow(
             if enqueue_options is not None
             else 0
         ),
-        "inputs": _serialization.serialize_args(inputs),
+        "inputs": dbos._serializer.serialize(inputs),
+        "queue_partition_key": (
+            enqueue_options["queue_partition_key"]
+            if enqueue_options is not None
+            else None
+        ),
+        "forked_from": None,
     }
 
     # Synchronously record the status and inputs for workflows
@@ -319,7 +316,8 @@ def _init_workflow(
                 "function_id": ctx.parent_workflow_fid,
                 "function_name": wf_name,
                 "output": None,
-                "error": _serialization.serialize_exception(e),
+                "error": dbos._serializer.serialize(e),
+                "started_at_epoch_ms": int(time.time() * 1000),
             }
             dbos._sys_db.record_operation_result(result)
         raise
@@ -378,7 +376,7 @@ def _get_wf_invoke_func(
                 dbos._sys_db.update_workflow_outcome(
                     status["workflow_uuid"],
                     "SUCCESS",
-                    output=_serialization.serialize(output),
+                    output=dbos._serializer.serialize(output),
                 )
             return output
         except DBOSWorkflowConflictIDError:
@@ -392,7 +390,7 @@ def _get_wf_invoke_func(
                 dbos._sys_db.update_workflow_outcome(
                     status["workflow_uuid"],
                     "ERROR",
-                    error=_serialization.serialize_exception(error),
+                    error=dbos._serializer.serialize(error),
                 )
             raise
         finally:
@@ -464,7 +462,7 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
-    inputs = _serialization.deserialize_args(status["inputs"])
+    inputs: WorkflowInputs = dbos._serializer.deserialize(status["inputs"])
     wf_func = dbos._registry.workflow_info_map.get(status["name"], None)
     if not wf_func:
         raise DBOSWorkflowFunctionNotFoundError(
@@ -572,6 +570,9 @@ def start_workflow(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
@@ -665,6 +666,9 @@ async def start_workflow_async(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
@@ -837,20 +841,15 @@ def workflow_wrapper(
             try:
                 r = func()
             except Exception as e:
-                serialized_e = _serialization.serialize_exception(e)
+                serialized_e = dbos._serializer.serialize(e)
                 assert workflow_id is not None
                 dbos._sys_db.record_get_result(workflow_id, None, serialized_e)
                 raise
-            serialized_r = _serialization.serialize(r)
+            serialized_r = dbos._serializer.serialize(r)
             assert workflow_id is not None
             dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
             return r
 
-        if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
-            dbos_logger.warning(
-                f"Sync workflow ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
-            )
-
         outcome = (
             wfOutcome.wrap(init_wf, dbos=dbos)
             .also(DBOSAssumeRole(rr))
@@ -948,15 +947,15 @@ def decorate_transaction(
                                         f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
                                     )
                                     if recorded_output["error"]:
-                                        deserialized_error = (
-                                            _serialization.deserialize_exception(
+                                        deserialized_error: Exception = (
+                                            dbos._serializer.deserialize(
                                                 recorded_output["error"]
                                             )
                                         )
                                         has_recorded_error = True
                                         raise deserialized_error
                                     elif recorded_output["output"]:
-                                        return _serialization.deserialize(
+                                        return dbos._serializer.deserialize(
                                             recorded_output["output"]
                                         )
                                     else:
@@ -969,7 +968,9 @@ def decorate_transaction(
                                     )
 
                                 output = func(*args, **kwargs)
-                                txn_output["output"] = _serialization.serialize(output)
+                                txn_output["output"] = dbos._serializer.serialize(
+                                    output
+                                )
                                 assert (
                                     ctx.sql_session is not None
                                 ), "Cannot find a database connection"
@@ -1010,8 +1011,8 @@ def decorate_transaction(
                         finally:
                             # Don't record the error if it was already recorded
                             if txn_error and not has_recorded_error:
-                                txn_output["error"] = (
-                                    _serialization.serialize_exception(txn_error)
+                                txn_output["error"] = dbos._serializer.serialize(
+                                    txn_error
                                 )
                                 dbos._app_db.record_transaction_error(txn_output)
             return output
@@ -1034,10 +1035,6 @@ def decorate_transaction(
                 assert (
                     ctx.is_workflow()
                 ), "Transactions must be called from within workflows"
-                if check_is_in_coroutine():
-                    dbos_logger.warning(
-                        f"Transaction function ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Use asyncio.to_thread instead."
-                    )
                 with DBOSAssumeRole(rr):
                     return invoke_tx(*args, **kwargs)
             else:
@@ -1123,15 +1120,16 @@ def decorate_step(
                     "function_name": step_name,
                     "output": None,
                     "error": None,
+                    "started_at_epoch_ms": int(time.time() * 1000),
                 }
 
                 try:
                     output = func()
                 except Exception as error:
-                    step_output["error"] = _serialization.serialize_exception(error)
+                    step_output["error"] = dbos._serializer.serialize(error)
                     dbos._sys_db.record_operation_result(step_output)
                     raise
-                step_output["output"] = _serialization.serialize(output)
+                step_output["output"] = dbos._serializer.serialize(output)
                 dbos._sys_db.record_operation_result(step_output)
                 return output
 
@@ -1147,13 +1145,13 @@ def decorate_step(
                         f"Replaying step, id: {ctx.function_id}, name: {attributes['name']}"
                     )
                     if recorded_output["error"] is not None:
-                        deserialized_error = _serialization.deserialize_exception(
+                        deserialized_error: Exception = dbos._serializer.deserialize(
                             recorded_output["error"]
                         )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
                         return cast(
-                            R, _serialization.deserialize(recorded_output["output"])
+                            R, dbos._serializer.deserialize(recorded_output["output"])
                         )
                     else:
                         raise Exception("Output and error are both None")
@@ -1182,10 +1180,6 @@ def decorate_step(
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
-            if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
-                dbos_logger.warning(
-                    f"Sync step ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
-                )
             # If the step is called from a workflow, run it as a step.
             # Otherwise, run it as a normal function.
             ctx = get_local_dbos_context()

dbos/_dbos.py

@@ -7,7 +7,6 @@ import os
 import sys
 import threading
 import time
-import uuid
 from concurrent.futures import ThreadPoolExecutor
 from logging import Logger
 from typing import (
@@ -31,8 +30,9 @@ from typing import (
 
 from dbos._conductor.conductor import ConductorWebsocket
 from dbos._debouncer import debouncer_workflow
+from dbos._serialization import DefaultSerializer, Serializer
 from dbos._sys_db import SystemDatabase, WorkflowStatus
-from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
+from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams, generate_uuid
 from dbos._workflow_commands import fork_workflow, list_queued_workflows, list_workflows
 
 from ._classproperty import classproperty
@@ -341,6 +341,8 @@ class DBOS:
         self.conductor_websocket: Optional[ConductorWebsocket] = None
         self._background_event_loop: BackgroundEventLoop = BackgroundEventLoop()
         self._active_workflows_set: set[str] = set()
+        serializer = config.get("serializer")
+        self._serializer: Serializer = serializer if serializer else DefaultSerializer()
 
         # Globally set the application version and executor ID.
         # In DBOS Cloud, instead use the values supplied through environment variables.
@@ -441,7 +443,7 @@ class DBOS:
             if GlobalParams.app_version == "":
                 GlobalParams.app_version = self._registry.compute_app_version()
             if self.conductor_key is not None:
-                GlobalParams.executor_id = str(uuid.uuid4())
+                GlobalParams.executor_id = generate_uuid()
             dbos_logger.info(f"Executor ID: {GlobalParams.executor_id}")
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=sys.maxsize)
@@ -449,28 +451,35 @@ class DBOS:
             assert self._config["database"]["sys_db_engine_kwargs"] is not None
             # Get the schema configuration, use "dbos" as default
             schema = self._config.get("dbos_system_schema", "dbos")
+            dbos_logger.debug("Creating system database")
             self._sys_db_field = SystemDatabase.create(
                 system_database_url=get_system_database_url(self._config),
                 engine_kwargs=self._config["database"]["sys_db_engine_kwargs"],
                 engine=self._config["system_database_engine"],
                 debug_mode=debug_mode,
                 schema=schema,
+                serializer=self._serializer,
+                executor_id=GlobalParams.executor_id,
             )
             assert self._config["database"]["db_engine_kwargs"] is not None
             if self._config["database_url"]:
+                dbos_logger.debug("Creating application database")
                 self._app_db_field = ApplicationDatabase.create(
                     database_url=self._config["database_url"],
                     engine_kwargs=self._config["database"]["db_engine_kwargs"],
                     debug_mode=debug_mode,
                     schema=schema,
+                    serializer=self._serializer,
                 )
 
             if debug_mode:
                 return
 
             # Run migrations for the system and application databases
+            dbos_logger.debug("Running system database migrations")
             self._sys_db.run_migrations()
             if self._app_db:
+                dbos_logger.debug("Running application database migrations")
                 self._app_db.run_migrations()
 
             admin_port = self._config.get("runtimeConfig", {}).get("admin_port")
@@ -481,25 +490,29 @@ class DBOS:
             )
             if run_admin_server:
                 try:
+                    dbos_logger.debug("Starting admin server")
                     self._admin_server_field = AdminServer(dbos=self, port=admin_port)
                 except Exception as e:
                     dbos_logger.warning(f"Failed to start admin server: {e}")
 
-            workflow_ids = self._sys_db.get_pending_workflows(
-                GlobalParams.executor_id, GlobalParams.app_version
-            )
-            if (len(workflow_ids)) > 0:
-                self.logger.info(
-                    f"Recovering {len(workflow_ids)} workflows from application version {GlobalParams.app_version}"
-                )
-            else:
-                self.logger.info(
-                    f"No workflows to recover from application version {GlobalParams.app_version}"
+            # Recover local workflows if not using a recovery service
+            if not self.conductor_key and not GlobalParams.dbos_cloud:
+                dbos_logger.debug("Retrieving local pending workflows for recovery")
+                workflow_ids = self._sys_db.get_pending_workflows(
+                    GlobalParams.executor_id, GlobalParams.app_version
                 )
-
-            self._executor.submit(startup_recovery_thread, self, workflow_ids)
+                if (len(workflow_ids)) > 0:
+                    self.logger.info(
+                        f"Recovering {len(workflow_ids)} workflows from application version {GlobalParams.app_version}"
+                    )
+                else:
+                    self.logger.info(
+                        f"No workflows to recover from application version {GlobalParams.app_version}"
+                    )
+                self._executor.submit(startup_recovery_thread, self, workflow_ids)
 
             # Listen to notifications
+            dbos_logger.debug("Starting notifications listener thread")
             notification_listener_thread = threading.Thread(
                 target=self._sys_db._notification_listener,
                 daemon=True,
@@ -511,6 +524,7 @@ class DBOS:
             self._registry.get_internal_queue()
 
             # Start the queue thread
+            dbos_logger.debug("Starting queue thread")
             evt = threading.Event()
             self.background_thread_stop_events.append(evt)
             bg_queue_thread = threading.Thread(
@@ -526,6 +540,7 @@ class DBOS:
                     self.conductor_url = f"wss://{dbos_domain}/conductor/v1alpha1"
                 evt = threading.Event()
                 self.background_thread_stop_events.append(evt)
+                dbos_logger.debug("Starting Conductor thread")
                 self.conductor_websocket = ConductorWebsocket(
                     self,
                     conductor_url=self.conductor_url,
@@ -536,6 +551,7 @@ class DBOS:
                 self._background_threads.append(self.conductor_websocket)
 
             # Grab any pollers that were deferred and start them
+            dbos_logger.debug("Starting event receivers")
             for evt, func, args, kwargs in self._registry.pollers:
                 self.poller_stop_events.append(evt)
                 poller_thread = threading.Thread(
@@ -1112,7 +1128,9 @@ class DBOS:
         end_time: Optional[str] = None,
         name: Optional[str] = None,
         app_version: Optional[str] = None,
+        forked_from: Optional[str] = None,
         user: Optional[str] = None,
+        queue_name: Optional[str] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         sort_desc: bool = False,
@@ -1129,6 +1147,7 @@ class DBOS:
                 end_time=end_time,
                 name=name,
                 app_version=app_version,
+                forked_from=forked_from,
                 user=user,
                 limit=limit,
                 offset=offset,
@@ -1136,6 +1155,7 @@ class DBOS:
                 workflow_id_prefix=workflow_id_prefix,
                 load_input=load_input,
                 load_output=load_output,
+                queue_name=queue_name,
             )
 
         return _get_dbos_instance()._sys_db.call_function_as_step(
@@ -1152,6 +1172,7 @@ class DBOS:
         end_time: Optional[str] = None,
         name: Optional[str] = None,
         app_version: Optional[str] = None,
+        forked_from: Optional[str] = None,
         user: Optional[str] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
@@ -1169,6 +1190,7 @@ class DBOS:
             end_time=end_time,
             name=name,
             app_version=app_version,
+            forked_from=forked_from,
             user=user,
             limit=limit,
             offset=offset,
@@ -1184,6 +1206,7 @@ class DBOS:
         *,
         queue_name: Optional[str] = None,
         status: Optional[Union[str, List[str]]] = None,
+        forked_from: Optional[str] = None,
         start_time: Optional[str] = None,
         end_time: Optional[str] = None,
         name: Optional[str] = None,
@@ -1197,6 +1220,7 @@ class DBOS:
                 _get_dbos_instance()._sys_db,
                 queue_name=queue_name,
                 status=status,
+                forked_from=forked_from,
                 start_time=start_time,
                 end_time=end_time,
                 name=name,
@@ -1216,6 +1240,7 @@ class DBOS:
         *,
         queue_name: Optional[str] = None,
         status: Optional[Union[str, List[str]]] = None,
+        forked_from: Optional[str] = None,
         start_time: Optional[str] = None,
         end_time: Optional[str] = None,
         name: Optional[str] = None,
@@ -1229,6 +1254,7 @@ class DBOS:
             cls.list_queued_workflows,
             queue_name=queue_name,
             status=status,
+            forked_from=forked_from,
             start_time=start_time,
             end_time=end_time,
             name=name,

dbos/_dbos_config.py

@@ -7,6 +7,8 @@ import sqlalchemy as sa
 import yaml
 from sqlalchemy import make_url
 
+from dbos._serialization import Serializer
+
 from ._error import DBOSInitializationError
 from ._logger import dbos_logger
 from ._schemas.system_database import SystemSchema
@@ -37,6 +39,7 @@ class DBOSConfig(TypedDict, total=False):
         enable_otlp (bool): If True, enable built-in DBOS OTLP tracing and logging.
         system_database_engine (sa.Engine): A custom system database engine. If provided, DBOS will not create an engine but use this instead.
         conductor_key (str): An API key for DBOS Conductor. Pass this in to connect your process to Conductor.
+        serializer (Serializer): A custom serializer and deserializer DBOS uses when storing program data in the system database
     """
 
     name: str
@@ -57,6 +60,7 @@ class DBOSConfig(TypedDict, total=False):
     enable_otlp: Optional[bool]
     system_database_engine: Optional[sa.Engine]
     conductor_key: Optional[str]
+    serializer: Optional[Serializer]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -67,16 +71,6 @@ class RuntimeConfig(TypedDict, total=False):
 
 
 class DatabaseConfig(TypedDict, total=False):
-    """
-    Internal data structure containing the DBOS database configuration.
-    Attributes:
-        sys_db_name (str): System database name
-        sys_db_pool_size (int): System database pool size
-        db_engine_kwargs (Dict[str, Any]): SQLAlchemy engine kwargs
-        migrate (List[str]): Migration commands to run on startup
-        dbos_system_schema (str): Schema name for DBOS system tables. Defaults to "dbos".
-    """
-
     sys_db_pool_size: Optional[int]
     db_engine_kwargs: Optional[Dict[str, Any]]
     sys_db_engine_kwargs: Optional[Dict[str, Any]]
@@ -270,8 +264,7 @@ def load_config(
                 data["telemetry"]["OTLPExporter"]["tracesEndpoint"]
             ]
 
-    data = cast(ConfigFile, data)
-    return data  # type: ignore
+    return cast(ConfigFile, data)
 
 
 def process_config(
@@ -450,6 +443,7 @@ def configure_db_engine_parameters(
 
     # Configure user database engine parameters
     app_engine_kwargs: dict[str, Any] = {
+        "connect_args": {"application_name": "dbos_transact"},
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 20,
@@ -483,8 +477,6 @@ def is_valid_database_url(database_url: str) -> bool:
         return True
     url = make_url(database_url)
     required_fields = [
-        ("username", "Username must be specified in the connection URL"),
-        ("host", "Host must be specified in the connection URL"),
         ("database", "Database name must be specified in the connection URL"),
     ]
     for field_name, error_message in required_fields:

dbos/_debouncer.py

@@ -2,7 +2,6 @@ import asyncio
 import math
 import time
 import types
-import uuid
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -39,7 +38,7 @@ from dbos._error import DBOSQueueDeduplicatedError
 from dbos._queue import Queue
 from dbos._registrations import get_dbos_func_name
 from dbos._serialization import WorkflowInputs
-from dbos._utils import INTERNAL_QUEUE_NAME
+from dbos._utils import INTERNAL_QUEUE_NAME, generate_uuid
 
 if TYPE_CHECKING:
     from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
@@ -209,7 +208,7 @@ class Debouncer(Generic[P, R]):
 
             # Deterministically generate the user workflow ID and message ID
             def assign_debounce_ids() -> tuple[str, str]:
-                return str(uuid.uuid4()), ctx.assign_workflow_id()
+                return generate_uuid(), ctx.assign_workflow_id()
 
             message_id, user_workflow_id = dbos._sys_db.call_function_as_step(
                 assign_debounce_ids, "DBOS.assign_debounce_ids"
@@ -320,14 +319,14 @@ class DebouncerClient:
             "workflow_id": (
                 self.workflow_options["workflow_id"]
                 if self.workflow_options.get("workflow_id")
-                else str(uuid.uuid4())
+                else generate_uuid()
             ),
             "app_version": self.workflow_options.get("app_version"),
             "deduplication_id": self.workflow_options.get("deduplication_id"),
             "priority": self.workflow_options.get("priority"),
             "workflow_timeout_sec": self.workflow_options.get("workflow_timeout"),
         }
-        message_id = str(uuid.uuid4())
+        message_id = generate_uuid()
         while True:
             try:
                 # Attempt to enqueue a debouncer for this workflow.

dbos/_fastapi.py

@@ -1,4 +1,3 @@
-import uuid
 from typing import Any, Callable, MutableMapping, cast
 
 from fastapi import FastAPI
@@ -9,7 +8,7 @@ from starlette.types import ASGIApp, Receive, Scope, Send
 from . import DBOS
 from ._context import EnterDBOSHandler, OperationType, SetWorkflowID, TracedAttributes
 from ._error import DBOSException
-from ._utils import request_id_header
+from ._utils import generate_uuid, request_id_header
 
 
 def _get_or_generate_request_id(request: FastAPIRequest) -> str:
@@ -17,7 +16,7 @@ def _get_or_generate_request_id(request: FastAPIRequest) -> str:
     if request_id is not None:
         return request_id
     else:
-        return str(uuid.uuid4())
+        return generate_uuid()
 
 
 async def _dbos_error_handler(request: FastAPIRequest, gexc: Exception) -> JSONResponse:

dbos/_flask.py

@@ -1,4 +1,3 @@
-import uuid
 from typing import Any
 from urllib.parse import urlparse
 
@@ -6,7 +5,7 @@ from flask import Flask
 from werkzeug.wrappers import Request as WRequest
 
 from ._context import EnterDBOSHandler, OperationType, SetWorkflowID, TracedAttributes
-from ._utils import request_id_header
+from ._utils import generate_uuid, request_id_header
 
 
 class FlaskMiddleware:
@@ -41,7 +40,7 @@ def _get_or_generate_request_id(request: WRequest) -> str:
     if request_id is not None:
         return request_id
     else:
-        return str(uuid.uuid4())
+        return generate_uuid()
 
 
 def setup_flask_middleware(app: Flask) -> None:

dbos/_kafka.py

@@ -1,6 +1,6 @@
 import re
 import threading
-from typing import TYPE_CHECKING, Any, Callable, NoReturn
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, NoReturn
 
 from confluent_kafka import Consumer, KafkaError, KafkaException
 
@@ -15,7 +15,9 @@ from ._kafka_message import KafkaMessage
 from ._logger import dbos_logger
 from ._registrations import get_dbos_func_name
 
-_KafkaConsumerWorkflow = Callable[[KafkaMessage], None]
+_KafkaConsumerWorkflow = (
+    Callable[[KafkaMessage], None] | Callable[[KafkaMessage], Coroutine[Any, Any, None]]
+)
 
 _kafka_queue: Queue
 _in_order_kafka_queues: dict[str, Queue] = {}
@@ -37,8 +39,8 @@ def _kafka_consumer_loop(
     in_order: bool,
 ) -> None:
 
-    def on_error(err: KafkaError) -> NoReturn:
-        raise KafkaException(err)
+    def on_error(err: KafkaError) -> None:
+        dbos_logger.error(f"Exception in Kafka consumer: {err}")
 
     config["error_cb"] = on_error
     if "auto.offset.reset" not in config: