dbos 1.15.0a1__tar.gz → 1.15.0a3__tar.gz

{dbos-1.15.0a1 → dbos-1.15.0a3}/PKG-INFO

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

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_app_db.py

@@ -241,6 +241,7 @@ class ApplicationDatabase(ABC):
     def create(
         database_url: str,
         engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
         debug_mode: bool = False,
     ) -> "ApplicationDatabase":
         """Factory method to create the appropriate ApplicationDatabase implementation based on URL."""
@@ -256,12 +257,32 @@ class ApplicationDatabase(ABC):
                 database_url=database_url,
                 engine_kwargs=engine_kwargs,
                 debug_mode=debug_mode,
+                schema=schema,
             )
 
 
 class PostgresApplicationDatabase(ApplicationDatabase):
     """PostgreSQL-specific implementation of ApplicationDatabase."""
 
+    def __init__(
+        self,
+        *,
+        database_url: str,
+        engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
+        debug_mode: bool = False,
+    ):
+        super().__init__(
+            database_url=database_url,
+            engine_kwargs=engine_kwargs,
+            debug_mode=debug_mode,
+        )
+        if schema is None:
+            self.schema = "dbos"
+        else:
+            self.schema = schema
+        ApplicationSchema.transaction_outputs.schema = schema
+
     def _create_engine(
         self, database_url: str, engine_kwargs: Dict[str, Any]
     ) -> sa.Engine:
@@ -271,9 +292,6 @@ class PostgresApplicationDatabase(ApplicationDatabase):
         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,
@@ -307,24 +325,18 @@ class PostgresApplicationDatabase(ApplicationDatabase):
                 sa.text(
                     "SELECT 1 FROM information_schema.schemata WHERE schema_name = :schema_name"
                 ),
-                parameters={"schema_name": ApplicationSchema.schema},
+                parameters={"schema_name": self.schema},
             ).scalar()
 
             if not schema_exists:
-                schema_creation_query = sa.text(
-                    f"CREATE SCHEMA {ApplicationSchema.schema}"
-                )
+                schema_creation_query = sa.text(f'CREATE SCHEMA "{self.schema}"')
                 conn.execute(schema_creation_query)
 
         inspector = inspect(self.engine)
-        if not inspector.has_table(
-            "transaction_outputs", schema=ApplicationSchema.schema
-        ):
+        if not inspector.has_table("transaction_outputs", schema=self.schema):
             ApplicationSchema.metadata_obj.create_all(self.engine)
         else:
-            columns = inspector.get_columns(
-                "transaction_outputs", schema=ApplicationSchema.schema
-            )
+            columns = inspector.get_columns("transaction_outputs", schema=self.schema)
             column_names = [col["name"] for col in columns]
 
             if "function_name" not in column_names:
