dbos 1.15.0a5__tar.gz → 1.15.0a7__tar.gz

{dbos-1.15.0a5 → dbos-1.15.0a7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 1.15.0a5
+Version: 1.15.0a7
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_client.py

@@ -14,6 +14,8 @@ from typing import (
     Union,
 )
 
+import sqlalchemy as sa
+
 from dbos import _serialization
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
@@ -122,6 +124,7 @@ class DBOSClient:
         database_url: Optional[str] = None,  # DEPRECATED
         *,
         system_database_url: Optional[str] = None,
+        system_database_engine: Optional[sa.Engine] = None,
         application_database_url: Optional[str] = None,
         dbos_system_schema: Optional[str] = "dbos",
     ):
@@ -145,6 +148,7 @@ class DBOSClient:
                 "max_overflow": 0,
                 "pool_size": 2,
             },
+            engine=system_database_engine,
             schema=dbos_system_schema,
         )
         self._sys_db.check_connection()

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_dbos.py

@@ -335,6 +335,8 @@ class DBOS:
         self._background_threads: List[threading.Thread] = []
         self.conductor_url: Optional[str] = conductor_url
         self.conductor_key: Optional[str] = conductor_key
+        if config.get("conductor_key"):
+            self.conductor_key = config.get("conductor_key")
         self.conductor_websocket: Optional[ConductorWebsocket] = None
         self._background_event_loop: BackgroundEventLoop = BackgroundEventLoop()
         self._active_workflows_set: set[str] = set()
@@ -449,6 +451,7 @@ class DBOS:
             self._sys_db_field = SystemDatabase.create(
                 system_database_url=get_system_database_url(self._config),
                 engine_kwargs=self._config["database"]["sys_db_engine_kwargs"],
+                engine=self._config["system_database_engine"],
                 debug_mode=debug_mode,
                 schema=schema,
             )

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_dbos_config.py

@@ -3,6 +3,7 @@ import re
 from importlib import resources
 from typing import Any, Dict, List, Optional, TypedDict, cast
 
+import sqlalchemy as sa
 import yaml
 from sqlalchemy import make_url
 
@@ -34,6 +35,8 @@ class DBOSConfig(TypedDict, total=False):
         executor_id (str): Executor ID, used to identify the application instance in distributed environments
         dbos_system_schema (str): Schema name for DBOS system tables. Defaults to "dbos".
         enable_otlp (bool): If True, enable built-in DBOS OTLP tracing and logging.
+        system_database_engine (sa.Engine): A custom system database engine. If provided, DBOS will not create an engine but use this instead.
+        conductor_key (str): An API key for DBOS Conductor. Pass this in to connect your process to Conductor.
     """
 
     name: str
@@ -52,6 +55,8 @@ class DBOSConfig(TypedDict, total=False):
     executor_id: Optional[str]
     dbos_system_schema: Optional[str]
     enable_otlp: Optional[bool]
+    system_database_engine: Optional[sa.Engine]
+    conductor_key: Optional[str]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -97,20 +102,7 @@ class TelemetryConfig(TypedDict, total=False):
 class ConfigFile(TypedDict, total=False):
     """
     Data structure containing the DBOS Configuration.
