dbos 1.15.0a5__py3-none-any.whl → 1.15.0a6__py3-none-any.whl

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/_dbos.py

@@ -449,6 +449,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/_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,7 @@ 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.
     """
 
     name: str
@@ -52,6 +54,7 @@ class DBOSConfig(TypedDict, total=False):
     executor_id: Optional[str]
     dbos_system_schema: Optional[str]
     enable_otlp: Optional[bool]
+    system_database_engine: Optional[sa.Engine]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -97,20 +100,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 +110,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 +179,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/_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/_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/_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/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.dist-info → dbos-1.15.0a6.dist-info}/METADATA

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

{dbos-1.15.0a5.dist-info → dbos-1.15.0a6.dist-info}/RECORD

@@ -1,20 +1,20 @@
-dbos-1.15.0a5.dist-info/METADATA,sha256=ySISelF_uzHOl3jJk5bLWR6-58WI00Ys5TdK9cEhiFk,13021
-dbos-1.15.0a5.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-dbos-1.15.0a5.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-1.15.0a5.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-1.15.0a6.dist-info/METADATA,sha256=zDqBOU1RunuC6aqSUT4vpD2X4xhrqglKUZo67RNlGTw,13021
+dbos-1.15.0a6.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+dbos-1.15.0a6.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-1.15.0a6.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=pT4BuNLDCrIQX27vQG8NlfxX6PZRU7r9miq4thJTszU,982
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=hubQJw5T8zGKCPNS6FQTXy8jQ8GTJxoYQaDTMlICl9k,16267
 dbos/_app_db.py,sha256=WJwUdKsTpSZPCIWVeSF5FQNf5y1PF_lJ96tiaCjvck8,16385
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
-dbos/_client.py,sha256=1M2PhMNodw7Cnfrs3D0_NbwY0VnDoA1OnAOkBnLzTGg,18805
+dbos/_client.py,sha256=ppkO3bJ_qpBzWjYf9BkGfy8yF76e91UuhoLMtbIwXnU,18933
 dbos/_conductor/conductor.py,sha256=3E_hL3c9g9yWqKZkvI6KA0-ZzPMPRo06TOzT1esMiek,24114
 dbos/_conductor/protocol.py,sha256=q3rgLxINFtWFigdOONc-4gX4vn66UmMlJQD6Kj8LnL4,7420
 dbos/_context.py,sha256=cJDxVbswTLXKE5MV4Hmg6gpIX3Dd5mBTG-4lmofWP9E,27668
 dbos/_core.py,sha256=13DNN_fpSIs42NquV80XsHV7yKwY_adKP03h_xhXok4,50493
 dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
-dbos/_dbos.py,sha256=wT4POKmTKwPYrapMiV7EhfWicFdT9qVOiw3bujcE0Lg,57261
-dbos/_dbos_config.py,sha256=ha0Qo4kRCCi1dQZZzm1e24iMqR2Zty9LejcLxUbzlUg,25286
+dbos/_dbos.py,sha256=X2HNrGnq5rVgK5WnsSSl8cwpWT4m4Y0C8e7B1MTGsNg,57324
+dbos/_dbos_config.py,sha256=JHogT0oUCxVJgd0TaQ_nQM1xDm4C1jOdIbSDou7q0P4,25004
 dbos/_debouncer.py,sha256=9-9dlXKLRHSUSylprCw18r-7MAdFwD8-w0KkY-bEG_I,15281
 dbos/_debug.py,sha256=0MfgNqutCUhI4PEmmra9x7f3DiFE_0nscfUCHdLimEY,1415
 dbos/_docker_pg_helper.py,sha256=xySum4hTA8TVMBODoG19u4cXQAB1vCock-jwM2pnmSI,7791
@@ -36,9 +36,9 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=SypAS9l9EsaBHFn9FR8jmnqt01M74d9AF1AMa4m2hhI,1040
 dbos/_schemas/system_database.py,sha256=aEkjRQDh9xjdke0d9uFx_20-c9UjQtvuLtHZ24aOypA,5497
 dbos/_serialization.py,sha256=GLgWLtHpvk7nSHyXukVQLE1ASNA3CJBtfF8w6iflBDw,3590
-dbos/_sys_db.py,sha256=Trlbf99KvH8YRthQ3diJQMSi0eDmky7X8dPG1wuv5VI,83098
-dbos/_sys_db_postgres.py,sha256=V3RGj2cXdCZ8xNxiMO-ECfKuHXtGZLVVNtOXtFSVHBw,7215
-dbos/_sys_db_sqlite.py,sha256=lEmJjyUHXo09q1sMpNpasurh4JyrH6DsmYzLjDC5k6Y,7091
+dbos/_sys_db.py,sha256=rIbwxIc4uglpgsg5ZQRmfHv7YPfFQKUCDOW07pGrD8w,83699
+dbos/_sys_db_postgres.py,sha256=GuyGVyZZD_Wl7LjRSkHnOuZ-hOROlO4Xs2UeDhKq10E,6963
+dbos/_sys_db_sqlite.py,sha256=ifjKdy-Z9vlVIBf5L6XnSaNjiBdvqPE73asVHim4A5Q,6998
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py.dbos,sha256=aQnBPSSQpkB8ERfhf7gB7P9tsU6OPKhZscfeh0yiaD8,2702
@@ -52,8 +52,8 @@ dbos/_workflow_commands.py,sha256=k-i1bCfNrux43BHLT8wQ-l-MVZX3D6LGZLH7-uuiDRo,49
 dbos/cli/_github_init.py,sha256=R_94Fnn40CAmPy-zM00lwHi0ndyfv57TmIooADjmag4,3378
 dbos/cli/_template_init.py,sha256=AltKk256VocgvxLpuTxpjJyACrdHFjbGoqYhHzeLae4,2649
 dbos/cli/cli.py,sha256=s-gGQvHvVPeQp68raQElGnbBlSCv69JZ3HNFj5Qt2bs,27686
-dbos/cli/migration.py,sha256=1Y52EMc2rX7PJgJa3_KXV7oiQEO569k7aHXRofwYipo,3608
+dbos/cli/migration.py,sha256=zJnDPUBnil5XZXFxc01EbZ0Radw_y8wtDGZExgelAxc,3633
 dbos/dbos-config.schema.json,sha256=47wofTZ5jlFynec7bG0L369tAXbRQQ2euBxBXvg4m9c,1730
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-1.15.0a5.dist-info/RECORD,,
+dbos-1.15.0a6.dist-info/RECORD,,