@@ -333,7 +345,7 @@ class PostgresApplicationDatabase(ApplicationDatabase):
                     conn.execute(
                         text(
                             f"""
-                        ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
+                        ALTER TABLE \"{self.schema}\".transaction_outputs
                         ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
                         """
                         )

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_client.py

@@ -22,11 +22,7 @@ from dbos._sys_db import SystemDatabase
 if TYPE_CHECKING:
     from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
 
-from dbos._dbos_config import (
-    get_application_database_url,
-    get_system_database_url,
-    is_valid_database_url,
-)
+from dbos._dbos_config import get_system_database_url, is_valid_database_url
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
 from dbos._serialization import WorkflowInputs
@@ -118,21 +114,20 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
 
 
 class DBOSClient:
+
+    _app_db: ApplicationDatabase | None = None
+
     def __init__(
         self,
         database_url: Optional[str] = None,  # DEPRECATED
         *,
         system_database_url: Optional[str] = None,
         application_database_url: Optional[str] = None,
+        dbos_system_schema: Optional[str] = "dbos",
         system_database: Optional[str] = None,  # DEPRECATED
     ):
-        application_database_url = get_application_database_url(
-            {
-                "system_database_url": system_database_url,
-                "database_url": (
-                    database_url if database_url else application_database_url
-                ),
-            }
+        application_database_url = (
+            database_url if database_url else application_database_url
         )
         system_database_url = get_system_database_url(
             {
@@ -142,7 +137,8 @@ class DBOSClient:
             }
         )
         assert is_valid_database_url(system_database_url)
-        assert is_valid_database_url(application_database_url)
+        if application_database_url:
+            assert is_valid_database_url(application_database_url)
         # We only create database connections but do not run migrations
         self._sys_db = SystemDatabase.create(
             system_database_url=system_database_url,
@@ -151,16 +147,19 @@ class DBOSClient:
                 "max_overflow": 0,
                 "pool_size": 2,
             },
+            schema=dbos_system_schema,
         )
         self._sys_db.check_connection()
-        self._app_db = ApplicationDatabase.create(
-            database_url=application_database_url,
-            engine_kwargs={
-                "pool_timeout": 30,
-                "max_overflow": 0,
-                "pool_size": 2,
-            },
-        )
+        if application_database_url:
+            self._app_db = ApplicationDatabase.create(
+                database_url=application_database_url,
+                engine_kwargs={
+                    "pool_timeout": 30,
+                    "max_overflow": 0,
+                    "pool_size": 2,
+                },
+                schema=dbos_system_schema,
+            )
 
     def destroy(self) -> None:
         self._sys_db.destroy()

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_core.py

@@ -896,7 +896,9 @@ def decorate_transaction(
                 raise DBOSWorkflowCancelledError(
                     f"Workflow {ctx.workflow_id} is cancelled. Aborting transaction {transaction_name}."
                 )
-
+            assert (
+                dbos._app_db
+            ), "Transactions can only be used if DBOS is configured with an application_database_url"
             with dbos._app_db.sessionmaker() as session:
                 attributes: TracedAttributes = {
                     "name": transaction_name,

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_dbos.py

@@ -409,13 +409,8 @@ class DBOS:
         return rv
 
     @property
-    def _app_db(self) -> ApplicationDatabase:
-        if self._app_db_field is None:
-            raise DBOSException(
-                "Application database accessed before DBOS was launched"
-            )
-        rv: ApplicationDatabase = self._app_db_field
-        return rv
+    def _app_db(self) -> ApplicationDatabase | None:
+        return self._app_db_field
 
     @property
     def _admin_server(self) -> AdminServer:
@@ -448,26 +443,31 @@ class DBOS:
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=sys.maxsize)
             self._background_event_loop.start()
-            assert self._config["database_url"] is not None
             assert self._config["database"]["sys_db_engine_kwargs"] is not None
+            # Get the schema configuration, use "dbos" as default
+            schema = self._config.get("dbos_system_schema", "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"],
                 debug_mode=debug_mode,
+                schema=schema,
             )
             assert self._config["database"]["db_engine_kwargs"] is not None
-            self._app_db_field = ApplicationDatabase.create(
-                database_url=self._config["database_url"],
-                engine_kwargs=self._config["database"]["db_engine_kwargs"],
-                debug_mode=debug_mode,
-            )
+            if self._config["database_url"]:
+                self._app_db_field = ApplicationDatabase.create(
+                    database_url=self._config["database_url"],
+                    engine_kwargs=self._config["database"]["db_engine_kwargs"],
+                    debug_mode=debug_mode,
+                    schema=schema,
+                )
 
             if debug_mode:
                 return
 
             # Run migrations for the system and application databases
             self._sys_db.run_migrations()
-            self._app_db.run_migrations()
+            if self._app_db:
+                self._app_db.run_migrations()
 
             admin_port = self._config.get("runtimeConfig", {}).get("admin_port")
             if admin_port is None:

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_dbos_config.py

@@ -1,4 +1,3 @@
-import json
 import os
 import re
 from importlib import resources
@@ -34,6 +33,7 @@ class DBOSConfig(TypedDict, total=False):
         otlp_attributes (dict[str, str]): A set of custom attributes to apply OTLP-exported logs and traces
         application_version (str): Application version
         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.
     """
 
@@ -52,6 +52,7 @@ class DBOSConfig(TypedDict, total=False):
     otlp_attributes: Optional[dict[str, str]]
     application_version: Optional[str]
     executor_id: Optional[str]
+    dbos_system_schema: Optional[str]
     enable_otlp: Optional[bool]
 
 
@@ -70,6 +71,7 @@ class DatabaseConfig(TypedDict, total=False):
         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
+        dbos_system_schema (str): Schema name for DBOS system tables. Defaults to "dbos".
     """
 
     sys_db_name: Optional[str]
@@ -113,6 +115,7 @@ class ConfigFile(TypedDict, total=False):
         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".
 
     """
 
@@ -123,6 +126,7 @@ class ConfigFile(TypedDict, total=False):
     system_database_url: Optional[str]
     telemetry: Optional[TelemetryConfig]
     env: Dict[str, str]
+    dbos_system_schema: Optional[str]
 
 
 def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
@@ -153,6 +157,9 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
     if "system_database_url" in config:
         translated_config["system_database_url"] = config.get("system_database_url")
 
+    if "dbos_system_schema" in config:
+        translated_config["dbos_system_schema"] = config.get("dbos_system_schema")
+
     # Runtime config
     translated_config["runtimeConfig"] = {"run_admin_server": True}
     if "admin_port" in config:
@@ -408,22 +415,20 @@ def process_config(
                 url = url.set(database=f"{url.database}{SystemSchema.sysdb_suffix}")
             data["system_database_url"] = url.render_as_string(hide_password=False)
 
-    # If a system database URL is provided but not an application database URL, set the
-    # application database URL to the system database URL.
+    # If a system database URL is provided but not an application database URL,
+    # do not create an application database.
     if data.get("system_database_url") and not data.get("database_url"):
         assert data["system_database_url"]
-        data["database_url"] = data["system_database_url"]
+        data["database_url"] = None
 
-    # If neither URL is provided, use a default SQLite database URL.
+    # If neither URL is provided, use a default SQLite system database URL.
     if not data.get("database_url") and not data.get("system_database_url"):
         _app_db_name = _app_name_to_db_name(data["name"])
-        data["system_database_url"] = data["database_url"] = (
-            f"sqlite:///{_app_db_name}.sqlite"
-        )
+        data["system_database_url"] = f"sqlite:///{_app_db_name}.sqlite"
+        data["database_url"] = None
 
     configure_db_engine_parameters(data["database"], connect_timeout=connect_timeout)
 
-    assert data["database_url"] is not None
     assert data["system_database_url"] is not None
     # Pretty-print connection information, respecting log level
     if not silent and logs["logLevel"] == "INFO" or logs["logLevel"] == "DEBUG":
@@ -431,7 +436,12 @@ def process_config(
             hide_password=True
         )
         print(f"DBOS system database URL: {printable_sys_db_url}")
-        if data["database_url"].startswith("sqlite"):
+        if data["database_url"]:
+            printable_app_db_url = make_url(data["database_url"]).render_as_string(
+                hide_password=True
+            )
+            print(f"DBOS application database URL: {printable_app_db_url}")
+        if data["system_database_url"].startswith("sqlite"):
             print(
                 f"Using SQLite as a system database. The SQLite system database is for development and testing. PostgreSQL is recommended for production use."
             )
@@ -543,6 +553,8 @@ def overwrite_config(provided_config: ConfigFile) -> ConfigFile:
             "DBOS_SYSTEM_DATABASE_URL environment variable is not set. This is required to connect to the database."
         )
     provided_config["system_database_url"] = system_db_url
+    # Always use the "dbos" schema when deploying to DBOS Cloud
+    provided_config["dbos_system_schema"] = "dbos"
 
     # Telemetry config
     if "telemetry" not in provided_config or provided_config["telemetry"] is None:
@@ -615,12 +627,11 @@ def get_system_database_url(config: ConfigFile) -> str:
         )
 
 
-def get_application_database_url(config: ConfigFile) -> str:
+def get_application_database_url(config: ConfigFile) -> str | None:
     # For backwards compatibility, the application database URL is "database_url"
     if config.get("database_url"):
         assert config["database_url"]
         return config["database_url"]
     else:
-        # If the application database URL is not specified, set it to the system database URL
-        assert config["system_database_url"]
-        return config["system_database_url"]
+        # If the application database URL is not specified, return None
+        return None

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_migration.py

@@ -5,7 +5,7 @@ import sqlalchemy as sa
 from ._logger import dbos_logger
 
 
-def ensure_dbos_schema(engine: sa.Engine) -> None:
+def ensure_dbos_schema(engine: sa.Engine, schema: str) -> None:
     """
     True if using DBOS migrations (DBOS schema and migrations table already exist or were created)
     False if using Alembic migrations (DBOS schema exists, but dbos_migrations table doesn't)
@@ -14,41 +14,46 @@ def ensure_dbos_schema(engine: sa.Engine) -> None:
         # Check if dbos schema exists
         schema_result = conn.execute(
             sa.text(
-                "SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'dbos'"
-            )
+                "SELECT schema_name FROM information_schema.schemata WHERE schema_name = :schema"
+            ),
+            {"schema": schema},
         )
         schema_exists = schema_result.fetchone() is not None
 
         # Create schema if it doesn't exist
         if not schema_exists:
-            conn.execute(sa.text("CREATE SCHEMA dbos"))
+            conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
 
         # Check if dbos_migrations table exists
         table_result = conn.execute(
             sa.text(
-                "SELECT table_name FROM information_schema.tables WHERE table_schema = 'dbos' AND table_name = 'dbos_migrations'"
-            )
+                "SELECT table_name FROM information_schema.tables WHERE table_schema = :schema AND table_name = 'dbos_migrations'"
+            ),
+            {"schema": schema},
         )
         table_exists = table_result.fetchone() is not None
 
         if not table_exists:
             conn.execute(
                 sa.text(
-                    "CREATE TABLE dbos.dbos_migrations (version BIGINT NOT NULL PRIMARY KEY)"
+                    f'CREATE TABLE "{schema}".dbos_migrations (version BIGINT NOT NULL PRIMARY KEY)'
                 )
             )
 
 
