dbos 0.27.0a1__tar.gz → 0.27.0a2__tar.gz

{dbos-0.27.0a1 → dbos-0.27.0a2}/PKG-INFO

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

{dbos-0.27.0a1 → dbos-0.27.0a2}/dbos/_app_db.py

@@ -77,12 +77,24 @@ class ApplicationDatabase:
         pool_size = database.get("app_db_pool_size")
         if pool_size is None:
             pool_size = 20
+
+        engine_kwargs = database.get("db_engine_kwargs")
+        if engine_kwargs is None:
+            engine_kwargs = {}
+
+        # Respect user-provided values. Otherwise, set defaults.
+        if "pool_size" not in engine_kwargs:
+            engine_kwargs["pool_size"] = pool_size
+        if "max_overflow" not in engine_kwargs:
+            engine_kwargs["max_overflow"] = 0
+        if "pool_timeout" not in engine_kwargs:
+            engine_kwargs["pool_timeout"] = 30
+        if "connect_args" not in engine_kwargs:
+            engine_kwargs["connect_args"] = connect_args
+
         self.engine = sa.create_engine(
             app_db_url,
-            pool_size=pool_size,
-            max_overflow=0,
-            pool_timeout=30,
-            connect_args=connect_args,
+            **engine_kwargs,
         )
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode

{dbos-0.27.0a1 → dbos-0.27.0a2}/dbos/_dbos_config.py

@@ -31,6 +31,7 @@ class DBOSConfig(TypedDict, total=False):
         app_db_pool_size (int): Application database pool size
         sys_db_name (str): System database name
         sys_db_pool_size (int): System database pool size
+        db_engine_kwargs (Dict[str, Any]): SQLAlchemy engine kwargs (See https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine)
         log_level (str): Log level
         otlp_traces_endpoints: List[str]: OTLP traces endpoints
         otlp_logs_endpoints: List[str]: OTLP logs endpoints
@@ -43,6 +44,7 @@ class DBOSConfig(TypedDict, total=False):
     app_db_pool_size: Optional[int]
     sys_db_name: Optional[str]
     sys_db_pool_size: Optional[int]
+    db_engine_kwargs: Optional[Dict[str, Any]]
     log_level: Optional[str]
     otlp_traces_endpoints: Optional[List[str]]
     otlp_logs_endpoints: Optional[List[str]]
@@ -64,6 +66,7 @@ class DatabaseConfig(TypedDict, total=False):
         app_db_pool_size (int): Application database pool size
         sys_db_name (str): System database name
         sys_db_pool_size (int): System database pool size
+        db_engine_kwargs (Dict[str, Any]): SQLAlchemy engine kwargs
         migrate (List[str]): Migration commands to run on startup
     """
 
@@ -76,6 +79,7 @@ class DatabaseConfig(TypedDict, total=False):
     app_db_pool_size: Optional[int]
     sys_db_name: Optional[str]
     sys_db_pool_size: Optional[int]
+    db_engine_kwargs: Optional[Dict[str, Any]]
     ssl: Optional[bool]  # Will be removed in a future version
     ssl_ca: Optional[str]  # Will be removed in a future version
     migrate: Optional[List[str]]
@@ -183,6 +187,8 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
         db_config["app_db_pool_size"] = config.get("app_db_pool_size")
     if "sys_db_pool_size" in config:
         db_config["sys_db_pool_size"] = config.get("sys_db_pool_size")
+    if "db_engine_kwargs" in config:
+        db_config["db_engine_kwargs"] = config.get("db_engine_kwargs")
     if db_config:
         translated_config["database"] = db_config
 

{dbos-0.27.0a1 → dbos-0.27.0a2}/dbos/_sys_db.py

@@ -267,12 +267,23 @@ class SystemDatabase:
         if pool_size is None:
             pool_size = 20
 
+        engine_kwargs = database.get("db_engine_kwargs")
+        if engine_kwargs is None:
+            engine_kwargs = {}
+
+        # Respect user-provided values. Otherwise, set defaults.
+        if "pool_size" not in engine_kwargs:
+            engine_kwargs["pool_size"] = pool_size
+        if "max_overflow" not in engine_kwargs:
+            engine_kwargs["max_overflow"] = 0
+        if "pool_timeout" not in engine_kwargs:
+            engine_kwargs["pool_timeout"] = 30
+        if "connect_args" not in engine_kwargs:
+            engine_kwargs["connect_args"] = {"connect_timeout": 10}
+
         self.engine = sa.create_engine(
             system_db_url,
-            pool_size=pool_size,
-            max_overflow=0,
-            pool_timeout=30,
-            connect_args={"connect_timeout": 10},
+            **engine_kwargs,
         )
 
         # Run a schema migration for the system database
@@ -378,7 +389,6 @@ class SystemDatabase:
         cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.workflow_deadline_epoch_ms, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
         results = conn.execute(cmd)
-
         row = results.fetchone()
         if row is not None:
             # Check the started workflow matches the expected name, class_name, config_name, and queue_name

{dbos-0.27.0a1 → dbos-0.27.0a2}/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.27.0a1"
+version = "0.27.0a2"
 
 [project.license]
 text = "MIT"

{dbos-0.27.0a1 → dbos-0.27.0a2}/tests/test_config.py

@@ -1251,6 +1251,8 @@ def test_configured_pool_sizes():
     dbos.launch()
     assert dbos._app_db.engine.pool._pool.maxsize == 42
     assert dbos._sys_db.engine.pool._pool.maxsize == 43
+    assert dbos._app_db.engine.pool._pre_ping == False
+    assert dbos._sys_db.engine.pool._pre_ping == False
     dbos.destroy()
 
 
@@ -1326,3 +1328,21 @@ def test_get_dbos_database_url(mocker):
         database="some_db",
     ).render_as_string(hide_password=False)
     assert get_dbos_database_url() == expected_url
+
+
+def test_db_engine_kwargs():
+    config: DBOSConfig = {
+        "name": "test-app",
+        "db_engine_kwargs": {
+            "pool_pre_ping": True,
+        },
+    }
+
+    dbos = DBOS(config=config)
+    dbos.launch()
+    assert dbos._app_db.engine.pool._pre_ping == True
+    assert dbos._sys_db.engine.pool._pre_ping == True
+    # Default values should be set
+    assert dbos._app_db.engine.pool._pool.maxsize == 20
+    assert dbos._sys_db.engine.pool._pool.maxsize == 20
+    dbos.destroy()