dbos 1.12.0a2__py3-none-any.whl → 1.13.0__py3-none-any.whl

dbos/_sys_db_postgres.py

@@ -0,0 +1,173 @@
+import time
+from typing import Any, Dict, Optional
+
+import psycopg
+import sqlalchemy as sa
+from sqlalchemy.exc import DBAPIError
+
+from dbos._migration import (
+    ensure_dbos_schema,
+    run_alembic_migrations,
+    run_dbos_migrations,
+)
+from dbos._schemas.system_database import SystemSchema
+
+from ._logger import dbos_logger
+from ._sys_db import SystemDatabase
+
+
+class PostgresSystemDatabase(SystemDatabase):
+    """PostgreSQL-specific implementation of SystemDatabase."""
+
+    def __init__(
+        self,
+        *,
+        system_database_url: str,
+        engine_kwargs: Dict[str, Any],
+        debug_mode: bool = False,
+    ):
+        super().__init__(
+            system_database_url=system_database_url,
+            engine_kwargs=engine_kwargs,
+            debug_mode=debug_mode,
+        )
+        self.notification_conn: Optional[psycopg.connection.Connection] = None
+
+    def _create_engine(
+        self, system_database_url: str, engine_kwargs: Dict[str, Any]
+    ) -> sa.Engine:
+        # TODO: Make the schema dynamic so this isn't needed
+        SystemSchema.workflow_status.schema = "dbos"
+        SystemSchema.operation_outputs.schema = "dbos"
+        SystemSchema.notifications.schema = "dbos"
+        SystemSchema.workflow_events.schema = "dbos"
+        SystemSchema.streams.schema = "dbos"
+        url = sa.make_url(system_database_url).set(drivername="postgresql+psycopg")
+        return sa.create_engine(url, **engine_kwargs)
+
+    def run_migrations(self) -> None:
+        """Run PostgreSQL-specific migrations."""
+        if self._debug_mode:
+            dbos_logger.warning("System database migrations are skipped in debug mode.")
+            return
+        system_db_url = self.engine.url
+        sysdb_name = system_db_url.database
+        # If the system database does not already exist, create it
+        engine = sa.create_engine(
+            system_db_url.set(database="postgres"), **self._engine_kwargs
+        )
+        with 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": sysdb_name},
+            ).scalar():
+                dbos_logger.info(f"Creating system database {sysdb_name}")
+                conn.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
+        engine.dispose()
+
+        using_dbos_migrations = ensure_dbos_schema(self.engine)
+        if not using_dbos_migrations:
+            # Complete the Alembic migrations, create the dbos_migrations table
+            run_alembic_migrations(self.engine)
+        run_dbos_migrations(self.engine)
+
+    def _cleanup_connections(self) -> None:
+        """Clean up PostgreSQL-specific connections."""
+        if self.notification_conn is not None:
+            self.notification_conn.close()
+
+    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_foreign_key_violation(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a foreign key violation in PostgreSQL."""
+        return dbapi_error.orig.sqlstate == "23503"  # type: ignore
+
+    @staticmethod
+    def _reset_system_database(database_url: str) -> None:
+        """Reset the PostgreSQL system database by dropping it."""
+        system_db_url = sa.make_url(database_url)
+        sysdb_name = system_db_url.database
+
+        if sysdb_name is None:
+            raise ValueError(f"System database name not found in URL {system_db_url}")
+
+        try:
+            # Connect to postgres default database
+            engine = sa.create_engine(
+                system_db_url.set(database="postgres", drivername="postgresql+psycopg"),
+                connect_args={"connect_timeout": 10},
+            )
+
+            with engine.connect() as conn:
+                # Set autocommit required for database dropping
+                conn.execution_options(isolation_level="AUTOCOMMIT")
+
+                # Drop the database
+                conn.execute(
+                    sa.text(f"DROP DATABASE IF EXISTS {sysdb_name} WITH (FORCE)")
+                )
+            engine.dispose()
+        except Exception as e:
+            dbos_logger.error(f"Error resetting PostgreSQL system database: {str(e)}")
+            raise e
+
+    def _notification_listener(self) -> None:
+        """Listen for PostgreSQL notifications using psycopg."""
+        while self._run_background_processes:
+            try:
+                # since we're using the psycopg connection directly, we need a url without the "+psycopg" suffix
+                url = sa.URL.create(
+                    "postgresql", **self.engine.url.translate_connect_args()
+                )
+                # Listen to notifications
+                self.notification_conn = psycopg.connect(
+                    url.render_as_string(hide_password=False), autocommit=True
+                )
+
+                self.notification_conn.execute("LISTEN dbos_notifications_channel")
+                self.notification_conn.execute("LISTEN dbos_workflow_events_channel")
+
+                while self._run_background_processes:
+                    gen = self.notification_conn.notifies()
+                    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:
+                                condition = self.notifications_map.get(notify.payload)
+                                if condition is None:
+                                    # No condition found for this payload
+                                    continue
+                                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:
+                                condition = self.workflow_events_map.get(notify.payload)
+                                if condition is None:
+                                    # No condition found for this payload
+                                    continue
+                                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.warning(f"Notification listener error: {e}")
+                    time.sleep(1)
+                    # Then the loop will try to reconnect and restart the listener
+            finally:
+                if self.notification_conn is not None:
+                    self.notification_conn.close()

dbos/_sys_db_sqlite.py

@@ -0,0 +1,182 @@
+import os
+import time
+from typing import Any, Dict, Optional, Tuple
+
+import sqlalchemy as sa
+from sqlalchemy.exc import DBAPIError
+
+from dbos._migration import sqlite_migrations
+from dbos._schemas.system_database import SystemSchema
+
+from ._logger import dbos_logger
+from ._sys_db import SystemDatabase
+
+
+class SQLiteSystemDatabase(SystemDatabase):
+    """SQLite-specific implementation of SystemDatabase."""
+
+    def _create_engine(
+        self, system_database_url: str, engine_kwargs: Dict[str, Any]
+    ) -> sa.Engine:
+        """Create a SQLite engine."""
+        # TODO: Make the schema dynamic so this isn't needed
+        SystemSchema.workflow_status.schema = None
+        SystemSchema.operation_outputs.schema = None
+        SystemSchema.notifications.schema = None
+        SystemSchema.workflow_events.schema = None
+        SystemSchema.streams.schema = None
+        return sa.create_engine(system_database_url)
+
+    def run_migrations(self) -> None:
+        """Run SQLite-specific migrations."""
+        if self._debug_mode:
+            dbos_logger.warning("System database migrations are skipped in debug mode.")
+            return
+
+        with self.engine.begin() as conn:
+            # Enable foreign keys for SQLite
+            conn.execute(sa.text("PRAGMA foreign_keys = ON"))
+
+            # Check if migrations table exists
+            result = conn.execute(
+                sa.text(
+                    "SELECT name FROM sqlite_master WHERE type='table' AND name='dbos_migrations'"
+                )
+            ).fetchone()
+
+            if result is None:
+                # Create migrations table
+                conn.execute(
+                    sa.text(
+                        "CREATE TABLE dbos_migrations (version INTEGER NOT NULL PRIMARY KEY)"
+                    )
+                )
+                last_applied = 0
+            else:
+                # Get current migration version
+                version_result = conn.execute(
+                    sa.text("SELECT version FROM dbos_migrations")
+                ).fetchone()
+                last_applied = version_result[0] if version_result else 0
+
+            # Apply migrations starting from the next version
+            for i, migration_sql in enumerate(sqlite_migrations, 1):
+                if i <= last_applied:
+                    continue
+
+                # Execute the migration
+                dbos_logger.info(
+                    f"Applying DBOS SQLite system database schema migration {i}"
+                )
+
+                # SQLite only allows one statement at a time, so split by semicolon
+                statements = [
+                    stmt.strip() for stmt in migration_sql.split(";") if stmt.strip()
+                ]
+                for statement in statements:
+                    conn.execute(sa.text(statement))
+
+                # Update the single row with the new version
+                if last_applied == 0:
+                    conn.execute(
+                        sa.text(
+                            "INSERT INTO dbos_migrations (version) VALUES (:version)"
+                        ),
+                        {"version": i},
+                    )
+                else:
+                    conn.execute(
+                        sa.text("UPDATE dbos_migrations SET version = :version"),
+                        {"version": i},
+                    )
+                last_applied = i
+
+    def _cleanup_connections(self) -> None:
+        # SQLite doesn't require special connection cleanup
+        pass
+
+    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_foreign_key_violation(self, dbapi_error: DBAPIError) -> bool:
+        """Check if the error is a foreign key violation in SQLite."""
+        return "FOREIGN KEY constraint failed" in str(dbapi_error.orig)
+
+    @staticmethod
+    def _reset_system_database(database_url: str) -> None:
+        """Reset the SQLite system database by deleting the database file."""
+
+        # Parse the SQLite database URL to get the file path
+        url = sa.make_url(database_url)
+        db_path = url.database
+
+        if db_path is None:
+            raise ValueError(f"System database path not found in URL {url}")
+
+        try:
+            if os.path.exists(db_path):
+                os.remove(db_path)
+                dbos_logger.info(f"Deleted SQLite database file: {db_path}")
+            else:
+                dbos_logger.info(f"SQLite database file does not exist: {db_path}")
+        except OSError as e:
+            dbos_logger.error(
+                f"Error deleting SQLite database file {db_path}: {str(e)}"
+            )
+            raise e
+
+    def _notification_listener(self) -> None:
+        """Poll for notifications and workflow events in SQLite."""
+
+        def split_payload(payload: str) -> Tuple[str, Optional[str]]:
+            """Split payload into components (first::second format)."""
+            if "::" in payload:
+                parts = payload.split("::", 1)
+                return parts[0], parts[1]
+            return payload, None
+
+        def signal_condition(condition_map: Any, payload: str) -> None:
+            """Signal a condition variable if it exists."""
+            condition = condition_map.get(payload)
+            if condition:
+                condition.acquire()
+                condition.notify_all()
+                condition.release()
+                dbos_logger.debug(f"Signaled condition for {payload}")
+
+        while self._run_background_processes:
+            try:
+                # Poll every second
+                time.sleep(1)
+
+                # Check all payloads in the notifications_map
+                for payload in list(self.notifications_map._dict.keys()):
+                    dest_uuid, topic = split_payload(payload)
+                    with self.engine.begin() as conn:
+                        result = conn.execute(
+                            sa.text(
+                                "SELECT 1 FROM notifications WHERE destination_uuid = :dest_uuid AND topic = :topic LIMIT 1"
+                            ),
+                            {"dest_uuid": dest_uuid, "topic": topic},
+                        )
+                        if result.fetchone():
+                            signal_condition(self.notifications_map, payload)
+
+                # Check all payloads in the workflow_events_map
+                for payload in list(self.workflow_events_map._dict.keys()):
+                    workflow_uuid, key = split_payload(payload)
+                    with self.engine.begin() as conn:
+                        result = conn.execute(
+                            sa.text(
+                                "SELECT 1 FROM workflow_events WHERE workflow_uuid = :workflow_uuid AND key = :key LIMIT 1"
+                            ),
+                            {"workflow_uuid": workflow_uuid, "key": key},
+                        )
+                        if result.fetchone():
+                            signal_condition(self.workflow_events_map, payload)
+
+            except Exception as e:
+                if self._run_background_processes:
+                    dbos_logger.warning(f"SQLite notification poller error: {e}")
+                    time.sleep(1)

dbos/_tracer.py

@@ -24,10 +24,14 @@ class DBOSTracer:
     def __init__(self) -> None:
         self.app_id = os.environ.get("DBOS__APPID", None)
         self.provider: Optional[TracerProvider] = None
+        self.disable_otlp: bool = False
 
     def config(self, config: ConfigFile) -> None:
         self.otlp_attributes = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
-        if not isinstance(trace.get_tracer_provider(), TracerProvider):
+        self.disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
+        if not self.disable_otlp and not isinstance(
+            trace.get_tracer_provider(), TracerProvider
+        ):
             resource = Resource(
                 attributes={
                     ResourceAttributes.SERVICE_NAME: config["name"],

dbos/_utils.py

@@ -20,7 +20,9 @@ class GlobalParams:
         dbos_version = "unknown"
 
 
-def retriable_postgres_exception(e: DBAPIError) -> bool:
+def retriable_postgres_exception(e: Exception) -> bool:
+    if not isinstance(e, DBAPIError):
+        return False
     if e.connection_invalidated:
         return True
     if isinstance(e.orig, psycopg.OperationalError):
@@ -48,3 +50,10 @@ def retriable_postgres_exception(e: DBAPIError) -> bool:
             return False
     else:
         return False
+
+
+def retriable_sqlite_exception(e: Exception) -> bool:
+    if "database is locked" in str(e):
+        return True
+    else:
+        return False