DBOS Client should not run migrations (#343)

This PR makes sure that DBOS Client doesn't run migrations on either the
system database or the application database. Added a test.

Author: qianl15

dbos/_app_db.py

@@ -8,6 +8,7 @@
 
 from . import _serialization
 from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
+from ._logger import dbos_logger
 from ._schemas.application_database import ApplicationSchema
 from ._sys_db import StepInfo
 
@@ -39,62 +40,68 @@ def __init__(
     ):
         app_db_url = sa.make_url(database_url).set(drivername="postgresql+psycopg")
 
-        # If the application database does not already exist, create it
-        if not debug_mode:
-            postgres_db_engine = sa.create_engine(
-                app_db_url.set(database="postgres"),
-                **engine_kwargs,
-            )
-            with postgres_db_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": app_db_url.database},
-                ).scalar():
-                    conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
-            postgres_db_engine.dispose()
-
         if engine_kwargs is None:
             engine_kwargs = {}
 
         self.engine = sa.create_engine(
             app_db_url,
             **engine_kwargs,
         )
+        self._engine_kwargs = engine_kwargs
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode
 
+    def run_migrations(self) -> None:
+        if self.debug_mode:
+            dbos_logger.warning(
+                "Application database migrations are skipped in debug mode."
+            )
+            return
+        # Check if the database exists
+        app_db_url = self.engine.url
+        postgres_db_engine = sa.create_engine(
+            app_db_url.set(database="postgres"),
+            **self._engine_kwargs,
+        )
+        with postgres_db_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": app_db_url.database},
+            ).scalar():
+                conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
+        postgres_db_engine.dispose()
+
         # Create the dbos schema and transaction_outputs table in the application database
-        if not debug_mode:
-            with self.engine.begin() as conn:
-                schema_creation_query = sa.text(
-                    f"CREATE SCHEMA IF NOT EXISTS {ApplicationSchema.schema}"
-                )
-                conn.execute(schema_creation_query)
+        with self.engine.begin() as conn:
+            schema_creation_query = sa.text(
+                f"CREATE SCHEMA IF NOT EXISTS {ApplicationSchema.schema}"
+            )
+            conn.execute(schema_creation_query)
 
-            inspector = inspect(self.engine)
-            if not inspector.has_table(
+        inspector = inspect(self.engine)
+        if not inspector.has_table(
+            "transaction_outputs", schema=ApplicationSchema.schema
+        ):
+            ApplicationSchema.metadata_obj.create_all(self.engine)
+        else:
+            columns = inspector.get_columns(
                 "transaction_outputs", schema=ApplicationSchema.schema
-            ):
-                ApplicationSchema.metadata_obj.create_all(self.engine)
-            else:
-                columns = inspector.get_columns(
-                    "transaction_outputs", schema=ApplicationSchema.schema
-                )
-                column_names = [col["name"] for col in columns]
+            )
+            column_names = [col["name"] for col in columns]
 