-
-    This configuration data is typically loaded from `dbos-config.yaml`.
-    See `https://docs.dbos.dev/python/reference/configuration#dbos-configuration-file`
-
-    Attributes:
-        name (str): Application name
-        runtimeConfig (RuntimeConfig): Configuration for DBOS Cloud
-        database (DatabaseConfig): Configure pool sizes, migrate commands
-        database_url (str): Application database URL
-        system_database_url (str): System database URL
-        telemetry (TelemetryConfig): Configuration for tracing / logging
-        env (Dict[str,str]): Environment variables
-        dbos_system_schema (str): Schema name for DBOS system tables. Defaults to "dbos".
-
+    The DBOSConfig object is parsed into this.
     """
 
     name: str
@@ -120,6 +112,7 @@ class ConfigFile(TypedDict, total=False):
     system_database_url: Optional[str]
     telemetry: Optional[TelemetryConfig]
     env: Dict[str, str]
+    system_database_engine: Optional[sa.Engine]
     dbos_system_schema: Optional[str]
 
 
@@ -188,6 +181,8 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
     if telemetry:
         translated_config["telemetry"] = telemetry
 
+    translated_config["system_database_engine"] = config.get("system_database_engine")
+
     return translated_config
 
 

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_sys_db.py

@@ -346,17 +346,32 @@ class SystemDatabase(ABC):
         *,
         system_database_url: str,
         engine_kwargs: Dict[str, Any],
+        engine: Optional[sa.Engine],
+        schema: Optional[str],
         debug_mode: bool = False,
     ):
         import sqlalchemy.dialects.postgresql as pg
         import sqlalchemy.dialects.sqlite as sq
 
         self.dialect = sq if system_database_url.startswith("sqlite") else pg
-        self.engine = self._create_engine(system_database_url, engine_kwargs)
+
+        if system_database_url.startswith("sqlite"):
+            self.schema = None
+        else:
+            self.schema = schema if schema else "dbos"
+        SystemSchema.set_schema(self.schema)
+
+        if engine:
+            self.engine = engine
+            self.created_engine = False
+        else:
+            self.engine = self._create_engine(system_database_url, engine_kwargs)
+            self.created_engine = True
         self._engine_kwargs = engine_kwargs
 
         self.notifications_map = ThreadSafeConditionDict()
         self.workflow_events_map = ThreadSafeConditionDict()
+        self._listener_thread_lock = threading.Lock()
 
         # Now we can run background processes
         self._run_background_processes = True
@@ -1443,6 +1458,7 @@ class SystemDatabase(ABC):
     def create(
         system_database_url: str,
         engine_kwargs: Dict[str, Any],
+        engine: Optional[sa.Engine],
         schema: Optional[str],
         debug_mode: bool = False,
     ) -> "SystemDatabase":
@@ -1453,6 +1469,8 @@ class SystemDatabase(ABC):
             return SQLiteSystemDatabase(
                 system_database_url=system_database_url,
                 engine_kwargs=engine_kwargs,
+                engine=engine,
+                schema=schema,
                 debug_mode=debug_mode,
             )
         else:
@@ -1461,8 +1479,9 @@ class SystemDatabase(ABC):
             return PostgresSystemDatabase(
                 system_database_url=system_database_url,
                 engine_kwargs=engine_kwargs,
-                debug_mode=debug_mode,
+                engine=engine,
                 schema=schema,
+                debug_mode=debug_mode,
             )
 
     @db_retry()

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_sys_db_postgres.py

@@ -1,12 +1,11 @@
 import time
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, cast
 
 import psycopg
 import sqlalchemy as sa
 from sqlalchemy.exc import DBAPIError
 
 from dbos._migration import ensure_dbos_schema, run_dbos_migrations
-from dbos._schemas.system_database import SystemSchema
 
 from ._logger import dbos_logger
 from ._sys_db import SystemDatabase
@@ -15,25 +14,7 @@ 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],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ):
-        super().__init__(
-            system_database_url=system_database_url,
-            engine_kwargs=engine_kwargs,
-            debug_mode=debug_mode,
-        )
-        if schema is None:
-            self.schema = "dbos"
-        else:
-            self.schema = schema
-        SystemSchema.set_schema(self.schema)
-        self.notification_conn: Optional[psycopg.connection.Connection] = None
+    notification_conn: Optional[sa.PoolProxiedConnection] = None
 
     def _create_engine(
         self, system_database_url: str, engine_kwargs: Dict[str, Any]
@@ -48,27 +29,35 @@ class PostgresSystemDatabase(SystemDatabase):
             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()
+        # Unless we were provided an engine, if the system database does not already exist, create it
+        if self.created_engine:
+            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()
+        else:
+            # If we were provided an engine, validate it can connect
+            with self.engine.connect() as conn:
+                conn.execute(sa.text("SELECT 1"))
 
+        assert self.schema
         ensure_dbos_schema(self.engine, self.schema)
         run_dbos_migrations(self.engine, self.schema)
 
     def _cleanup_connections(self) -> None:
         """Clean up PostgreSQL-specific connections."""
-        if self.notification_conn is not None:
-            self.notification_conn.close()
+        with self._listener_thread_lock:
+            if self.notification_conn and self.notification_conn.dbapi_connection:
+                self.notification_conn.dbapi_connection.close()
+                self.notification_conn.invalidate()
 
     def _is_unique_constraint_violation(self, dbapi_error: DBAPIError) -> bool:
         """Check if the error is a unique constraint violation in PostgreSQL."""
@@ -111,20 +100,18 @@ class PostgresSystemDatabase(SystemDatabase):
         """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")
-
+                with self._listener_thread_lock:
+                    self.notification_conn = self.engine.raw_connection()
+                    self.notification_conn.detach()
+                    psycopg_conn = cast(
+                        psycopg.connection.Connection, self.notification_conn
+                    )
+                    psycopg_conn.set_autocommit(True)
+
+                    psycopg_conn.execute("LISTEN dbos_notifications_channel")
+                    psycopg_conn.execute("LISTEN dbos_workflow_events_channel")
                 while self._run_background_processes:
-                    gen = self.notification_conn.notifies()
+                    gen = psycopg_conn.notifies()
                     for notify in gen:
                         channel = notify.channel
                         dbos_logger.debug(
@@ -162,5 +149,4 @@ class PostgresSystemDatabase(SystemDatabase):
                     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()
+                self._cleanup_connections()

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/_sys_db_sqlite.py

@@ -6,7 +6,6 @@ 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
@@ -19,7 +18,6 @@ class SQLiteSystemDatabase(SystemDatabase):
         self, system_database_url: str, engine_kwargs: Dict[str, Any]
     ) -> sa.Engine:
         """Create a SQLite engine."""
-        SystemSchema.set_schema(None)
         return sa.create_engine(system_database_url)
 
     def run_migrations(self) -> None:

{dbos-1.15.0a5 → dbos-1.15.0a7}/dbos/cli/migration.py

