dbos 0.25.0a14__tar.gz → 0.25.1__tar.gz

{dbos-0.25.0a14 → dbos-0.25.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.25.0a14
+Version: 0.25.1
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/__init__.py

@@ -1,4 +1,5 @@
 from . import _error as error
+from ._client import DBOSClient, EnqueueOptions
 from ._context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
 from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle
 from ._dbos_config import ConfigFile, DBOSConfig, get_dbos_database_url, load_config
@@ -11,9 +12,11 @@ __all__ = [
     "ConfigFile",
     "DBOSConfig",
     "DBOS",
+    "DBOSClient",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
     "DBOSContextSetAuth",
+    "EnqueueOptions",
     "GetWorkflowsInput",
     "KafkaMessage",
     "SetWorkflowID",

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_app_db.py

@@ -27,19 +27,18 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase:
 
-    def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
-        self.config = config
+    def __init__(self, database: DatabaseConfig, *, debug_mode: bool = False):
 
-        app_db_name = config["database"]["app_db_name"]
+        app_db_name = database["app_db_name"]
 
         # If the application database does not already exist, create it
         if not debug_mode:
             postgres_db_url = sa.URL.create(
                 "postgresql+psycopg",
-                username=config["database"]["username"],
-                password=config["database"]["password"],
-                host=config["database"]["hostname"],
-                port=config["database"]["port"],
+                username=database["username"],
+                password=database["password"],
+                host=database["hostname"],
+                port=database["port"],
                 database="postgres",
             )
             postgres_db_engine = sa.create_engine(postgres_db_url)
@@ -55,25 +54,25 @@ class ApplicationDatabase:
         # Create a connection pool for the application database
         app_db_url = sa.URL.create(
             "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
+            username=database["username"],
+            password=database["password"],
+            host=database["hostname"],
+            port=database["port"],
             database=app_db_name,
         )
 
         connect_args = {}
         if (
-            "connectionTimeoutMillis" in config["database"]
-            and config["database"]["connectionTimeoutMillis"]
+            "connectionTimeoutMillis" in database
+            and database["connectionTimeoutMillis"]
         ):
             connect_args["connect_timeout"] = int(
-                config["database"]["connectionTimeoutMillis"] / 1000
+                database["connectionTimeoutMillis"] / 1000
             )
 
         self.engine = sa.create_engine(
             app_db_url,
-            pool_size=config["database"]["app_db_pool_size"],
+            pool_size=database["app_db_pool_size"],
             max_overflow=0,
             pool_timeout=30,
             connect_args=connect_args,

dbos-0.25.1/dbos/_client.py

@@ -0,0 +1,206 @@
+import asyncio
+import sys
+import uuid
+from typing import Any, Generic, Optional, TypedDict, TypeVar
+
+if sys.version_info < (3, 11):
+    from typing_extensions import NotRequired
+else:
+    from typing import NotRequired
+
+from dbos import _serialization
+from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
+from dbos._dbos_config import parse_database_url_to_dbconfig
+from dbos._error import DBOSNonExistentWorkflowError
+from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
+from dbos._serialization import WorkflowInputs
+from dbos._sys_db import SystemDatabase, WorkflowStatusInternal, WorkflowStatusString
+from dbos._workflow_commands import WorkflowStatus, get_workflow
+
+R = TypeVar("R", covariant=True)  # A generic type for workflow return values
+
+
+class EnqueueOptions(TypedDict):
+    workflow_name: str
+    workflow_class_name: NotRequired[str]
+    queue_name: str
+    app_version: NotRequired[str]
+    workflow_id: NotRequired[str]
+
+
+class WorkflowHandleClientPolling(Generic[R]):
+
+    def __init__(self, workflow_id: str, sys_db: SystemDatabase):
+        self.workflow_id = workflow_id
+        self._sys_db = sys_db
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    def get_result(self) -> R:
+        res: R = self._sys_db.await_workflow_result(self.workflow_id)
+        return res
+
+    def get_status(self) -> "WorkflowStatus":
+        status = get_workflow(self._sys_db, self.workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(self.workflow_id)
+        return status
+
+
+class WorkflowHandleClientAsyncPolling(Generic[R]):
+
+    def __init__(self, workflow_id: str, sys_db: SystemDatabase):
+        self.workflow_id = workflow_id
+        self._sys_db = sys_db
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    async def get_result(self) -> R:
+        res: R = await asyncio.to_thread(
+            self._sys_db.await_workflow_result, self.workflow_id
+        )
+        return res
+
+    async def get_status(self) -> "WorkflowStatus":
+        status = await asyncio.to_thread(
+            get_workflow, self._sys_db, self.workflow_id, True
+        )
+        if status is None:
+            raise DBOSNonExistentWorkflowError(self.workflow_id)
+        return status
+
+
+class DBOSClient:
+    def __init__(self, database_url: str, *, system_database: Optional[str] = None):
+        db_config = parse_database_url_to_dbconfig(database_url)
+        if system_database is not None:
+            db_config["sys_db_name"] = system_database
+        self._sys_db = SystemDatabase(db_config)
+
+    def destroy(self) -> None:
+        self._sys_db.destroy()
+
+    def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
+        workflow_name = options["workflow_name"]
+        queue_name = options["queue_name"]
+
+        workflow_class_name = options.get("workflow_class_name")
+        app_version = options.get("app_version")
+        max_recovery_attempts = options.get("max_recovery_attempts")
+        if max_recovery_attempts is None:
+            max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
+        workflow_id = options.get("workflow_id")
+        if workflow_id is None:
+            workflow_id = str(uuid.uuid4())
+
+        status: WorkflowStatusInternal = {
+            "workflow_uuid": workflow_id,
+            "status": WorkflowStatusString.ENQUEUED.value,
+            "name": workflow_name,
+            "class_name": workflow_class_name,
+            "queue_name": queue_name,
+            "app_version": app_version,
+            "config_name": None,
+            "authenticated_user": None,
+            "assumed_role": None,
+            "authenticated_roles": None,
+            "request": None,
+            "output": None,
+            "error": None,
+            "created_at": None,
+            "updated_at": None,
+            "executor_id": None,
+            "recovery_attempts": None,
+            "app_id": None,
+        }
+
+        inputs: WorkflowInputs = {
+            "args": args,
+            "kwargs": kwargs,
+        }
+
+        wf_status = self._sys_db.insert_workflow_status(status)
+        self._sys_db.update_workflow_inputs(
+            workflow_id, _serialization.serialize_args(inputs)
+        )
+        if wf_status == WorkflowStatusString.ENQUEUED.value:
+            self._sys_db.enqueue(workflow_id, queue_name)
+        return workflow_id
+
+    def enqueue(
+        self, options: EnqueueOptions, *args: Any, **kwargs: Any
+    ) -> WorkflowHandle[R]:
+        workflow_id = self._enqueue(options, *args, **kwargs)
+        return WorkflowHandleClientPolling[R](workflow_id, self._sys_db)
+
+    async def enqueue_async(
+        self, options: EnqueueOptions, *args: Any, **kwargs: Any
+    ) -> WorkflowHandleAsync[R]:
+        workflow_id = await asyncio.to_thread(self._enqueue, options, *args, **kwargs)
+        return WorkflowHandleClientAsyncPolling[R](workflow_id, self._sys_db)
+
+    def retrieve_workflow(self, workflow_id: str) -> WorkflowHandle[R]:
+        status = get_workflow(self._sys_db, workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(workflow_id)
+        return WorkflowHandleClientPolling[R](workflow_id, self._sys_db)
+
+    async def retrieve_workflow_async(self, workflow_id: str) -> WorkflowHandleAsync[R]:
+        status = asyncio.to_thread(get_workflow, self._sys_db, workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(workflow_id)
+        return WorkflowHandleClientAsyncPolling[R](workflow_id, self._sys_db)
+
+    def send(
+        self,
+        destination_id: str,
+        message: Any,
+        topic: Optional[str] = None,
+        idempotency_key: Optional[str] = None,
+    ) -> None:
+        idempotency_key = idempotency_key if idempotency_key else str(uuid.uuid4())
+        status: WorkflowStatusInternal = {
+            "workflow_uuid": f"{destination_id}-{idempotency_key}",
+            "status": WorkflowStatusString.SUCCESS.value,
+            "name": "temp_workflow-send-client",
+            "class_name": None,
+            "queue_name": None,
+            "config_name": None,
+            "authenticated_user": None,
+            "assumed_role": None,
+            "authenticated_roles": None,
+            "request": None,
+            "output": None,
+            "error": None,
+            "created_at": None,
+            "updated_at": None,
+            "executor_id": None,
+            "recovery_attempts": None,
+            "app_id": None,
+            "app_version": None,
+        }
+        self._sys_db.insert_workflow_status(status)
+        self._sys_db.send(status["workflow_uuid"], 0, destination_id, message, topic)
+
+    async def send_async(
+        self,
+        destination_id: str,
+        message: Any,
+        topic: Optional[str] = None,
+        idempotency_key: Optional[str] = None,
+    ) -> None:
+        return await asyncio.to_thread(
+            self.send, destination_id, message, topic, idempotency_key
+        )
+
+    def get_event(self, workflow_id: str, key: str, timeout_seconds: float = 60) -> Any:
+        return self._sys_db.get_event(workflow_id, key, timeout_seconds)
+
+    async def get_event_async(
+        self, workflow_id: str, key: str, timeout_seconds: float = 60
+    ) -> Any:
+        return await asyncio.to_thread(
+            self.get_event, workflow_id, key, timeout_seconds
+        )

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_context.py

@@ -195,8 +195,10 @@ class DBOSContext:
     def end_handler(self, exc_value: Optional[BaseException]) -> None:
         self._end_span(exc_value)
 
-    def get_current_span(self) -> Span:
-        return self.spans[-1]
+    def get_current_span(self) -> Optional[Span]:
+        if len(self.spans):
+            return self.spans[-1]
+        return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
         attributes["operationUUID"] = (

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_core.py

@@ -892,10 +892,12 @@ def decorate_transaction(
                         except DBAPIError as dbapi_error:
                             if dbapi_error.orig.sqlstate == "40001":  # type: ignore
                                 # Retry on serialization failure
-                                ctx.get_current_span().add_event(
-                                    "Transaction Serialization Failure",
-                                    {"retry_wait_seconds": retry_wait_seconds},
-                                )
+                                span = ctx.get_current_span()
+                                if span:
+                                    span.add_event(
+                                        "Transaction Serialization Failure",
+                                        {"retry_wait_seconds": retry_wait_seconds},
+                                    )
                                 time.sleep(retry_wait_seconds)
                                 retry_wait_seconds = min(
                                     retry_wait_seconds * backoff_factor,
@@ -1004,13 +1006,15 @@ def decorate_step(
                     f"Step being automatically retried. (attempt {attempt + 1} of {attempts}). {traceback.format_exc()}"
                 )
                 ctx = assert_current_dbos_context()
-                ctx.get_current_span().add_event(
-                    f"Step attempt {attempt} failed",
-                    {
-                        "error": str(error),
-                        "retryIntervalSeconds": interval_seconds,
-                    },
-                )
+                span = ctx.get_current_span()
+                if span:
+                    span.add_event(
+                        f"Step attempt {attempt} failed",
+                        {
+                            "error": str(error),
+                            "retryIntervalSeconds": interval_seconds,
+                        },
+                    )
                 return min(
                     interval_seconds * (backoff_rate**attempt),
                     max_retry_interval_seconds,

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_dbos.py

@@ -445,8 +445,12 @@ class DBOS:
             dbos_logger.info(f"Executor ID: {GlobalParams.executor_id}")
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=64)
-            self._sys_db_field = SystemDatabase(self.config, debug_mode=debug_mode)
-            self._app_db_field = ApplicationDatabase(self.config, debug_mode=debug_mode)
+            self._sys_db_field = SystemDatabase(
+                self.config["database"], debug_mode=debug_mode
+            )
+            self._app_db_field = ApplicationDatabase(
+                self.config["database"], debug_mode=debug_mode
+            )
 
             if debug_mode:
                 return
@@ -1079,7 +1083,9 @@ class DBOS:
     def span(cls) -> Span:
         """Return the tracing `Span` associated with the current context."""
         ctx = assert_current_dbos_context()
-        return ctx.get_current_span()
+        span = ctx.get_current_span()
+        assert span
+        return span
 
     @classproperty
     def request(cls) -> Optional["Request"]:

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_logger.py

@@ -7,6 +7,7 @@ from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
 from opentelemetry.sdk.resources import Resource
+from opentelemetry.trace.span import format_trace_id
 
 from dbos._utils import GlobalParams
 
@@ -26,6 +27,19 @@ class DBOSLogTransformer(logging.Filter):
         record.applicationID = self.app_id
         record.applicationVersion = GlobalParams.app_version
         record.executorID = GlobalParams.executor_id
+
+        # If available, decorate the log entry with Workflow ID and Trace ID
+        from dbos._context import get_local_dbos_context
+
+        ctx = get_local_dbos_context()
+        if ctx:
+            if ctx.is_within_workflow():
+                record.operationUUID = ctx.workflow_id
+            span = ctx.get_current_span()
+            if span:
+                trace_id = format_trace_id(span.get_span_context().trace_id)
+                record.traceId = trace_id
+
         return True
 
 

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_queue.py

@@ -85,7 +85,7 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
         for _, queue in dbos._registry.queue_info_map.items():
             try:
                 wf_ids = dbos._sys_db.start_queued_workflows(
-                    queue, GlobalParams.executor_id
+                    queue, GlobalParams.executor_id, GlobalParams.app_version
                 )
                 for id in wf_ids:
                     execute_workflow_by_id(dbos, id)

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/_sys_db.py

@@ -29,7 +29,7 @@ from dbos._utils import GlobalParams
 
 from . import _serialization
 from ._context import get_local_dbos_context
-from ._dbos_config import ConfigFile
+from ._dbos_config import ConfigFile, DatabaseConfig
 from ._error import (
     DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
@@ -170,23 +170,21 @@ _dbos_null_topic = "__null__topic__"
 
 class SystemDatabase:
 
-    def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
-        self.config = config
-
+    def __init__(self, database: DatabaseConfig, *, debug_mode: bool = False):
         sysdb_name = (
-            config["database"]["sys_db_name"]
-            if "sys_db_name" in config["database"] and config["database"]["sys_db_name"]
-            else config["database"]["app_db_name"] + SystemSchema.sysdb_suffix
+            database["sys_db_name"]
+            if "sys_db_name" in database and database["sys_db_name"]
+            else database["app_db_name"] + SystemSchema.sysdb_suffix
         )
 
         if not debug_mode:
             # If the system database does not already exist, create it
             postgres_db_url = sa.URL.create(
                 "postgresql+psycopg",
-                username=config["database"]["username"],
-                password=config["database"]["password"],
-                host=config["database"]["hostname"],
-                port=config["database"]["port"],
+                username=database["username"],
+                password=database["password"],
+                host=database["hostname"],
+                port=database["port"],
                 database="postgres",
                 # fills the "application_name" column in pg_stat_activity
                 query={"application_name": f"dbos_transact_{GlobalParams.executor_id}"},
@@ -203,19 +201,23 @@ class SystemDatabase:
 
         system_db_url = sa.URL.create(
             "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
+            username=database["username"],
+            password=database["password"],
+            host=database["hostname"],
+            port=database["port"],
             database=sysdb_name,
             # fills the "application_name" column in pg_stat_activity
             query={"application_name": f"dbos_transact_{GlobalParams.executor_id}"},
         )
 
         # Create a connection pool for the system database
+        pool_size = database.get("sys_db_pool_size")
+        if pool_size is None:
+            pool_size = 20
+
         self.engine = sa.create_engine(
             system_db_url,
-            pool_size=config["database"]["sys_db_pool_size"],
+            pool_size=pool_size,
             max_overflow=0,
             pool_timeout=30,
             connect_args={"connect_timeout": 10},
@@ -1264,7 +1266,9 @@ class SystemDatabase:
                 .on_conflict_do_nothing()
             )
 
-    def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
+    def start_queued_workflows(
+        self, queue: "Queue", executor_id: str, app_version: str
+    ) -> List[str]:
         if self._debug_mode:
             return []
 
@@ -1379,26 +1383,36 @@ class SystemDatabase:
                         break
 
                 # To start a function, first set its status to PENDING and update its executor ID
-                c.execute(
+                res = c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
                     .where(
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.ENQUEUED.value
                     )
+                    .where(
+                        sa.or_(
+                            SystemSchema.workflow_status.c.application_version
+                            == app_version,
+                            SystemSchema.workflow_status.c.application_version.is_(
+                                None
+                            ),
+                        )
+                    )
                     .values(
                         status=WorkflowStatusString.PENDING.value,
+                        application_version=app_version,
                         executor_id=executor_id,
                     )
                 )
-
-                # Then give it a start time and assign the executor ID
-                c.execute(
-                    SystemSchema.workflow_queue.update()
-                    .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
-                    .values(started_at_epoch_ms=start_time_ms)
-                )
-                ret_ids.append(id)
+                if res.rowcount > 0:
+                    # Then give it a start time and assign the executor ID
+                    c.execute(
+                        SystemSchema.workflow_queue.update()
+                        .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
+                        .values(started_at_epoch_ms=start_time_ms)
+                    )
+                    ret_ids.append(id)
 
             # If we have a limiter, garbage-collect all completed functions started
             # before the period. If there's no limiter, there's no need--they were

{dbos-0.25.0a14 → dbos-0.25.1}/dbos/cli/cli.py

@@ -189,8 +189,8 @@ def migrate() -> None:
     app_db = None
     sys_db = None
     try:
-        sys_db = SystemDatabase(config)
-        app_db = ApplicationDatabase(config)
+        sys_db = SystemDatabase(config["database"])
+        app_db = ApplicationDatabase(config["database"])
     except Exception as e:
         typer.echo(f"DBOS system schema migration failed: {e}")
     finally:
@@ -314,7 +314,7 @@ def list(
     ] = False,
 ) -> None:
     config = load_config(silent=True)
-    sys_db = SystemDatabase(config)
+    sys_db = SystemDatabase(config["database"])
     workflows = list_workflows(
         sys_db,
         limit=limit,
@@ -338,7 +338,7 @@ def get(
     ] = False,
 ) -> None:
     config = load_config(silent=True)
-    sys_db = SystemDatabase(config)
+    sys_db = SystemDatabase(config["database"])
     print(
         jsonpickle.encode(get_workflow(sys_db, workflow_id, request), unpicklable=False)
     )
@@ -349,7 +349,7 @@ def steps(
     workflow_id: Annotated[str, typer.Argument()],
 ) -> None:
     config = load_config(silent=True)
-    sys_db = SystemDatabase(config)
+    sys_db = SystemDatabase(config["database"])
     print(
         jsonpickle.encode(list_workflow_steps(sys_db, workflow_id), unpicklable=False)
     )
@@ -475,7 +475,7 @@ def list_queue(
     ] = False,
 ) -> None:
     config = load_config(silent=True)
-    sys_db = SystemDatabase(config)
+    sys_db = SystemDatabase(config["database"])
     workflows = list_queued_workflows(
         sys_db=sys_db,
         limit=limit,

{dbos-0.25.0a14 → dbos-0.25.1}/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.25.0a14"
+version = "0.25.1"
 
 [project.license]
 text = "MIT"

dbos-0.25.1/tests/client_collateral.py

@@ -0,0 +1,38 @@
+import json
+from typing import Optional, TypedDict, cast
+
+from dbos import DBOS, Queue
+
+
+class Person(TypedDict):
+    first: str
+    last: str
+    age: int
+
+
+queue = Queue("test_queue")
+
+
+@DBOS.workflow()
+def enqueue_test(numVal: int, strVal: str, person: Person) -> str:
+    return f"{numVal}-{strVal}-{json.dumps(person)}"
+
+
+@DBOS.workflow()
+def send_test(topic: Optional[str] = None) -> str:
+    return cast(str, DBOS.recv(topic, 60))
+
+
+@DBOS.workflow()
+def retrieve_test(value: str) -> str:
+    DBOS.sleep(5)
+    return value
+
+
+@DBOS.workflow()
+def event_test(key: str, value: str, update: Optional[int] = None) -> str:
+    DBOS.set_event(key, value)
+    if update is not None:
+        DBOS.sleep(update)
+        DBOS.set_event(key, f"updated-{value}")
+    return f"{key}-{value}"

dbos-0.25.1/tests/client_worker.py

@@ -0,0 +1,30 @@
+import os
+import sys
+
+from dbos import DBOS, SetWorkflowID
+from tests import client_collateral as cc
+from tests.conftest import default_config
+
+app_vers = os.environ.get("DBOS__APPVERSION")
+if app_vers is None:
+    DBOS.logger.error("DBOS__APPVERSION not set")
+    os._exit(1)
+else:
+    DBOS.logger.info(f"DBOS__APPVERSION: {app_vers}")
+
+if len(sys.argv) < 2:
+    DBOS.logger.error("Usage: client_worker wfid <topic>")
+    os._exit(1)
+
+wfid = sys.argv[1]
+topic = sys.argv[2] if len(sys.argv) > 2 else None
+
+config = default_config()
+DBOS(config=config)
+DBOS.launch()
+
+DBOS.logger.info(f"Starting send_test with WF ID: {wfid}")
+with SetWorkflowID(wfid):
+    DBOS.start_workflow(cc.send_test, topic)
+
+os._exit(0)