-                if "function_name" not in column_names:
-                    # Column missing, alter table to add it
-                    with self.engine.connect() as conn:
-                        conn.execute(
-                            text(
-                                f"""
-                            ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
-                            ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
-                            """
-                            )
+            if "function_name" not in column_names:
+                # Column missing, alter table to add it
+                with self.engine.connect() as conn:
+                    conn.execute(
+                        text(
+                            f"""
+                        ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
+                        ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
+                        """
                         )
-                        conn.commit()
+                    )
+                    conn.commit()
 
     def destroy(self) -> None:
         self.engine.dispose()

dbos/_client.py

@@ -99,6 +99,7 @@ async def get_status(self) -> WorkflowStatus:
 
 class DBOSClient:
     def __init__(self, database_url: str, *, system_database: Optional[str] = None):
+        # We only create database connections but do not run migrations
         self._sys_db = SystemDatabase(
             database_url=database_url,
             engine_kwargs={

dbos/_dbos.py

@@ -433,6 +433,10 @@ def _launch(self, *, debug_mode: bool = False) -> None:
             if debug_mode:
                 return
 
+            # Run migrations for the system and application databases
+            self._sys_db.run_migrations()
+            self._app_db.run_migrations()
+
             admin_port = self._config.get("runtimeConfig", {}).get("admin_port")
             if admin_port is None:
                 admin_port = 3001

dbos/_sys_db.py

@@ -241,34 +241,63 @@ def __init__(
             sysdb_name = system_db_url.database + SystemSchema.sysdb_suffix
         system_db_url = system_db_url.set(database=sysdb_name)
 
-        if not debug_mode:
-            # If the system database does not already exist, create it
-            engine = sa.create_engine(
-                system_db_url.set(database="postgres"), **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()
-
         self.engine = sa.create_engine(
             system_db_url,
             **engine_kwargs,
         )
+        self._engine_kwargs = engine_kwargs
+
+        self.notification_conn: Optional[psycopg.connection.Connection] = None
+        self.notifications_map: Dict[str, threading.Condition] = {}
+        self.workflow_events_map: Dict[str, threading.Condition] = {}
+
+        # Now we can run background processes
+        self._run_background_processes = True
+        self._debug_mode = debug_mode
+
+    # Run migrations
+    def run_migrations(self) -> None:
+        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()
 
         # Run a schema migration for the system database
-        if not debug_mode:
-            migration_dir = os.path.join(
-                os.path.dirname(os.path.realpath(__file__)), "_migrations"
+        migration_dir = os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "_migrations"
+        )
+        alembic_cfg = Config()
+        alembic_cfg.set_main_option("script_location", migration_dir)
+        logging.getLogger("alembic").setLevel(logging.WARNING)
+        # Alembic requires the % in URL-escaped parameters to itself be escaped to %%.
+        escaped_conn_string = re.sub(
+            r"%(?=[0-9A-Fa-f]{2})",
+            "%%",
+            self.engine.url.render_as_string(hide_password=False),
+        )
+        alembic_cfg.set_main_option("sqlalchemy.url", escaped_conn_string)
+        try:
+            command.upgrade(alembic_cfg, "head")
+        except Exception as e:
+            dbos_logger.warning(
+                f"Exception during system database construction. This is most likely because the system database was configured using a later version of DBOS: {e}"
             )
             alembic_cfg = Config()
             alembic_cfg.set_main_option("script_location", migration_dir)
-            logging.getLogger("alembic").setLevel(logging.WARNING)
             # Alembic requires the % in URL-escaped parameters to itself be escaped to %%.
             escaped_conn_string = re.sub(
                 r"%(?=[0-9A-Fa-f]{2})",
@@ -282,29 +311,6 @@ def __init__(
                 dbos_logger.warning(
                     f"Exception during system database construction. This is most likely because the system database was configured using a later version of DBOS: {e}"
                 )
-                alembic_cfg = Config()
-                alembic_cfg.set_main_option("script_location", migration_dir)
-                # Alembic requires the % in URL-escaped parameters to itself be escaped to %%.
-                escaped_conn_string = re.sub(
-                    r"%(?=[0-9A-Fa-f]{2})",
-                    "%%",
-                    self.engine.url.render_as_string(hide_password=False),
-                )
-                alembic_cfg.set_main_option("sqlalchemy.url", escaped_conn_string)
-                try:
-                    command.upgrade(alembic_cfg, "head")
-                except Exception as e:
-                    dbos_logger.warning(
-                        f"Exception during system database construction. This is most likely because the system database was configured using a later version of DBOS: {e}"
-                    )
-
-        self.notification_conn: Optional[psycopg.connection.Connection] = None
-        self.notifications_map: Dict[str, threading.Condition] = {}
-        self.workflow_events_map: Dict[str, threading.Condition] = {}
-
-        # Now we can run background processes
-        self._run_background_processes = True
-        self._debug_mode = debug_mode
 
     # Destroy the pool when finished
     def destroy(self) -> None:

dbos/cli/cli.py

@@ -258,6 +258,8 @@ def migrate(
                 "pool_size": 2,
             },
         )
+        sys_db.run_migrations()
+        app_db.run_migrations()
     except Exception as e:
         typer.echo(f"DBOS system schema migration failed: {e}")
     finally:

tests/conftest.py

@@ -49,6 +49,7 @@ def sys_db(config: DBOSConfig) -> Generator[SystemDatabase, Any, None]:
             "connect_args": {"connect_timeout": 30},
         },
     )
+    sys_db.run_migrations()
     yield sys_db
     sys_db.destroy()
 
@@ -65,6 +66,7 @@ def app_db(config: DBOSConfig) -> Generator[ApplicationDatabase, Any, None]:
             "connect_args": {"connect_timeout": 30},
         },
     )
+    app_db.run_migrations()
     yield app_db
     app_db.destroy()
 

tests/test_client.py

@@ -31,6 +31,20 @@ def run_client_collateral() -> None:
     runpy.run_path(filename)
 
 
+def test_client_no_migrate(dbos: DBOS, config: DBOSConfig) -> None:
+    # Drop the system database
+    DBOS.destroy()
+    DBOS(config=config)
+    DBOS.reset_system_database()
+
+    # The client should not be able to connect to the system database
+    with pytest.raises(Exception) as exc_info:
+        assert config["database_url"] is not None
+        client = DBOSClient(config["database_url"])
+        client.list_workflows()
+    assert f'database "dbostestpy_dbos_sys" does not exist' in str(exc_info.value)
+
+
 def test_client_enqueue_and_get_result(dbos: DBOS, client: DBOSClient) -> None:
     run_client_collateral()