-def run_dbos_migrations(engine: sa.Engine) -> None:
+def run_dbos_migrations(engine: sa.Engine, schema: str) -> None:
     """Run DBOS-managed migrations by executing each SQL command in dbos_migrations."""
     with engine.begin() as conn:
         # Get current migration version
-        result = conn.execute(sa.text("SELECT version FROM dbos.dbos_migrations"))
+        result = conn.execute(
+            sa.text(f'SELECT version FROM "{schema}".dbos_migrations')
+        )
         current_version = result.fetchone()
         last_applied = current_version[0] if current_version else 0
 
         # Apply migrations starting from the next version
-        for i, migration_sql in enumerate(dbos_migrations, 1):
+        migrations = get_dbos_migrations(schema)
+        for i, migration_sql in enumerate(migrations, 1):
             if i <= last_applied:
                 continue
 
@@ -60,23 +65,26 @@ def run_dbos_migrations(engine: sa.Engine) -> None:
             if last_applied == 0:
                 conn.execute(
                     sa.text(
-                        "INSERT INTO dbos.dbos_migrations (version) VALUES (:version)"
+                        f'INSERT INTO "{schema}".dbos_migrations (version) VALUES (:version)'
                     ),
                     {"version": i},
                 )
             else:
                 conn.execute(
-                    sa.text("UPDATE dbos.dbos_migrations SET version = :version"),
+                    sa.text(
+                        f'UPDATE "{schema}".dbos_migrations SET version = :version'
+                    ),
                     {"version": i},
                 )
             last_applied = i
 
 
