dbos 0.6.1__tar.gz → 0.7.0__tar.gz

{dbos-0.6.1 → dbos-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.6.1
+Version: 0.7.0
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT
@@ -8,7 +8,6 @@ Requires-Python: >=3.9
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: jsonschema>=4.23.0
 Requires-Dist: alembic>=1.13.2
-Requires-Dist: psycopg2-binary>=2.9.9
 Requires-Dist: typing-extensions>=4.12.2; python_version < "3.10"
 Requires-Dist: typer>=0.12.3
 Requires-Dist: jsonpickle>=3.2.2
@@ -19,6 +18,7 @@ Requires-Dist: python-dateutil>=2.9.0.post0
 Requires-Dist: fastapi[standard]>=0.112.1
 Requires-Dist: psutil>=6.0.0
 Requires-Dist: tomlkit>=0.13.2
+Requires-Dist: psycopg>=3.2.1
 Description-Content-Type: text/markdown
 
 

{dbos-0.6.1 → dbos-0.7.0}/dbos/__init__.py

@@ -3,6 +3,7 @@ from .context import DBOSContextEnsure, SetWorkflowID
 from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
 from .dbos_config import ConfigFile, get_dbos_database_url, load_config
 from .kafka_message import KafkaMessage
+from .queue import Queue
 from .system_database import GetWorkflowsInput, WorkflowStatusString
 
 __all__ = [
@@ -19,4 +20,5 @@ __all__ = [
     "load_config",
     "get_dbos_database_url",
     "error",
+    "Queue",
 ]

{dbos-0.6.1 → dbos-0.7.0}/dbos/application_database.py

@@ -2,7 +2,6 @@ from typing import Optional, TypedDict, cast
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
-import sqlalchemy.exc as sa_exc
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -36,7 +35,7 @@ class ApplicationDatabase:
 
         # If the application database does not already exist, create it
         postgres_db_url = sa.URL.create(
-            "postgresql",
+            "postgresql+psycopg",
             username=config["database"]["username"],
             password=config["database"]["password"],
             host=config["database"]["hostname"],
@@ -55,7 +54,7 @@ class ApplicationDatabase:
 
         # Create a connection pool for the application database
         app_db_url = sa.URL.create(
-            "postgresql",
+            "postgresql+psycopg",
             username=config["database"]["username"],
             password=config["database"]["password"],
             host=config["database"]["hostname"],
@@ -97,11 +96,9 @@ class ApplicationDatabase:
                 )
             )
         except DBAPIError as dbapi_error:
-            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
-            raise dbapi_error
-        except Exception as e:
-            raise e
+            raise
 
     def record_transaction_error(self, output: TransactionResultInternal) -> None:
         try:
@@ -122,11 +119,9 @@ class ApplicationDatabase:
                     )
                 )
         except DBAPIError as dbapi_error:
