dbos 0.7.0a0__tar.gz → 0.7.0a5__tar.gz

{dbos-0.7.0a0 → dbos-0.7.0a5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.7.0a0
+Version: 0.7.0a5
 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
 
 
@@ -91,7 +91,7 @@ def step_two():
     print("Step two completed!")
 
 @DBOS.workflow()
-def workflow():
+def dbos_workflow():
     step_one()
     for _ in range(5):
         print("Press Control + \ to stop the app...")
@@ -99,8 +99,8 @@ def workflow():
     step_two()
 
 @app.get("/")
-def endpoint():
-    workflow()
+def fastapi_endpoint():
+    dbos_workflow()
 ```
 
 Save the program into `main.py`, edit `dbos-config.yaml` to configure your Postgres connection settings, and start it with `fastapi run`.

{dbos-0.7.0a0 → dbos-0.7.0a5}/README.md

@@ -68,7 +68,7 @@ def step_two():
     print("Step two completed!")
 
 @DBOS.workflow()
-def workflow():
+def dbos_workflow():
     step_one()
     for _ in range(5):
         print("Press Control + \ to stop the app...")
@@ -76,8 +76,8 @@ def workflow():
     step_two()
 
 @app.get("/")
-def endpoint():
-    workflow()
+def fastapi_endpoint():
+    dbos_workflow()
 ```
 
 Save the program into `main.py`, edit `dbos-config.yaml` to configure your Postgres connection settings, and start it with `fastapi run`.

{dbos-0.7.0a0 → dbos-0.7.0a5}/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.7.0a0 → dbos-0.7.0a5}/dbos/cli.py

@@ -12,7 +12,7 @@ from typing import Any
 import tomlkit
 import typer
 from rich import print
-from rich.prompt import Confirm, Prompt
+from rich.prompt import Prompt
 from typing_extensions import Annotated
 
 from dbos import load_config
@@ -30,6 +30,7 @@ def on_windows() -> bool:
 def start() -> None:
     config = load_config()
     start_commands = config["runtimeConfig"]["start"]
+    typer.echo("Executing start commands from 'dbos-config.yaml'")
     for command in start_commands:
         typer.echo(f"Executing: {command}")
 
@@ -129,9 +130,12 @@ def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
         "project_name": project_name,
         "package_name": package_name,
         "db_name": db_name,
+        "migration_command": "alembic upgrade head",
     }
 
     if config_mode:
+        ctx["package_name"] = "."
+        ctx["migration_command"] = "echo 'No migrations specified'"
         copy_dbos_template(
             os.path.join(src_dir, "dbos-config.yaml.dbos"),
             os.path.join(dst_dir, "dbos-config.yaml"),
@@ -244,6 +248,7 @@ def migrate() -> None:
             app_db.destroy()
 
     # Next, run any custom migration commands specified in the configuration
+    typer.echo("Executing migration commands from 'dbos-config.yaml'")
     try:
         migrate_commands = (
             config["database"]["migrate"]

{dbos-0.7.0a0 → dbos-0.7.0a5}/dbos/core.py

@@ -192,7 +192,7 @@ def _execute_workflow(
         status["status"] = "ERROR"
         status["error"] = utils.serialize(error)
         dbos.sys_db.update_workflow_status(status)
-        raise error
+        raise
 
     return output
 
@@ -217,7 +217,7 @@ 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]":
@@ -493,7 +493,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 +505,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
+                            raise
             return output
 
         fi = get_or_create_func_info(func)

{dbos-0.7.0a0 → dbos-0.7.0a5}/dbos/dbos.py

@@ -277,32 +277,9 @@ class DBOS:
 
         # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
-            from fastapi.requests import Request as FARequest
-            from fastapi.responses import JSONResponse
-
-            async def dbos_error_handler(
-                request: FARequest, 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__),
-                    },
-                )
-
-            self.fastapi.add_exception_handler(DBOSException, dbos_error_handler)
-
             from dbos.fastapi import setup_fastapi_middleware
 
-            setup_fastapi_middleware(self.fastapi)
-            self.fastapi.on_event("startup")(self._launch)
-            self.fastapi.on_event("shutdown")(self._destroy)
+            setup_fastapi_middleware(self.fastapi, _get_dbos_instance())
 
         # If using Flask, set up middleware
         if self.flask is not None:

{dbos-0.7.0a0 → dbos-0.7.0a5}/dbos/dbos_config.py

