DBOS Client (#265)

DBOS Client, but this time in Python!

Author: devhawk

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 @@
     "ConfigFile",
     "DBOSConfig",
     "DBOS",
+    "DBOSClient",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
     "DBOSContextSetAuth",
+    "EnqueueOptions",
     "GetWorkflowsInput",
     "KafkaMessage",
     "SetWorkflowID",

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 @@ def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
         # 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/_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/_dbos.py

@@ -445,8 +445,12 @@ def _launch(self, *, debug_mode: bool = False) -> None:
             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

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)