-            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
-            raise dbapi_error
-        except Exception as e:
-            raise e
+            raise
 
     @staticmethod
     def check_transaction_execution(

{dbos-0.6.1 → dbos-0.7.0}/dbos/context.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import json
 import os
 import uuid
+from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from enum import Enum
 from types import TracebackType
@@ -344,7 +345,7 @@ class SetWorkflowRecovery:
         return False  # Did not handle
 
 
-class EnterDBOSWorkflow:
+class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.created_ctx = False
         self.attributes = attributes
@@ -377,7 +378,7 @@ class EnterDBOSWorkflow:
         return False  # Did not handle
 
 
-class EnterDBOSChildWorkflow:
+class EnterDBOSChildWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.parent_ctx: Optional[DBOSContext] = None
         self.child_ctx: Optional[DBOSContext] = None

{dbos-0.6.1 → dbos-0.7.0}/dbos/core.py

@@ -63,6 +63,7 @@ from dbos.system_database import (
     OperationResultInternal,
     WorkflowInputs,
     WorkflowStatusInternal,
+    WorkflowStatusString,
 )
 
 if TYPE_CHECKING:
@@ -108,7 +109,7 @@ class _WorkflowHandlePolling(Generic[R]):
         return self.workflow_id
 
     def get_result(self) -> R:
-        res: R = self.dbos.sys_db.await_workflow_result(self.workflow_id)
+        res: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
         return res
 
     def get_status(self) -> "WorkflowStatus":
@@ -126,6 +127,7 @@ def _init_workflow(
     class_name: Optional[str],
     config_name: Optional[str],
     temp_wf_type: Optional[str],
+    queue: Optional[str] = None,
 ) -> WorkflowStatusInternal:
     wfid = (
         ctx.workflow_id
@@ -134,7 +136,11 @@ def _init_workflow(
     )
     status: WorkflowStatusInternal = {
         "workflow_uuid": wfid,
-        "status": "PENDING",
+        "status": (
+            WorkflowStatusString.PENDING.value
+            if queue is None
+            else WorkflowStatusString.ENQUEUED.value
+        ),
         "name": wf_name,
         "class_name": class_name,
         "config_name": config_name,
@@ -150,20 +156,25 @@ def _init_workflow(
             json.dumps(ctx.authenticated_roles) if ctx.authenticated_roles else None
         ),
         "assumed_role": ctx.assumed_role,
+        "queue_name": queue,
     }
 
     # If we have a class name, the first arg is the instance and do not serialize
     if class_name is not None:
         inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
 
-    if temp_wf_type != "transaction":
+    if temp_wf_type != "transaction" or queue is not None:
         # Synchronously record the status and inputs for workflows and single-step workflows
         # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
-        dbos.sys_db.update_workflow_status(status, False, ctx.in_recovery)
-        dbos.sys_db.update_workflow_inputs(wfid, utils.serialize(inputs))
+        # TODO: Make this transactional (and with the queue step below)
+        dbos._sys_db.update_workflow_status(status, False, ctx.in_recovery)
+        dbos._sys_db.update_workflow_inputs(wfid, utils.serialize(inputs))
     else:
         # Buffer the inputs for single-transaction workflows, but don't buffer the status
-        dbos.sys_db.buffer_workflow_inputs(wfid, utils.serialize(inputs))
+        dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize(inputs))
+
+    if queue is not None:
+        dbos._sys_db.enqueue(wfid, queue)
 
     return status
 
@@ -179,7 +190,9 @@ def _execute_workflow(
         output = func(*args, **kwargs)
         status["status"] = "SUCCESS"
         status["output"] = utils.serialize(output)
-        dbos.sys_db.buffer_workflow_status(status)
+        if status["queue_name"] is not None:
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
+        dbos._sys_db.buffer_workflow_status(status)
     except DBOSWorkflowConflictIDError:
         # Retrieve the workflow handle and wait for the result.
         # Must use existing_workflow=False because workflow status might not be set yet for single transaction workflows.
@@ -191,8 +204,10 @@ def _execute_workflow(
     except Exception as error:
         status["status"] = "ERROR"
         status["error"] = utils.serialize(error)
-        dbos.sys_db.update_workflow_status(status)
-        raise error
+        if status["queue_name"] is not None:
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
+        dbos._sys_db.update_workflow_status(status)
+        raise
 
     return output
 
@@ -217,14 +232,14 @@ def _execute_workflow_wthread(
                 dbos.logger.error(
                     f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
                 )
-                raise e
+                raise
 
 
 def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]":
-    status = dbos.sys_db.get_workflow_status(workflow_id)
+    status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
-    inputs = dbos.sys_db.get_workflow_inputs(workflow_id)
+    inputs = dbos._sys_db.get_workflow_inputs(workflow_id)
     if not inputs:
         raise DBOSRecoveryError(workflow_id, "Workflow inputs not found")
     wf_func = dbos._registry.workflow_info_map.get(status["name"], None)
@@ -249,6 +264,8 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
                 return _start_workflow(
                     dbos,
                     wf_func,
+                    status["queue_name"],
+                    True,
                     dbos._registry.instance_info_map[iname],
                     *inputs["args"],
                     **inputs["kwargs"],
@@ -264,6 +281,8 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
                 return _start_workflow(
                     dbos,
                     wf_func,
+                    status["queue_name"],
+                    True,
                     dbos._registry.class_info_map[class_name],
                     *inputs["args"],
                     **inputs["kwargs"],
@@ -271,7 +290,12 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
         else:
             with SetWorkflowID(workflow_id):
                 return _start_workflow(
-                    dbos, wf_func, *inputs["args"], **inputs["kwargs"]
+                    dbos,
+                    wf_func,
+                    status["queue_name"],
+                    True,
+                    *inputs["args"],
+                    **inputs["kwargs"],
                 )
 
 
@@ -298,34 +322,22 @@ def _workflow_wrapper(dbosreg: "_DBOSRegistry", func: F) -> F:
             "kwargs": kwargs,
         }
         ctx = get_local_dbos_context()
-        if ctx and ctx.is_workflow():
-            with EnterDBOSChildWorkflow(attributes), DBOSAssumeRole(rr):
-                ctx = assert_current_dbos_context()  # Now the child ctx
-                status = _init_workflow(
-                    dbos,
-                    ctx,
-                    inputs=inputs,
-                    wf_name=get_dbos_func_name(func),
-                    class_name=get_dbos_class_name(fi, func, args),
-                    config_name=get_config_name(fi, func, args),
-                    temp_wf_type=get_temp_workflow_type(func),
-                )
-
-                return _execute_workflow(dbos, status, func, *args, **kwargs)
-        else:
-            with EnterDBOSWorkflow(attributes), DBOSAssumeRole(rr):
-                ctx = assert_current_dbos_context()
-                status = _init_workflow(
-                    dbos,
-                    ctx,
-                    inputs=inputs,
-                    wf_name=get_dbos_func_name(func),
-                    class_name=get_dbos_class_name(fi, func, args),
-                    config_name=get_config_name(fi, func, args),
-                    temp_wf_type=get_temp_workflow_type(func),
-                )
+        enterWorkflowCtxMgr = (
+            EnterDBOSChildWorkflow if ctx and ctx.is_workflow() else EnterDBOSWorkflow
+        )
+        with enterWorkflowCtxMgr(attributes), DBOSAssumeRole(rr):
+            ctx = assert_current_dbos_context()  # Now the child ctx
+            status = _init_workflow(
+                dbos,
+                ctx,
+                inputs=inputs,
+                wf_name=get_dbos_func_name(func),
+                class_name=get_dbos_class_name(fi, func, args),
+                config_name=get_config_name(fi, func, args),
+                temp_wf_type=get_temp_workflow_type(func),
+            )
 
-                return _execute_workflow(dbos, status, func, *args, **kwargs)
+            return _execute_workflow(dbos, status, func, *args, **kwargs)
 
     wrapped_func = cast(F, wrapper)
     return wrapped_func
@@ -343,6 +355,8 @@ def _workflow(reg: "_DBOSRegistry") -> Callable[[F], F]:
 def _start_workflow(
     dbos: "DBOS",
     func: "Workflow[P, R]",
+    queue_name: Optional[str],
+    execute_workflow: bool,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
@@ -396,10 +410,14 @@ def _start_workflow(
         class_name=get_dbos_class_name(fi, func, gin_args),
         config_name=get_config_name(fi, func, gin_args),
         temp_wf_type=get_temp_workflow_type(func),
+        queue=queue_name,
     )
 
+    if not execute_workflow:
+        return _WorkflowHandlePolling(new_wf_id, dbos)
+
     if fself is not None:
-        future = dbos.executor.submit(
+        future = dbos._executor.submit(
             cast(Callable[..., R], _execute_workflow_wthread),
             dbos,
             status,
@@ -410,7 +428,7 @@ def _start_workflow(
             **kwargs,
         )
     else:
-        future = dbos.executor.submit(
+        future = dbos._executor.submit(
             cast(Callable[..., R], _execute_workflow_wthread),
             dbos,
             status,
@@ -432,7 +450,7 @@ def _transaction(
                     f"Function {func.__name__} invoked before DBOS initialized"
                 )
             dbos = dbosreg.dbos
-            with dbos.app_db.sessionmaker() as session:
+            with dbos._app_db.sessionmaker() as session:
                 attributes: TracedAttributes = {
                     "name": func.__name__,
                     "operationType": OperationType.TRANSACTION.value,
@@ -493,7 +511,7 @@ def _transaction(
                                 )
                                 break
                         except DBAPIError as dbapi_error:
-                            if dbapi_error.orig.pgcode == "40001":  # type: ignore
+                            if dbapi_error.orig.sqlstate == "40001":  # type: ignore
                                 # Retry on serialization failure
                                 ctx.get_current_span().add_event(
                                     "Transaction Serialization Failure",
@@ -505,13 +523,13 @@ def _transaction(
                                     max_retry_wait_seconds,
                                 )
                                 continue
-                            raise dbapi_error
+                            raise
                         except Exception as error:
                             # Don't record the error if it was already recorded
                             if not has_recorded_error:
                                 txn_output["error"] = utils.serialize(error)
-                                dbos.app_db.record_transaction_error(txn_output)
-                            raise error
+                                dbos._app_db.record_transaction_error(txn_output)
+                            raise
             return output
 
         fi = get_or_create_func_info(func)
@@ -541,6 +559,7 @@ def _transaction(
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "transaction")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        wrapper.__orig_func = temp_wf  # type: ignore
 
         return cast(F, wrapper)
 
@@ -575,7 +594,7 @@ def _step(
                     "output": None,
                     "error": None,
                 }
-                recorded_output = dbos.sys_db.check_operation_execution(
+                recorded_output = dbos._sys_db.check_operation_execution(
                     ctx.workflow_id, ctx.function_id
                 )
                 if recorded_output:
@@ -622,7 +641,7 @@ def _step(
                 step_output["error"] = (
                     utils.serialize(error) if error is not None else None
                 )
-                dbos.sys_db.record_operation_result(step_output)
+                dbos._sys_db.record_operation_result(step_output)
 
                 if error is not None:
                     raise error
@@ -657,6 +676,7 @@ def _step(
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "step")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        wrapper.__orig_func = temp_wf  # type: ignore
 
         return cast(F, wrapper)
 
@@ -671,7 +691,7 @@ def _send(
             "name": "send",
         }
         with EnterDBOSStep(attributes) as ctx:
-            dbos.sys_db.send(
+            dbos._sys_db.send(
                 ctx.workflow_id,
                 ctx.curr_step_function_id,
                 destination_id,
@@ -702,7 +722,7 @@ def _recv(
         with EnterDBOSStep(attributes) as ctx:
             ctx.function_id += 1  # Reserve for the sleep
             timeout_function_id = ctx.function_id
-            return dbos.sys_db.recv(
+            return dbos._sys_db.recv(
                 ctx.workflow_id,
                 ctx.curr_step_function_id,
                 timeout_function_id,
@@ -725,7 +745,7 @@ def _set_event(dbos: "DBOS", key: str, value: Any) -> None:
             "name": "set_event",
         }
         with EnterDBOSStep(attributes) as ctx:
-            dbos.sys_db.set_event(
+            dbos._sys_db.set_event(
                 ctx.workflow_id, ctx.curr_step_function_id, key, value
             )
     else:
@@ -753,7 +773,7 @@ def _get_event(
                 "function_id": ctx.curr_step_function_id,
                 "timeout_function_id": timeout_function_id,
             }
-            return dbos.sys_db.get_event(workflow_id, key, timeout_seconds, caller_ctx)
+            return dbos._sys_db.get_event(workflow_id, key, timeout_seconds, caller_ctx)
     else:
         # Directly call it outside of a workflow
-        return dbos.sys_db.get_event(workflow_id, key, timeout_seconds)
+        return dbos._sys_db.get_event(workflow_id, key, timeout_seconds)