-dbos_migration_one = """
+def get_dbos_migration_one(schema: str) -> str:
+    return f"""
 -- Enable uuid extension for generating UUIDs
 CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
 
-CREATE TABLE dbos.workflow_status (
+CREATE TABLE \"{schema}\".workflow_status (
     workflow_uuid TEXT PRIMARY KEY,
     status TEXT,
     name TEXT,
@@ -103,15 +111,15 @@ CREATE TABLE dbos.workflow_status (
     priority INTEGER NOT NULL DEFAULT 0
 );
 
-CREATE INDEX workflow_status_created_at_index ON dbos.workflow_status (created_at);
-CREATE INDEX workflow_status_executor_id_index ON dbos.workflow_status (executor_id);
-CREATE INDEX workflow_status_status_index ON dbos.workflow_status (status);
+CREATE INDEX workflow_status_created_at_index ON \"{schema}\".workflow_status (created_at);
+CREATE INDEX workflow_status_executor_id_index ON \"{schema}\".workflow_status (executor_id);
+CREATE INDEX workflow_status_status_index ON \"{schema}\".workflow_status (status);
 
-ALTER TABLE dbos.workflow_status 
+ALTER TABLE \"{schema}\".workflow_status 
 ADD CONSTRAINT uq_workflow_status_queue_name_dedup_id 
 UNIQUE (queue_name, deduplication_id);
 
-CREATE TABLE dbos.operation_outputs (
+CREATE TABLE \"{schema}\".operation_outputs (
     workflow_uuid TEXT NOT NULL,
     function_id INTEGER NOT NULL,
     function_name TEXT NOT NULL DEFAULT '',
@@ -119,23 +127,23 @@ CREATE TABLE dbos.operation_outputs (
     error TEXT,
     child_workflow_id TEXT,
     PRIMARY KEY (workflow_uuid, function_id),
-    FOREIGN KEY (workflow_uuid) REFERENCES dbos.workflow_status(workflow_uuid) 
+    FOREIGN KEY (workflow_uuid) REFERENCES \"{schema}\".workflow_status(workflow_uuid) 
         ON UPDATE CASCADE ON DELETE CASCADE
 );
 
-CREATE TABLE dbos.notifications (
+CREATE TABLE \"{schema}\".notifications (
     destination_uuid TEXT NOT NULL,
     topic TEXT,
     message TEXT NOT NULL,
     created_at_epoch_ms BIGINT NOT NULL DEFAULT (EXTRACT(epoch FROM now()) * 1000::numeric)::bigint,
     message_uuid TEXT NOT NULL DEFAULT gen_random_uuid(), -- Built-in function
-    FOREIGN KEY (destination_uuid) REFERENCES dbos.workflow_status(workflow_uuid) 
+    FOREIGN KEY (destination_uuid) REFERENCES \"{schema}\".workflow_status(workflow_uuid) 
         ON UPDATE CASCADE ON DELETE CASCADE
 );
-CREATE INDEX idx_workflow_topic ON dbos.notifications (destination_uuid, topic);
+CREATE INDEX idx_workflow_topic ON \"{schema}\".notifications (destination_uuid, topic);
 
 -- Create notification function
-CREATE OR REPLACE FUNCTION dbos.notifications_function() RETURNS TRIGGER AS $$
+CREATE OR REPLACE FUNCTION \"{schema}\".notifications_function() RETURNS TRIGGER AS $$
 DECLARE
     payload text := NEW.destination_uuid || '::' || NEW.topic;
 BEGIN
@@ -146,20 +154,20 @@ $$ LANGUAGE plpgsql;
 
 -- Create notification trigger
 CREATE TRIGGER dbos_notifications_trigger
-AFTER INSERT ON dbos.notifications
-FOR EACH ROW EXECUTE FUNCTION dbos.notifications_function();
+AFTER INSERT ON \"{schema}\".notifications
+FOR EACH ROW EXECUTE FUNCTION \"{schema}\".notifications_function();
 
-CREATE TABLE dbos.workflow_events (
+CREATE TABLE \"{schema}\".workflow_events (
     workflow_uuid TEXT NOT NULL,
     key TEXT NOT NULL,
     value TEXT NOT NULL,
     PRIMARY KEY (workflow_uuid, key),
-    FOREIGN KEY (workflow_uuid) REFERENCES dbos.workflow_status(workflow_uuid) 
+    FOREIGN KEY (workflow_uuid) REFERENCES \"{schema}\".workflow_status(workflow_uuid) 
         ON UPDATE CASCADE ON DELETE CASCADE
 );
 
 -- Create events function
-CREATE OR REPLACE FUNCTION dbos.workflow_events_function() RETURNS TRIGGER AS $$
+CREATE OR REPLACE FUNCTION \"{schema}\".workflow_events_function() RETURNS TRIGGER AS $$
 DECLARE
     payload text := NEW.workflow_uuid || '::' || NEW.key;
 BEGIN
@@ -170,20 +178,20 @@ $$ LANGUAGE plpgsql;
 
 -- Create events trigger
 CREATE TRIGGER dbos_workflow_events_trigger
-AFTER INSERT ON dbos.workflow_events
-FOR EACH ROW EXECUTE FUNCTION dbos.workflow_events_function();
+AFTER INSERT ON \"{schema}\".workflow_events
+FOR EACH ROW EXECUTE FUNCTION \"{schema}\".workflow_events_function();
 
-CREATE TABLE dbos.streams (
+CREATE TABLE \"{schema}\".streams (
     workflow_uuid TEXT NOT NULL,
     key TEXT NOT NULL,
     value TEXT NOT NULL,
     "offset" INTEGER NOT NULL,
     PRIMARY KEY (workflow_uuid, key, "offset"),
-    FOREIGN KEY (workflow_uuid) REFERENCES dbos.workflow_status(workflow_uuid) 
+    FOREIGN KEY (workflow_uuid) REFERENCES \"{schema}\".workflow_status(workflow_uuid) 
         ON UPDATE CASCADE ON DELETE CASCADE
 );
 
-CREATE TABLE dbos.event_dispatch_kv (
+CREATE TABLE \"{schema}\".event_dispatch_kv (
     service_name TEXT NOT NULL,
     workflow_fn_name TEXT NOT NULL,
     key TEXT NOT NULL,
@@ -195,6 +203,10 @@ CREATE TABLE dbos.event_dispatch_kv (
 """
 
 
+def get_dbos_migrations(schema: str) -> list[str]:
+    return [get_dbos_migration_one(schema)]
+
+
 def get_sqlite_timestamp_expr() -> str:
     """Get SQLite timestamp expression with millisecond precision for Python >= 3.12."""
     if sys.version_info >= (3, 12):
@@ -281,5 +293,4 @@ CREATE TABLE streams (
 );
 """
 
-dbos_migrations = [dbos_migration_one]
 sqlite_migrations = [sqlite_migration_one]

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_schemas/system_database.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from sqlalchemy import (
     BigInteger,
     Column,
@@ -19,6 +21,21 @@ class SystemSchema:
     metadata_obj = MetaData(schema="dbos")
     sysdb_suffix = "_dbos_sys"
 
+    @classmethod
+    def set_schema(cls, schema_name: Optional[str]) -> None:
+        """
+        Set the schema for all DBOS system tables.
+
+        Args:
+            schema_name: The name of the schema to use for system tables
+        """
+        cls.metadata_obj.schema = schema_name
+        cls.workflow_status.schema = schema_name
+        cls.operation_outputs.schema = schema_name
+        cls.notifications.schema = schema_name
+        cls.workflow_events.schema = schema_name
+        cls.streams.schema = schema_name
+
     workflow_status = Table(
         "workflow_status",
         metadata_obj,

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_sys_db.py

@@ -1443,6 +1443,7 @@ class SystemDatabase(ABC):
     def create(
         system_database_url: str,
         engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
         debug_mode: bool = False,
     ) -> "SystemDatabase":
         """Factory method to create the appropriate SystemDatabase implementation based on URL."""
@@ -1461,6 +1462,7 @@ class SystemDatabase(ABC):
                 system_database_url=system_database_url,
                 engine_kwargs=engine_kwargs,
                 debug_mode=debug_mode,
+                schema=schema,
             )
 
     @db_retry()

{dbos-1.15.0a1 → dbos-1.15.0a3}/dbos/_sys_db_postgres.py

@@ -20,6 +20,7 @@ class PostgresSystemDatabase(SystemDatabase):
         *,
         system_database_url: str,
         engine_kwargs: Dict[str, Any],
+        schema: Optional[str],
         debug_mode: bool = False,
     ):
         super().__init__(
@@ -27,17 +28,16 @@ class PostgresSystemDatabase(SystemDatabase):
             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
 
     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)
 
@@ -62,8 +62,8 @@ class PostgresSystemDatabase(SystemDatabase):
                 conn.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
         engine.dispose()
 
-        ensure_dbos_schema(self.engine)
-        run_dbos_migrations(self.engine)
+        ensure_dbos_schema(self.engine, self.schema)
+        run_dbos_migrations(self.engine, self.schema)
 
     def _cleanup_connections(self) -> None:
         """Clean up PostgreSQL-specific connections."""