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

dbos/system_database.py

@@ -1,12 +1,11 @@
 import datetime
 import os
-import select
 import threading
 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
@@ -14,11 +13,7 @@ from alembic.config import Config
 from sqlalchemy.exc import DBAPIError
 
 import dbos.utils as utils
-from dbos.error import (
-    DBOSDuplicateWorkflowEventError,
-    DBOSNonExistentWorkflowError,
-    DBOSWorkflowConflictIDError,
-)
+from dbos.error import DBOSNonExistentWorkflowError, DBOSWorkflowConflictIDError
 
 from .dbos_config import ConfigFile
 from .logger import dbos_logger
@@ -33,10 +28,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 +57,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 +151,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 +169,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 +193,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 +244,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 +318,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 +339,7 @@ class SystemDatabase:
                 "authenticated_user": row[6],
                 "authenticated_roles": row[7],
                 "assumed_role": row[8],
+                "queue_name": row[9],
             }
             return status
 
@@ -379,6 +379,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 +400,7 @@ class SystemDatabase:
                 "authenticated_user": row[7],
                 "authenticated_roles": row[8],
                 "assumed_role": row[9],
+                "queue_name": row[10],
             }
             return status
 
@@ -565,11 +567,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 +623,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 +727,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()
 
@@ -839,20 +827,18 @@ class SystemDatabase:
             if recorded_output is not None:
                 return  # Already sent before
 
-            try:
-                c.execute(
-                    pg.insert(SystemSchema.workflow_events).values(
-                        workflow_uuid=workflow_uuid,
-                        key=key,
-                        value=utils.serialize(message),
-                    )
+            c.execute(
+                pg.insert(SystemSchema.workflow_events)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=utils.serialize(message),
                 )
-            except DBAPIError as dbapi_error:
-                if dbapi_error.orig.pgcode == "23505":  # type: ignore
-                    raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
-                raise dbapi_error
-            except Exception as e:
-                raise e
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key"],
+                    set_={"value": utils.serialize(message)},
+                )
+            )
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
@@ -1026,3 +1012,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.6.1.dist-info → dbos-0.7.0.dist-info}/METADATA

@@ -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.dist-info → dbos-0.7.0.dist-info}/RECORD

@@ -1,38 +1,40 @@
-dbos-0.6.1.dist-info/METADATA,sha256=3WgP4NAQT1ZNQbEkLV9hupdBqevptUrwS2k6WpJ6dss,5016
-dbos-0.6.1.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
-dbos-0.6.1.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
-dbos-0.6.1.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
-dbos/__init__.py,sha256=heuB3bqRXlVdfea9sKHBIVKSqFP6UuwhQecQfV4hyas,642
+dbos-0.7.0.dist-info/METADATA,sha256=l5ADRiFrTFSL2RzqUeXd6R8_CI9guOA5D_Zu75-Wexg,5008
+dbos-0.7.0.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
+dbos-0.7.0.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.7.0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos/__init__.py,sha256=jjlBFzSAzO2e-LD5IKJw7bRqZjrxpF5Sn-_JUJJptHU,680
 dbos/admin_sever.py,sha256=Qg5T3YRrbPW05PR_99yAaxgo1ugQrAp_uTeTqSfjm_k,3397
-dbos/application_database.py,sha256=1K3kE96BgGi_QWOd2heXluyNTwFAwlUVuAR6JKKUqf0,5659
+dbos/application_database.py,sha256=knFK8We8y6WrIpnFCKvFq5hvSuFQqUuJqOqDpSVMCPI,5521
 dbos/cli.py,sha256=z5dXbbnGWzSC3E1rfS8Lp1_OIImzcDKM7jP-iu_Q4aI,8602
-dbos/context.py,sha256=NVMGyvAa2RIiBVspvDz-8MBk_BQyGyYdPdorgO-GSng,16407
-dbos/core.py,sha256=jeqO8DABPAUrFlJXOfRfFDSnA8BGwiPnMa1JNbGuYLs,28584
+dbos/context.py,sha256=1Xp0i-ZP72Vj6eMdHuyfXi9RHnoT2w4MZ2Kx1CrKnQ8,16559
+dbos/core.py,sha256=LnQRHLjfpdlzlRkElduay8nObNMyRNC8SJSZEzqei7o,29161
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
-dbos/dbos.py,sha256=wzL51K7bY3J8grHXi0k0F0N09vFVBYWKwI0DGxyN4MY,29730
-dbos/dbos_config.py,sha256=ih_TD_1zTKhPKxk8TPdEIp3ihu82R06SGKg-s4rHxws,5344
+dbos/dbos.py,sha256=RtDcvKe4sm1TlnCGU4cyex-UI7hMMlhgzmOl1NuRLo4,29294
+dbos/dbos_config.py,sha256=NJVze2GkKgYUmcPP31Unb-QpsA0TzImEeQGJgVq6W6k,5352
 dbos/decorators.py,sha256=lbPefsLK6Cya4cb7TrOcLglOpGT3pc6qjZdsQKlfZLg,629
