dbos 1.12.0a3__py3-none-any.whl → 1.13.0a5__py3-none-any.whl

dbos/_alembic_migrations/versions/471b60d64126_dbos_migrations.py

@@ -0,0 +1,35 @@
+"""dbos_migrations
+
+Revision ID: 471b60d64126
+Revises: 01ce9f07bd10
+Create Date: 2025-08-21 14:22:31.455266
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "471b60d64126"
+down_revision: Union[str, None] = "01ce9f07bd10"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Create dbos_migrations table
+    op.create_table(
+        "dbos_migrations",
+        sa.Column("version", sa.BigInteger(), nullable=False),
+        sa.PrimaryKeyConstraint("version"),
+        schema="dbos",
+    )
+
+    # Insert initial version 1
+    op.execute("INSERT INTO dbos.dbos_migrations (version) VALUES (1)")
+
+
+def downgrade() -> None:
+    op.drop_table("dbos_migrations", schema="dbos")

dbos/_app_db.py

@@ -1,11 +1,14 @@
+from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, TypedDict
 
+import psycopg
 import sqlalchemy as sa
-import sqlalchemy.dialects.postgresql as pg
 from sqlalchemy import inspect, text
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
+from dbos._migration import get_sqlite_timestamp_expr
+
 from . import _serialization
 from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
 from ._logger import dbos_logger
@@ -29,7 +32,7 @@ class RecordedResult(TypedDict):
     error: Optional[str]  # JSON (jsonpickle)
 
 
-class ApplicationDatabase:
+class ApplicationDatabase(ABC):
 
     def __init__(
         self,
@@ -38,95 +41,37 @@ class ApplicationDatabase:
         engine_kwargs: Dict[str, Any],
         debug_mode: bool = False,
     ):
-        app_db_url = sa.make_url(database_url).set(drivername="postgresql+psycopg")
-
-        if engine_kwargs is None:
-            engine_kwargs = {}
-
-        self.engine = sa.create_engine(
-            app_db_url,
-            **engine_kwargs,
-        )
+        self.engine = self._create_engine(database_url, engine_kwargs)
         self._engine_kwargs = engine_kwargs
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode
 
-    def run_migrations(self) -> None:
-        if self.debug_mode:
-            dbos_logger.warning(
-                "Application database migrations are skipped in debug mode."
-            )
-            return
-        # Check if the database exists
-        app_db_url = self.engine.url
-        postgres_db_engine = sa.create_engine(
-            app_db_url.set(database="postgres"),
-            **self._engine_kwargs,
-        )
-        with postgres_db_engine.connect() as conn:
-            conn.execution_options(isolation_level="AUTOCOMMIT")
-            if not conn.execute(
-                sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
-                parameters={"db_name": app_db_url.database},
-            ).scalar():
-                conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
-        postgres_db_engine.dispose()
-
-        # Create the dbos schema and transaction_outputs table in the application database
-        with self.engine.begin() as conn:
-            # Check if schema exists first
-            schema_exists = conn.execute(
-                sa.text(
-                    "SELECT 1 FROM information_schema.schemata WHERE schema_name = :schema_name"
-                ),
-                parameters={"schema_name": ApplicationSchema.schema},
-            ).scalar()
-
-            if not schema_exists:
-                schema_creation_query = sa.text(
-                    f"CREATE SCHEMA {ApplicationSchema.schema}"
-                )
-                conn.execute(schema_creation_query)
-
-        inspector = inspect(self.engine)
-        if not inspector.has_table(
-            "transaction_outputs", schema=ApplicationSchema.schema
-        ):
-            ApplicationSchema.metadata_obj.create_all(self.engine)
-        else:
-            columns = inspector.get_columns(
-                "transaction_outputs", schema=ApplicationSchema.schema
-            )
-            column_names = [col["name"] for col in columns]
+    @abstractmethod
+    def _create_engine(
+        self, database_url: str, engine_kwargs: Dict[str, Any]
+    ) -> sa.Engine:
+        """Create a database engine specific to the database type."""
+        pass
 
-            if "function_name" not in column_names:
-                # Column missing, alter table to add it
-                with self.engine.connect() as conn:
-                    conn.execute(
-                        text(
-                            f"""
-                        ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
-                        ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
-                        """
-                        )
-                    )
-                    conn.commit()
+    @abstractmethod
+    def run_migrations(self) -> None:
+        """Run database migrations specific to the database type."""
+        pass
 
     def destroy(self) -> None:
         self.engine.dispose()
 
-    @staticmethod
     def record_transaction_output(
-        session: Session, output: TransactionResultInternal
+        self, session: Session, output: TransactionResultInternal
     ) -> None:
         try:
             session.execute(
-                pg.insert(ApplicationSchema.transaction_outputs).values(
+                sa.insert(ApplicationSchema.transaction_outputs).values(
                     workflow_uuid=output["workflow_uuid"],
                     function_id=output["function_id"],
                     output=output["output"],
                     error=None,
-                    txn_id=sa.text("(select pg_current_xact_id_if_assigned()::text)"),
+                    txn_id="",
                     txn_snapshot=output["txn_snapshot"],
                     executor_id=(
                         output["executor_id"] if output["executor_id"] else None
@@ -135,7 +80,7 @@ class ApplicationDatabase:
                 )
             )
         except DBAPIError as dbapi_error:
-            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
+            if self._is_unique_constraint_violation(dbapi_error):
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
             raise
 
@@ -145,14 +90,12 @@ class ApplicationDatabase:
         try:
             with self.engine.begin() as conn:
                 conn.execute(
-                    pg.insert(ApplicationSchema.transaction_outputs).values(
+                    sa.insert(ApplicationSchema.transaction_outputs).values(
                         workflow_uuid=output["workflow_uuid"],
                         function_id=output["function_id"],
                         output=None,
                         error=output["error"],
-                        txn_id=sa.text(
-                            "(select pg_current_xact_id_if_assigned()::text)"
-                        ),
+                        txn_id="",
                         txn_snapshot=output["txn_snapshot"],
                         executor_id=(
                             output["executor_id"] if output["executor_id"] else None
@@ -161,7 +104,7 @@ class ApplicationDatabase:
                     )
                 )
         except DBAPIError as dbapi_error:
-            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
+            if self._is_unique_constraint_violation(dbapi_error):
                 raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
             raise
 
@@ -283,3 +226,197 @@ class ApplicationDatabase:
                 )
 
             c.execute(delete_query)
+
+    @abstractmethod
+    def _is_unique_constraint_violation(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a unique constraint violation."""
+        pass
+
+    @abstractmethod
+    def _is_serialization_error(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a serialization/concurrency error."""
+        pass
+
+    @staticmethod
+    def create(
+        database_url: str,
+        engine_kwargs: Dict[str, Any],
+        debug_mode: bool = False,
+    ) -> "ApplicationDatabase":
+        """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
+        if database_url.startswith("sqlite"):
+            return SQLiteApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                debug_mode=debug_mode,
+            )
+        else:
+            # Default to PostgreSQL for postgresql://, postgres://, or other URLs
+            return PostgresApplicationDatabase(
+                database_url=database_url,
+                engine_kwargs=engine_kwargs,
+                debug_mode=debug_mode,
+            )
+
+
+class PostgresApplicationDatabase(ApplicationDatabase):
+    """PostgreSQL-specific implementation of ApplicationDatabase."""
+
+    def _create_engine(
+        self, database_url: str, engine_kwargs: Dict[str, Any]
+    ) -> sa.Engine:
+        """Create a PostgreSQL engine."""
+        app_db_url = sa.make_url(database_url).set(drivername="postgresql+psycopg")
+
+        if engine_kwargs is None:
+            engine_kwargs = {}
+
+        # TODO: Make the schema dynamic so this isn't needed
+        ApplicationSchema.transaction_outputs.schema = "dbos"
+
+        return sa.create_engine(
+            app_db_url,
+            **engine_kwargs,
+        )
+
+    def run_migrations(self) -> None:
+        if self.debug_mode:
+            dbos_logger.warning(
+                "Application database migrations are skipped in debug mode."
+            )
+            return
+        # Check if the database exists
+        app_db_url = self.engine.url
+        postgres_db_engine = sa.create_engine(
+            app_db_url.set(database="postgres"),
+            **self._engine_kwargs,
+        )
+        with postgres_db_engine.connect() as conn:
+            conn.execution_options(isolation_level="AUTOCOMMIT")
+            if not conn.execute(
+                sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
+                parameters={"db_name": app_db_url.database},
+            ).scalar():
+                conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
+        postgres_db_engine.dispose()
+
+        # Create the dbos schema and transaction_outputs table in the application database
+        with self.engine.begin() as conn:
+            # Check if schema exists first
+            schema_exists = conn.execute(
+                sa.text(
+                    "SELECT 1 FROM information_schema.schemata WHERE schema_name = :schema_name"
+                ),
+                parameters={"schema_name": ApplicationSchema.schema},
+            ).scalar()
+
+            if not schema_exists:
+                schema_creation_query = sa.text(
+                    f"CREATE SCHEMA {ApplicationSchema.schema}"
+                )
+                conn.execute(schema_creation_query)
+
+        inspector = inspect(self.engine)
+        if not inspector.has_table(
+            "transaction_outputs", schema=ApplicationSchema.schema
+        ):
+            ApplicationSchema.metadata_obj.create_all(self.engine)
+        else:
+            columns = inspector.get_columns(
+                "transaction_outputs", schema=ApplicationSchema.schema
+            )
+            column_names = [col["name"] for col in columns]
+
+            if "function_name" not in column_names:
+                # Column missing, alter table to add it
+                with self.engine.connect() as conn:
+                    conn.execute(
+                        text(
+                            f"""
+                        ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
+                        ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
+                        """
+                        )
+                    )
+                    conn.commit()
+
+    def _is_unique_constraint_violation(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a unique constraint violation in PostgreSQL."""
+        return dbapi_error.orig.sqlstate == "23505"  # type: ignore
+
+    def _is_serialization_error(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a serialization/concurrency error in PostgreSQL."""
+        # 40001: serialization_failure (MVCC conflict)
+        # 40P01: deadlock_detected
+        driver_error = dbapi_error.orig
+        return (
+            driver_error is not None
+            and isinstance(driver_error, psycopg.OperationalError)
+            and driver_error.sqlstate in ("40001", "40P01")
+        )
+
+
+class SQLiteApplicationDatabase(ApplicationDatabase):
+    """SQLite-specific implementation of ApplicationDatabase."""
+
+    def _create_engine(
+        self, database_url: str, engine_kwargs: Dict[str, Any]
+    ) -> sa.Engine:
+        """Create a SQLite engine."""
+        # TODO: Make the schema dynamic so this isn't needed
+        ApplicationSchema.transaction_outputs.schema = None
+        return sa.create_engine(database_url)
+
+    def run_migrations(self) -> None:
+        if self.debug_mode:
+            dbos_logger.warning(
+                "Application database migrations are skipped in debug mode."
+            )
+            return
+
+        with self.engine.begin() as conn:
+            # Check if table exists
+            result = conn.execute(
+                sa.text(
+                    "SELECT name FROM sqlite_master WHERE type='table' AND name='transaction_outputs'"
+                )
+            ).fetchone()
+
+            if result is None:
+                # Create the table with proper SQLite syntax
+                conn.execute(
+                    sa.text(
+                        f"""
+                        CREATE TABLE transaction_outputs (
+                            workflow_uuid TEXT NOT NULL,
+                            function_id INTEGER NOT NULL,
+                            output TEXT,
+                            error TEXT,
+                            txn_id TEXT,
+                            txn_snapshot TEXT NOT NULL,
+                            executor_id TEXT,
+                            function_name TEXT NOT NULL DEFAULT '',
+                            created_at BIGINT NOT NULL DEFAULT {get_sqlite_timestamp_expr()},
+                            PRIMARY KEY (workflow_uuid, function_id)
+                        )
+                        """
+                    )
+                )
+                # Create the index
+                conn.execute(
+                    sa.text(
+                        "CREATE INDEX transaction_outputs_created_at_index ON transaction_outputs (created_at)"
+                    )
+                )
+
+    def _is_unique_constraint_violation(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a unique constraint violation in SQLite."""
+        return "UNIQUE constraint failed" in str(dbapi_error.orig)
+
+    def _is_serialization_error(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a serialization/concurrency error in SQLite."""
+        # SQLite database is locked or busy errors
+        error_msg = str(dbapi_error.orig).lower()
+        return (
+            "database is locked" in error_msg or "database table is locked" in error_msg
+        )

dbos/_client.py

@@ -16,6 +16,7 @@ from typing import (
 
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
+from dbos._sys_db import SystemDatabase
 
 if sys.version_info < (3, 11):
     from typing_extensions import NotRequired
@@ -119,7 +120,7 @@ class DBOSClient:
     ):
         assert is_valid_database_url(database_url)
         # We only create database connections but do not run migrations
-        self._sys_db = SystemDatabase(
+        self._sys_db = SystemDatabase.create(
             system_database_url=get_system_database_url(
                 {
                     "system_database_url": system_database_url,
@@ -134,7 +135,7 @@ class DBOSClient:
             },
         )
         self._sys_db.check_connection()
-        self._app_db = ApplicationDatabase(
+        self._app_db = ApplicationDatabase.create(
             database_url=database_url,
             engine_kwargs={
                 "pool_timeout": 30,

dbos/_context.py

@@ -221,6 +221,8 @@ class DBOSContext:
         return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
+        if dbos_tracer.disable_otlp:
+            return
         attributes["operationUUID"] = (
             self.workflow_id if len(self.workflow_id) > 0 else None
         )
@@ -246,6 +248,8 @@ class DBOSContext:
         cm.__enter__()
 
     def _end_span(self, exc_value: Optional[BaseException]) -> None:
+        if dbos_tracer.disable_otlp:
+            return
         context_span = self.context_spans.pop()
         if exc_value is None:
             context_span.span.set_status(Status(StatusCode.OK))

dbos/_core.py

@@ -356,6 +356,7 @@ def _get_wf_invoke_func(
             )
             return recorded_result
         try:
+            dbos._active_workflows_set.add(status["workflow_uuid"])
             output = func()
             if not dbos.debug_mode:
                 dbos._sys_db.update_workflow_outcome(
@@ -378,6 +379,8 @@ def _get_wf_invoke_func(
                     error=_serialization.serialize_exception(error),
                 )
             raise
+        finally:
+            dbos._active_workflows_set.discard(status["workflow_uuid"])
 
     return persist
 
@@ -947,18 +950,14 @@ def decorate_transaction(
                                 assert (
                                     ctx.sql_session is not None
                                 ), "Cannot find a database connection"
-                                ApplicationDatabase.record_transaction_output(
+                                dbos._app_db.record_transaction_output(
                                     ctx.sql_session, txn_output
                                 )
                                 break
                         except DBAPIError as dbapi_error:
-                            driver_error = cast(
-                                Optional[psycopg.OperationalError], dbapi_error.orig
-                            )
-                            if retriable_postgres_exception(dbapi_error) or (
-                                driver_error is not None
-                                and driver_error.sqlstate == "40001"
-                            ):
+                            if retriable_postgres_exception(
+                                dbapi_error
+                            ) or dbos._app_db._is_serialization_error(dbapi_error):
                                 # Retry on serialization failure
                                 span = ctx.get_current_span()
                                 if span:

dbos/_dbos.py

@@ -32,7 +32,7 @@ from opentelemetry.trace import Span
 from rich import print
 
 from dbos._conductor.conductor import ConductorWebsocket
-from dbos._sys_db import WorkflowStatus
+from dbos._sys_db import SystemDatabase, WorkflowStatus
 from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
 from dbos._workflow_commands import fork_workflow, list_queued_workflows, list_workflows
 
@@ -70,7 +70,6 @@ from ._sys_db import (
     SystemDatabase,
     WorkflowStatus,
     _dbos_stream_closed_sentinel,
-    reset_system_database,
     workflow_is_active,
 )
 from ._tracer import DBOSTracer, dbos_tracer
@@ -80,7 +79,6 @@ if TYPE_CHECKING:
     from ._kafka import _KafkaConsumerWorkflow
     from flask import Flask
 
-from sqlalchemy import make_url
 from sqlalchemy.orm import Session
 
 if sys.version_info < (3, 10):
@@ -293,16 +291,24 @@ class DBOS:
         return _dbos_global_instance
 
     @classmethod
-    def destroy(cls, *, destroy_registry: bool = False) -> None:
+    def destroy(
+        cls,
+        *,
+        destroy_registry: bool = False,
+        workflow_completion_timeout_sec: int = 0,
+    ) -> None:
         global _dbos_global_instance
         if _dbos_global_instance is not None:
-            _dbos_global_instance._destroy()
+            _dbos_global_instance._destroy(
+                workflow_completion_timeout_sec=workflow_completion_timeout_sec,
+            )
         _dbos_global_instance = None
         if destroy_registry:
             global _dbos_global_registry
             _dbos_global_registry = None
         GlobalParams.app_version = os.environ.get("DBOS__APPVERSION", "")
         GlobalParams.executor_id = os.environ.get("DBOS__VMID", "local")
+        dbos_logger.info("DBOS successfully shut down")
 
     def __init__(
         self,
@@ -337,6 +343,7 @@ class DBOS:
         self.conductor_key: Optional[str] = conductor_key
         self.conductor_websocket: Optional[ConductorWebsocket] = None
         self._background_event_loop: BackgroundEventLoop = BackgroundEventLoop()
+        self._active_workflows_set: set[str] = set()
 
         # Globally set the application version and executor ID.
         # In DBOS Cloud, instead use the values supplied through environment variables.
@@ -448,13 +455,13 @@ class DBOS:
             self._background_event_loop.start()
             assert self._config["database_url"] is not None
             assert self._config["database"]["sys_db_engine_kwargs"] is not None
-            self._sys_db_field = SystemDatabase(
+            self._sys_db_field = SystemDatabase.create(
                 system_database_url=get_system_database_url(self._config),
                 engine_kwargs=self._config["database"]["sys_db_engine_kwargs"],
                 debug_mode=debug_mode,
             )
             assert self._config["database"]["db_engine_kwargs"] is not None
-            self._app_db_field = ApplicationDatabase(
+            self._app_db_field = ApplicationDatabase.create(
                 database_url=self._config["database_url"],
                 engine_kwargs=self._config["database"]["db_engine_kwargs"],
                 debug_mode=debug_mode,
@@ -580,20 +587,25 @@ class DBOS:
             not self._launched
         ), "The system database cannot be reset after DBOS is launched. Resetting the system database is a destructive operation that should only be used in a test environment."
 
-        sysdb_name = self._config["database"]["sys_db_name"]
-        assert sysdb_name is not None
+        SystemDatabase.reset_system_database(get_system_database_url(self._config))
 
-        assert self._config["database_url"] is not None
-        pg_db_url = make_url(self._config["database_url"]).set(database="postgres")
-
-        reset_system_database(pg_db_url, sysdb_name)
-
-    def _destroy(self) -> None:
+    def _destroy(self, *, workflow_completion_timeout_sec: int) -> None:
         self._initialized = False
         for event in self.poller_stop_events:
             event.set()
         for event in self.background_thread_stop_events:
             event.set()
+        if workflow_completion_timeout_sec > 0:
+            deadline = time.time() + workflow_completion_timeout_sec
+            while time.time() < deadline:
+                time.sleep(1)
+                active_workflows = len(self._active_workflows_set)
+                if active_workflows > 0:
+                    dbos_logger.info(
+                        f"Attempting to shut down DBOS. {active_workflows} workflows remain active. IDs: {self._active_workflows_set}"
+                    )
+                else:
+                    break
         self._background_event_loop.stop()
         if self._sys_db_field is not None:
             self._sys_db_field.destroy()
@@ -609,10 +621,8 @@ class DBOS:
             and self.conductor_websocket.websocket is not None
         ):
             self.conductor_websocket.websocket.close()
-        # CB - This needs work, some things ought to stop before DBs are tossed out,
-        #  on the other hand it hangs to move it
         if self._executor_field is not None:
-            self._executor_field.shutdown(cancel_futures=True)
+            self._executor_field.shutdown(wait=False, cancel_futures=True)
             self._executor_field = None
         for bg_thread in self._background_threads:
             bg_thread.join()