dbos 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

dbos/error.py

@@ -32,7 +32,6 @@ class DBOSErrorCode(Enum):
     InitializationError = 3
     WorkflowFunctionNotFound = 4
     NonExistentWorkflowError = 5
-    DuplicateWorkflowEventError = 6
     MaxStepRetriesExceeded = 7
     NotAuthorized = 8
 
@@ -87,16 +86,6 @@ class DBOSNonExistentWorkflowError(DBOSException):
         )
 
 
-class DBOSDuplicateWorkflowEventError(DBOSException):
-    """Exception raised when a workflow attempts to set an event value more than once per key."""
-
-    def __init__(self, workflow_id: str, key: str):
-        super().__init__(
-            f"Workflow {workflow_id} has already emitted an event with key {key}",
-            dbos_error_code=DBOSErrorCode.DuplicateWorkflowEventError.value,
-        )
-
-
 class DBOSNotAuthorizedError(DBOSException):
     """Exception raised by DBOS role-based security when the user is not authorized to access a function."""
 

dbos/fastapi.py

@@ -1,8 +1,13 @@
 import uuid
-from typing import Any, Callable
+from typing import Any, Callable, cast
 
 from fastapi import FastAPI
 from fastapi import Request as FastAPIRequest
+from fastapi.responses import JSONResponse
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from dbos import DBOS
+from dbos.error import DBOSException
 
 from .context import (
     EnterDBOSHandler,
@@ -35,7 +40,46 @@ def make_request(request: FastAPIRequest) -> Request:
     )
 
 