@@ -20,6 +20,7 @@ def migrate_dbos_databases(
                 "max_overflow": 0,
                 "pool_size": 2,
             },
+            engine=None,
             schema=schema,
         )
         sys_db.run_migrations()

{dbos-1.15.0a5 → dbos-1.15.0a7}/pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
 ]
 requires-python = ">=3.10"
 readme = "README.md"
-version = "1.15.0a5"
+version = "1.15.0a7"
 
 [project.license]
 text = "MIT"

{dbos-1.15.0a5 → dbos-1.15.0a7}/tests/conftest.py

@@ -52,7 +52,7 @@ def default_config() -> DBOSConfig:
         "system_database_url": (
             "sqlite:///test.sqlite"
             if using_sqlite()
-            else f"postgresql://postgres:{quote(os.environ.get('PGPASSWORD', 'dbos'), safe='')}@localhost:5432/dbostestpy_dbos_sys"
+            else f"postgresql+psycopg://postgres:{quote(os.environ.get('PGPASSWORD', 'dbos'), safe='')}@localhost:5432/dbostestpy_dbos_sys"
         ),
         "enable_otlp": True,
     }

{dbos-1.15.0a5 → dbos-1.15.0a7}/tests/test_dbos.py

@@ -11,6 +11,7 @@ from typing import Any, Optional
 
 import pytest
 import sqlalchemy as sa
+from sqlalchemy.exc import OperationalError
 
 # Public API
 from dbos import (
@@ -1827,3 +1828,65 @@ def test_custom_schema(
     steps = client.list_workflow_steps(handle.workflow_id)
     assert len(steps) == 4
     assert "transaction" in steps[0]["function_name"]
+
+
+def test_custom_engine(
+    config: DBOSConfig,
+    cleanup_test_databases: None,
+    db_engine: sa.Engine,
+    skip_with_sqlite: None,
+) -> None:
+    DBOS.destroy(destroy_registry=True)
+    assert config["system_database_url"]
+    config["application_database_url"] = None
+    system_database_url = config["system_database_url"]
+
+    # Create a custom engine
+    engine = sa.create_engine(system_database_url)
+    config["system_database_engine"] = engine
+
+    # Launch DBOS with the engine. It should fail because the database does not exist.
+    dbos = DBOS(config=config)
+    with pytest.raises(OperationalError):
+        DBOS.launch()
+    DBOS.destroy(destroy_registry=True)
+
+    # Create the database
+    with db_engine.connect() as c:
+        c.execution_options(isolation_level="AUTOCOMMIT")
+        sysdb_name = sa.make_url(config["system_database_url"]).database
+        c.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
+
+    # Launch DBOS again using the custom pool. It should succeed despite the bogus URL.
+    config["system_database_url"] = "postgresql://bogus:url@not:42/fake"
+    dbos = DBOS(config=config)
+    DBOS.launch()
+
+    key = "key"
+    val = "val"
+
+    @DBOS.workflow()
+    def recv_workflow() -> Any:
+        DBOS.set_event(key, val)
+        return DBOS.recv()
+
+    assert dbos._sys_db.engine == engine
+    handle = DBOS.start_workflow(recv_workflow)
+    assert DBOS.get_event(handle.workflow_id, key) == val
+    DBOS.send(handle.workflow_id, val)
+    assert handle.get_result() == val
+    assert len(DBOS.list_workflows()) == 2
+    steps = DBOS.list_workflow_steps(handle.workflow_id)
+    assert len(steps) == 3
+    assert "setEvent" in steps[0]["function_name"]
+    DBOS.destroy(destroy_registry=True)
+
+    # Test custom engine with client
+    client = DBOSClient(
+        system_database_url=config["system_database_url"],
+        system_database_engine=config["system_database_engine"],
+    )
+    assert len(client.list_workflows()) == 2
+    steps = client.list_workflow_steps(handle.workflow_id)
+    assert len(steps) == 3
+    assert "setEvent" in steps[0]["function_name"]

{dbos-1.15.0a5 → dbos-1.15.0a7}/tests/test_failures.py

@@ -136,12 +136,10 @@ def test_notification_errors(dbos: DBOS, skip_with_sqlite: None) -> None:
     system_database = cast(PostgresSystemDatabase, dbos._sys_db)
     while system_database.notification_conn is None:
         time.sleep(1)
-    system_database.notification_conn.close()
-    assert system_database.notification_conn.closed == 1
+    system_database._cleanup_connections()
 
-    # Wait for the connection to be re-established
-    while system_database.notification_conn.closed != 0:
-        time.sleep(1)
+    # Wait for the connection to re-establish
+    time.sleep(3)
 
     dest_uuid = str("sruuid1")
     with SetWorkflowID(dest_uuid):

{dbos-1.15.0a5 → dbos-1.15.0a7}/tests/test_schema_migration.py

@@ -132,8 +132,9 @@ def test_sqlite_systemdb_migration() -> None:
         sys_db = SystemDatabase.create(
             system_database_url=sqlite_url,
             engine_kwargs={},
-            debug_mode=False,
+            engine=None,
             schema=None,
+            debug_mode=False,
         )
 
         # Run migrations