-dbos/error.py,sha256=DDhB0VHmoZE_CP51ICdFMZSL2gmVS3Dm0aPNWncci94,3876
-dbos/fastapi.py,sha256=s7LnwwYVpJm_QZZwBW5um8NV2Q2Qx85uVZqGcKlSZAo,1881
+dbos/error.py,sha256=c2y7d3Cbb-ZOO-M9txcCxuyPE1bdnwKrJnXgJXYb-pQ,3437
+dbos/fastapi.py,sha256=gx9hlpxYOiwbuhSlbY9bn5C-F_FsCbrJvkX9ZAvDG6U,3418
 dbos/flask.py,sha256=azr4geMEGuuTBCyxIZmgDmmP-6s_pTIF-lGyp9Q4IB8,2430
-dbos/kafka.py,sha256=FtngQHBu2TKfyDF7GFsKJAawFQJiOFxgKEUlNNxrdrw,3055
+dbos/kafka.py,sha256=LH3hbNapnkjLcuXNUtdGU0398JafWb-t0GwUl3LOzkc,3645
 dbos/kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
 dbos/logger.py,sha256=D-aFSZUCHBP34J1IZ5YNkTrJW-rDiH3py_v9jLU4Yrk,3565
 dbos/migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
 dbos/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
 dbos/migrations/versions/5c361fc04708_added_system_tables.py,sha256=QMgFMb0aLgC25YicsvPSr6AHRCA6Zd66hyaRUhwKzrQ,6404
 dbos/migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEULYsPfnp4YzaO5L198gDTgYY39AVA,2022
+dbos/migrations/versions/eab0cc1d9a14_job_queue.py,sha256=_9-FCW-zOpCQfblTS_yRLtFiUaWlC1tM4BoKBTDeH9k,1395
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
-dbos/recovery.py,sha256=xfwQFWNuD8DXg5HD5_-3tG7Neo9j-x1lrqiwtn5FSh8,2015
+dbos/queue.py,sha256=5unMPX1ThoVWfXOR2IUSeIjgATgcX_YmZS3Fz1Uoc7o,1228
+dbos/recovery.py,sha256=zqtO_ExGoIErLMVnbneU3VeHLVWvhV4jnfqssAVlQQk,2016
 dbos/registrations.py,sha256=gMI-u05tv5bpvyddQGtoUgCsqARx51aOY7p0JXPafQo,6539
 dbos/request.py,sha256=-FIwtknayvRl6OjvqO4V2GySVzSdP1Ft3cc9ZBS-PLY,928
 dbos/roles.py,sha256=7Lh7uwUq1dpa6TXCOHre4mPTd5qmXzK_QPkvYR52DXg,2285
 dbos/scheduler/croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
-dbos/scheduler/scheduler.py,sha256=uO4_9jmWW2rLv1ODL3lc1cE_37ZaVTgnvmFx_FAlN50,1472
+dbos/scheduler/scheduler.py,sha256=Sz4EIpAtur7so2YajTic64GrTpa4qPw8QxXn0M34v80,1360
 dbos/schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/schemas/application_database.py,sha256=q_Wr2XbiZNBYFkOtu7uKavo1T_cSOBblxKGHThYGGsY,962
-dbos/schemas/system_database.py,sha256=5V3vqnEzry0Hn7ZbVS9Gs_dJKia8uX8p7mGC82Ru8rk,4303
-dbos/system_database.py,sha256=BZ8yE79FH9NZhEt8AAYNXy0-wMcKLBAq8VPdZbfidR4,40564
+dbos/schemas/system_database.py,sha256=ed4c1UntsD-cqXD0ekM4jvcYYEViavDh_G6c0pVDe7k,4938
+dbos/system_database.py,sha256=VSGFSPubbMokGYsZfRb6cQPltLfjoWM-Re_2Gj9qkRc,41844
 dbos/templates/hello/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/templates/hello/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/templates/hello/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -46,4 +48,4 @@ dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKs
 dbos/tracer.py,sha256=GaXDhdKKF_IQp5SAMipGXiDVwteRKjNbrXyYCH1mor0,2520
 dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.6.1.dist-info/RECORD,,
+dbos-0.7.0.dist-info/RECORD,,

{dbos-0.6.1.dist-info → dbos-0.7.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: pdm-backend (2.3.3)
+Generator: pdm-backend (2.4.1)
 Root-Is-Purelib: true
 Tag: py3-none-any

{dbos-0.6.1.dist-info → dbos-0.7.0.dist-info}/entry_points.txt

@@ -1,3 +1,5 @@
 [console_scripts]
 dbos = dbos.cli:app
 
+[gui_scripts]
+