-def setup_fastapi_middleware(app: FastAPI) -> None:
+async def dbos_error_handler(request: FastAPIRequest, gexc: Exception) -> JSONResponse:
+    exc: DBOSException = cast(DBOSException, gexc)
+    status_code = 500
+    if exc.status_code is not None:
+        status_code = exc.status_code
+    return JSONResponse(
+        status_code=status_code,
+        content={
+            "message": str(exc.message),
+            "dbos_error_code": str(exc.dbos_error_code),
+            "dbos_error": str(exc.__class__.__name__),
+        },
+    )
+
+
+class LifespanMiddleware:
+    def __init__(self, app: ASGIApp, dbos: DBOS):
+        self.app = app
+        self.dbos = dbos
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "lifespan":
+            while True:
+                message = await receive()
+                if message["type"] == "lifespan.startup":
+                    self.dbos._launch()
+                    await send({"type": "lifespan.startup.complete"})
+                elif message["type"] == "lifespan.shutdown":
+                    self.dbos._destroy()
+                    await send({"type": "lifespan.shutdown.complete"})
+                    break
+        else:
+            await self.app(scope, receive, send)
+
+
+def setup_fastapi_middleware(app: FastAPI, dbos: DBOS) -> None:
+
+    app.add_middleware(LifespanMiddleware, dbos=dbos)
+    app.add_exception_handler(DBOSException, dbos_error_handler)
+
     @app.middleware("http")
     async def dbos_fastapi_middleware(
         request: FastAPIRequest, call_next: Callable[..., Any]

dbos/kafka.py

@@ -1,26 +1,30 @@
 import threading
-import traceback
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generator, NoReturn, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, NoReturn
 
 from confluent_kafka import Consumer, KafkaError, KafkaException
-from confluent_kafka import Message as CTypeMessage
+
+from dbos.queue import Queue
 
 if TYPE_CHECKING:
     from dbos.dbos import _DBOSRegistry
 
 from .context import SetWorkflowID
+from .error import DBOSInitializationError
 from .kafka_message import KafkaMessage
 from .logger import dbos_logger
 
 KafkaConsumerWorkflow = Callable[[KafkaMessage], None]
 
+kafka_queue: Queue
+in_order_kafka_queues: dict[str, Queue] = {}
+
 
 def _kafka_consumer_loop(
     func: KafkaConsumerWorkflow,
     config: dict[str, Any],
     topics: list[str],
     stop_event: threading.Event,
+    in_order: bool,
 ) -> None:
 
     def on_error(err: KafkaError) -> NoReturn:
@@ -70,24 +74,35 @@ def _kafka_consumer_loop(
                 with SetWorkflowID(
                     f"kafka-unique-id-{msg.topic}-{msg.partition}-{msg.offset}"
                 ):
-                    try:
-                        func(msg)
-                    except Exception as e:
-                        dbos_logger.error(
-                            f"Exception encountered in Kafka consumer: {traceback.format_exc()}"
-                        )
+                    if in_order:
+                        assert msg.topic is not None
+                        queue = in_order_kafka_queues[msg.topic]
+                        queue.enqueue(func, msg)
+                    else:
+                        kafka_queue.enqueue(func, msg)
 
     finally:
         consumer.close()
 
 
 def kafka_consumer(
-    dbosreg: "_DBOSRegistry", config: dict[str, Any], topics: list[str]
+    dbosreg: "_DBOSRegistry", config: dict[str, Any], topics: list[str], in_order: bool
 ) -> Callable[[KafkaConsumerWorkflow], KafkaConsumerWorkflow]:
     def decorator(func: KafkaConsumerWorkflow) -> KafkaConsumerWorkflow:
+        if in_order:
+            for topic in topics:
+                if topic.startswith("^"):
+                    raise DBOSInitializationError(
+                        f"Error: in-order processing is not supported for regular expression topic selectors ({topic})"
+                    )
+                queue = Queue(f"_dbos_kafka_queue_topic_{topic}", concurrency=1)
+                in_order_kafka_queues[topic] = queue
+        else:
+            global kafka_queue
+            kafka_queue = Queue("_dbos_internal_queue")
         stop_event = threading.Event()
         dbosreg.register_poller(
-            stop_event, _kafka_consumer_loop, func, config, topics, stop_event
+            stop_event, _kafka_consumer_loop, func, config, topics, stop_event, in_order
         )
         return func
 

dbos/migrations/versions/eab0cc1d9a14_job_queue.py

@@ -0,0 +1,55 @@
+"""job_queue
+
+Revision ID: eab0cc1d9a14
+Revises: a3b18ad34abe
+Create Date: 2024-09-13 14:50:00.531294
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "eab0cc1d9a14"
+down_revision: Union[str, None] = "a3b18ad34abe"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "job_queue",
+        sa.Column("workflow_uuid", sa.Text(), nullable=False),
+        sa.Column("queue_name", sa.Text(), nullable=False),
+        sa.Column(
+            "created_at_epoch_ms",
+            sa.BigInteger(),
+            server_default=sa.text(
+                "(EXTRACT(epoch FROM now()) * 1000::numeric)::bigint"
+            ),
+            nullable=False,
+            primary_key=True,
+        ),
+        sa.ForeignKeyConstraint(
+            ["workflow_uuid"],
+            ["dbos.workflow_status.workflow_uuid"],
+            onupdate="CASCADE",
+            ondelete="CASCADE",
+        ),
+        schema="dbos",
+    )
+    op.add_column(
+        "workflow_status",
+        sa.Column(
+            "queue_name",
+            sa.Text(),
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("job_queue", schema="dbos")
+    op.drop_column("workflow_status", "queue_name", schema="dbos")

dbos/queue.py

@@ -0,0 +1,36 @@
+import threading
+import time
+from typing import TYPE_CHECKING, Optional
+
+from dbos.core import P, R, _execute_workflow_id, _start_workflow
+from dbos.error import DBOSInitializationError
+
+if TYPE_CHECKING:
+    from dbos.dbos import DBOS, Workflow, WorkflowHandle
+
+
+class Queue:
+    def __init__(self, name: str, concurrency: Optional[int] = None) -> None:
+        self.name = name
+        self.concurrency = concurrency
+        from dbos.dbos import _get_or_create_dbos_registry
+
+        registry = _get_or_create_dbos_registry()
+        registry.queue_info_map[self.name] = self
+
+    def enqueue(
+        self, func: "Workflow[P, R]", *args: P.args, **kwargs: P.kwargs
+    ) -> "WorkflowHandle[R]":
+        from dbos.dbos import _get_dbos_instance
+
+        dbos = _get_dbos_instance()
+        return _start_workflow(dbos, func, self.name, False, *args, **kwargs)
+
+
+def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
+    while not stop_event.is_set():
+        time.sleep(1)
+        for queue_name, queue in dbos._registry.queue_info_map.items():
+            wf_ids = dbos._sys_db.start_queued_workflows(queue_name, queue.concurrency)
+            for id in wf_ids:
+                _execute_workflow_id(dbos, id)

dbos/recovery.py

@@ -41,7 +41,7 @@ def _recover_pending_workflows(
                 f"Skip local recovery because it's running in a VM: {os.environ.get('DBOS__VMID')}"
             )
         dbos.logger.debug(f"Recovering pending workflows for executor: {executor_id}")
-        workflow_ids = dbos.sys_db.get_pending_workflows(executor_id)
+        workflow_ids = dbos._sys_db.get_pending_workflows(executor_id)
         dbos.logger.debug(f"Pending workflows: {workflow_ids}")
 
         for workflowID in workflow_ids:

dbos/scheduler/scheduler.py

@@ -1,41 +1,39 @@
 import threading
-import traceback
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Callable
 
+from dbos.queue import Queue
+
 if TYPE_CHECKING:
     from dbos.dbos import _DBOSRegistry
 
 from ..context import SetWorkflowID
-from ..logger import dbos_logger
 from .croniter import croniter  # type: ignore
 
 ScheduledWorkflow = Callable[[datetime, datetime], None]
 
+scheduler_queue: Queue
+
 
 def scheduler_loop(
     func: ScheduledWorkflow, cron: str, stop_event: threading.Event
 ) -> None:
-    iter = croniter(cron, datetime.now(timezone.utc))
+    iter = croniter(cron, datetime.now(timezone.utc), second_at_beginning=True)
     while not stop_event.is_set():
         nextExecTime = iter.get_next(datetime)
         sleepTime = nextExecTime - datetime.now(timezone.utc)
         if stop_event.wait(timeout=sleepTime.total_seconds()):
             return
         with SetWorkflowID(f"sched-{func.__qualname__}-{nextExecTime.isoformat()}"):
-            try:
-                func(nextExecTime, datetime.now(timezone.utc))
-            except Exception as e:
-                dbos_logger.error(
-                    f"Exception encountered in scheduled workflow: {traceback.format_exc()}"
-                )
-                pass  # Let the thread keep running
+            scheduler_queue.enqueue(func, nextExecTime, datetime.now(timezone.utc))
 
 
 def scheduled(
     dbosreg: "_DBOSRegistry", cron: str
 ) -> Callable[[ScheduledWorkflow], ScheduledWorkflow]:
     def decorator(func: ScheduledWorkflow) -> ScheduledWorkflow:
+        global scheduler_queue
+        scheduler_queue = Queue("_dbos_internal_queue")
         stop_event = threading.Event()
         dbosreg.register_poller(stop_event, scheduler_loop, func, cron, stop_event)
         return func

dbos/schemas/system_database.py

@@ -1,5 +1,6 @@
 from sqlalchemy import (
     BigInteger,
+    Boolean,
     Column,
     ForeignKey,
     Index,
@@ -53,6 +54,7 @@ class SystemSchema:
             nullable=True,
             server_default=text("'0'::bigint"),
         ),
+        Column("queue_name", Text),
         Index("workflow_status_created_at_index", "created_at"),
         Index("workflow_status_executor_id_index", "executor_id"),
     )
@@ -139,3 +141,24 @@ class SystemSchema:
         Column("workflow_fn_name", Text, primary_key=True, nullable=False),
         Column("last_run_time", BigInteger, nullable=False),
     )
+
+    job_queue = Table(
+        "job_queue",
+        metadata_obj,
+        Column(
+            "workflow_uuid",
+            Text,
+            ForeignKey(
+                "workflow_status.workflow_uuid", onupdate="CASCADE", ondelete="CASCADE"
+            ),
+            nullable=False,
+            primary_key=True,
+        ),
+        Column("queue_name", Text, nullable=False),
+        Column(
+            "created_at_epoch_ms",
+            BigInteger,
+            nullable=False,
+            server_default=text("(EXTRACT(epoch FROM now()) * 1000::numeric)::bigint"),
+        ),
+    )