@@ -105,7 +105,7 @@ def get_dbos_database_url(config_file_path: str = "dbos-config.yaml") -> str:
     """
     dbos_config = load_config(config_file_path)
     db_url = URL.create(
-        "postgresql",
+        "postgresql+psycopg",
         username=dbos_config["database"]["username"],
         password=dbos_config["database"]["password"],
         host=dbos_config["database"]["hostname"],

{dbos-0.7.0a0 → dbos-0.7.0a5}/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-0.7.0a0 → dbos-0.7.0a5}/dbos/scheduler/scheduler.py

@@ -16,7 +16,7 @@ ScheduledWorkflow = Callable[[datetime, datetime], None]
 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)

{dbos-0.7.0a0 → dbos-0.7.0a5}/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
@@ -154,7 +154,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 +172,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 +196,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] = {}
 
@@ -565,11 +565,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 +621,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 +725,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 +834,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,

{dbos-0.7.0a0 → dbos-0.7.0a5}/dbos/templates/hello/dbos-config.yaml.dbos

@@ -15,7 +15,7 @@ database:
   password: ${PGPASSWORD}
   app_db_name: ${db_name}
   migrate:
-    - alembic upgrade head
+    - ${migration_command}
 telemetry:
   logs:
     logLevel: INFO

{dbos-0.7.0a0 → dbos-0.7.0a5}/pyproject.toml

@@ -9,7 +9,6 @@ dependencies = [
     "pyyaml>=6.0.2",
     "jsonschema>=4.23.0",
     "alembic>=1.13.2",
-    "psycopg2-binary>=2.9.9",
     "typing-extensions>=4.12.2; python_version < \"3.10\"",
     "typer>=0.12.3",
     "jsonpickle>=3.2.2",
@@ -20,10 +19,11 @@ dependencies = [
     "fastapi[standard]>=0.112.1",
     "psutil>=6.0.0",
     "tomlkit>=0.13.2",
+    "psycopg>=3.2.1",
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.7.0a0"
+version = "0.7.0a5"
 
 [project.license]
 text = "MIT"
@@ -58,7 +58,6 @@ dev = [
     "requests>=2.32.3",
     "types-requests>=2.32.0.20240712",
     "httpx>=0.27.0",
-    "types-psycopg2>=2.9.21.20240417",
     "pytz>=2024.1",
     "GitPython>=3.1.43",
     "confluent-kafka>=2.5.3",

{dbos-0.7.0a0 → dbos-0.7.0a5}/tests/conftest.py

@@ -48,7 +48,7 @@ def config() -> ConfigFile:
 def postgres_db_engine() -> sa.Engine:
     cfg = default_config()
     postgres_db_url = sa.URL.create(
-        "postgresql",
+        "postgresql+psycopg",
         username=cfg["database"]["username"],
         password=cfg["database"]["password"],
         host=cfg["database"]["hostname"],

{dbos-0.7.0a0 → dbos-0.7.0a5}/tests/scheduler/test_scheduler.py

@@ -9,14 +9,14 @@ from dbos import DBOS
 def test_scheduled_workflow(dbos: DBOS) -> None:
     wf_counter: int = 0
 
-    @DBOS.scheduled("* * * * * *")
+    @DBOS.scheduled("*/2 * * * * *")
     @DBOS.workflow()
     def test_workflow(scheduled: datetime, actual: datetime) -> None:
         nonlocal wf_counter
         wf_counter += 1
 
-    time.sleep(2)
-    assert wf_counter >= 1 and wf_counter <= 3
+    time.sleep(4)
+    assert wf_counter > 1 and wf_counter <= 3
 
 
 def test_scheduled_workflow_exception(dbos: DBOS) -> None:

{dbos-0.7.0a0 → dbos-0.7.0a5}/tests/test_failures.py

@@ -4,7 +4,8 @@ import uuid
 
 import pytest
 import sqlalchemy as sa
-from sqlalchemy.exc import DBAPIError
+from psycopg.errors import SerializationFailure
+from sqlalchemy.exc import OperationalError
 
 # Public API
 from dbos import DBOS, GetWorkflowsInput, SetWorkflowID
@@ -18,10 +19,9 @@ def test_transaction_errors(dbos: DBOS) -> None:
         nonlocal retry_counter
         if retry_counter < max_retry:
             retry_counter += 1
-            base_err = BaseException()
-            base_err.pgcode = "40001"  # type: ignore
-            err = DBAPIError("Serialization test error", {}, base_err)
-            raise err
+            raise OperationalError(
+                "Serialization test error", {}, SerializationFailure()
+            )
         return max_retry
 
     @DBOS.transaction()
@@ -36,7 +36,7 @@ def test_transaction_errors(dbos: DBOS) -> None:
 
     with pytest.raises(Exception) as exc_info:
         test_noretry_transaction()
-    assert exc_info.value.orig.pgcode == "42601"  # type: ignore
+    assert exc_info.value.orig.sqlstate == "42601"  # type: ignore
     assert retry_counter == 11
 
 
@@ -101,7 +101,7 @@ def test_buffer_flush_errors(dbos: DBOS) -> None:
     # Crash the system database connection and make sure the buffer flush works on time.
     backup_engine = dbos.sys_db.engine
     dbos.sys_db.engine = sa.create_engine(
-        "postgresql+psycopg2://fake:database@localhost/fake_db"
+        "postgresql+psycopg://fake:database@localhost/fake_db"
     )
 
     res = test_transaction("bob")