dbos 0.7.0a1__py3-none-any.whl → 0.7.0a8__py3-none-any.whl

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"),
+        ),
+    )

dbos/system_database.py

@@ -6,7 +6,7 @@ import time
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Sequence, Set, TypedDict, cast
 
-import psycopg2
+import psycopg
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from alembic import command
@@ -33,10 +33,11 @@ class WorkflowStatusString(Enum):
     ERROR = "ERROR"
     RETRIES_EXCEEDED = "RETRIES_EXCEEDED"
     CANCELLED = "CANCELLED"
+    ENQUEUED = "ENQUEUED"
 
 
 WorkflowStatuses = Literal[
-    "PENDING", "SUCCESS", "ERROR", "RETRIES_EXCEEDED", "CANCELLED"
+    "PENDING", "SUCCESS", "ERROR", "RETRIES_EXCEEDED", "CANCELLED", "ENQUEUED"
 ]
 
 
@@ -61,6 +62,7 @@ class WorkflowStatusInternal(TypedDict):
     authenticated_user: Optional[str]
     assumed_role: Optional[str]
     authenticated_roles: Optional[str]  # JSON list of roles.
+    queue_name: Optional[str]
 
 
 class RecordedResult(TypedDict):
@@ -154,7 +156,7 @@ class SystemDatabase:
 
         # If the system 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"],
@@ -172,7 +174,7 @@ class SystemDatabase:
         engine.dispose()
 
         system_db_url = sa.URL.create(
-            "postgresql",
+            "postgresql+psycopg",
             username=config["database"]["username"],
             password=config["database"]["password"],
             host=config["database"]["hostname"],
@@ -196,7 +198,7 @@ class SystemDatabase:
         )
         command.upgrade(alembic_cfg, "head")
 
-        self.notification_conn: Optional[psycopg2.extensions.connection] = None
+        self.notification_conn: Optional[psycopg.connection.Connection] = None
         self.notifications_map: Dict[str, threading.Condition] = {}
         self.workflow_events_map: Dict[str, threading.Condition] = {}
 
@@ -247,6 +249,7 @@ class SystemDatabase:
             authenticated_user=status["authenticated_user"],
             authenticated_roles=status["authenticated_roles"],
             assumed_role=status["assumed_role"],
+            queue_name=status["queue_name"],
         )
         if replace:
             cmd = cmd.on_conflict_do_update(
@@ -320,6 +323,7 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.authenticated_user,
                     SystemSchema.workflow_status.c.authenticated_roles,
                     SystemSchema.workflow_status.c.assumed_role,
+                    SystemSchema.workflow_status.c.queue_name,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -340,6 +344,7 @@ class SystemDatabase:
                 "authenticated_user": row[6],
                 "authenticated_roles": row[7],
                 "assumed_role": row[8],
+                "queue_name": row[9],
             }
             return status
 
@@ -379,6 +384,7 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.authenticated_user,
                     SystemSchema.workflow_status.c.authenticated_roles,
                     SystemSchema.workflow_status.c.assumed_role,
+                    SystemSchema.workflow_status.c.queue_name,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -399,6 +405,7 @@ class SystemDatabase:
                 "authenticated_user": row[7],
                 "authenticated_roles": row[8],
                 "assumed_role": row[9],
+                "queue_name": row[10],
             }
             return status
 
@@ -565,11 +572,9 @@ class SystemDatabase:
                 with self.engine.begin() as c:
                     c.execute(sql)
         except DBAPIError as dbapi_error:
-            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
                 raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
-            raise dbapi_error
-        except Exception as e:
-            raise e
+            raise
 
     def check_operation_execution(
         self, workflow_uuid: str, function_id: int, conn: Optional[sa.Connection] = None
@@ -623,11 +628,9 @@ class SystemDatabase:
                 )
             except DBAPIError as dbapi_error:
                 # Foreign key violation
-                if dbapi_error.orig.pgcode == "23503":  # type: ignore
+                if dbapi_error.orig.sqlstate == "23503":  # type: ignore
                     raise DBOSNonExistentWorkflowError(destination_uuid)
-                raise dbapi_error
-            except Exception as e:
-                raise e
+                raise
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
@@ -729,69 +732,59 @@ class SystemDatabase:
         return message
 
     def _notification_listener(self) -> None:
-        notification_cursor: Optional[psycopg2.extensions.cursor] = None
         while self._run_background_processes:
             try:
-                # Listen to notifications
-                self.notification_conn = psycopg2.connect(
-                    self.engine.url.render_as_string(hide_password=False)
+                # since we're using the psycopg connection directly, we need a url without the "+pycopg" suffix
+                url = sa.URL.create(
+                    "postgresql", **self.engine.url.translate_connect_args()
                 )
-                self.notification_conn.set_isolation_level(
-                    psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT
+                # Listen to notifications
+                self.notification_conn = psycopg.connect(
+                    url.render_as_string(hide_password=False), autocommit=True
                 )
-                notification_cursor = self.notification_conn.cursor()
 
-                notification_cursor.execute("LISTEN dbos_notifications_channel")
-                notification_cursor.execute("LISTEN dbos_workflow_events_channel")
+                self.notification_conn.execute("LISTEN dbos_notifications_channel")
+                self.notification_conn.execute("LISTEN dbos_workflow_events_channel")
+
                 while self._run_background_processes:
-                    if select.select([self.notification_conn], [], [], 60) == (
-                        [],
-                        [],
-                        [],
-                    ):
-                        continue
-                    else:
-                        self.notification_conn.poll()
-                        while self.notification_conn.notifies:
-                            notify = self.notification_conn.notifies.pop(0)
-                            channel = notify.channel
-                            dbos_logger.debug(
-                                f"Received notification on channel: {channel}, payload: {notify.payload}"
-                            )
-                            if channel == "dbos_notifications_channel":
-                                if (
-                                    notify.payload
-                                    and notify.payload in self.notifications_map
-                                ):
-                                    condition = self.notifications_map[notify.payload]
-                                    condition.acquire()
-                                    condition.notify_all()
-                                    condition.release()
-                                    dbos_logger.debug(
-                                        f"Signaled notifications condition for {notify.payload}"
-                                    )
-                            elif channel == "dbos_workflow_events_channel":
-                                if (
-                                    notify.payload
-                                    and notify.payload in self.workflow_events_map
-                                ):
-                                    condition = self.workflow_events_map[notify.payload]
-                                    condition.acquire()
-                                    condition.notify_all()
-                                    condition.release()
-                                    dbos_logger.debug(
-                                        f"Signaled workflow_events condition for {notify.payload}"
-                                    )
-                            else:
-                                dbos_logger.error(f"Unknown channel: {channel}")
+                    gen = self.notification_conn.notifies(timeout=60)
+                    for notify in gen:
+                        channel = notify.channel
+                        dbos_logger.debug(
+                            f"Received notification on channel: {channel}, payload: {notify.payload}"
+                        )
+                        if channel == "dbos_notifications_channel":
+                            if (
+                                notify.payload
+                                and notify.payload in self.notifications_map
+                            ):
+                                condition = self.notifications_map[notify.payload]
+                                condition.acquire()
+                                condition.notify_all()
+                                condition.release()
+                                dbos_logger.debug(
+                                    f"Signaled notifications condition for {notify.payload}"
+                                )
+                        elif channel == "dbos_workflow_events_channel":
+                            if (
+                                notify.payload
+                                and notify.payload in self.workflow_events_map
+                            ):
+                                condition = self.workflow_events_map[notify.payload]
+                                condition.acquire()
+                                condition.notify_all()
+                                condition.release()
+                                dbos_logger.debug(
+                                    f"Signaled workflow_events condition for {notify.payload}"
+                                )
+                        else:
+                            dbos_logger.error(f"Unknown channel: {channel}")
             except Exception as e:
                 if self._run_background_processes:
                     dbos_logger.error(f"Notification listener error: {e}")
                     time.sleep(1)
                     # Then the loop will try to reconnect and restart the listener
             finally:
-                if notification_cursor is not None:
-                    notification_cursor.close()
                 if self.notification_conn is not None:
                     self.notification_conn.close()
 
@@ -848,11 +841,9 @@ class SystemDatabase:
                     )
                 )
             except DBAPIError as dbapi_error:
-                if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                if dbapi_error.orig.sqlstate == "23505":  # type: ignore
                     raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
-                raise dbapi_error
-            except Exception as e:
-                raise e
+                raise
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
@@ -1026,3 +1017,50 @@ class SystemDatabase:
             len(self._workflow_status_buffer) == 0
             and len(self._workflow_inputs_buffer) == 0
         )
+
+    def enqueue(self, workflow_id: str, queue_name: str) -> None:
+        with self.engine.begin() as c:
+            c.execute(
+                pg.insert(SystemSchema.job_queue)
+                .values(
+                    workflow_uuid=workflow_id,
+                    queue_name=queue_name,
+                )
+                .on_conflict_do_nothing()
+            )
+
+    def start_queued_workflows(
+        self, queue_name: str, concurrency: Optional[int]
+    ) -> List[str]:
+        with self.engine.begin() as c:
+            query = sa.select(SystemSchema.job_queue.c.workflow_uuid).where(
+                SystemSchema.job_queue.c.queue_name == queue_name
+            )
+            if concurrency is not None:
+                query = query.order_by(
+                    SystemSchema.job_queue.c.created_at_epoch_ms.asc()
+                ).limit(concurrency)
+            rows = c.execute(query).fetchall()
+            dequeued_ids: List[str] = [row[0] for row in rows]
+            ret_ids = []
+            for id in dequeued_ids:
+                result = c.execute(
+                    SystemSchema.workflow_status.update()
+                    .where(SystemSchema.workflow_status.c.workflow_uuid == id)
+                    .where(
+                        SystemSchema.workflow_status.c.status
+                        == WorkflowStatusString.ENQUEUED.value
+                    )
+                    .values(status=WorkflowStatusString.PENDING.value)
+                )
+                if result.rowcount > 0:
+                    ret_ids.append(id)
+            return ret_ids
+
+    def remove_from_queue(self, workflow_id: str) -> None:
+        with self.engine.begin() as c:
+            c.execute(
+                sa.delete(SystemSchema.job_queue).where(
+                    SystemSchema.job_queue.c.workflow_uuid == workflow_id
+                )
+            )

{dbos-0.7.0a1.dist-info → dbos-0.7.0a8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.7.0a1
+Version: 0.7.0a8
 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