feat: support postgres schema overrides

.all-contributorsrc

@@ -1222,6 +1222,17 @@
       "contributions": [
         "code"
       ]
+    },
+    {
+      "login": "hongzexin",
+      "name": "Jason HONG",
+      "avatar_url": "https://avatars.githubusercontent.com/u/136784169?v=4",
+      "profile": "https://github.com/hongzexin",
+      "contributions": [
+        "code",
+        "test",
+        "doc"
+      ]
     }
   ],
   "contributorsPerLine": 7,

.env.example

@@ -15,6 +15,8 @@
 
 # Database. SQLite (default) needs nothing. For PostgreSQL:
 # CODEX_LB_DATABASE_URL=postgresql+asyncpg://codex_lb:codex_lb@127.0.0.1:5432/codex_lb
+# Optional shared-DB schema isolation (runtime "<schema>,public"; migrations "<schema>"):
+# CODEX_LB_DATABASE_POSTGRES_SCHEMA=codex_lb_prod
 
 # Encryption key file (pin for Docker volumes; must be shared across replicas).
 # CODEX_LB_ENCRYPTION_KEY_FILE=/var/lib/codex-lb/encryption.key

README.md

@@ -283,6 +283,9 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/e
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/glopyglerky"><img src="https://avatars.githubusercontent.com/u/189872235?v=4?s=100" width="100px;" alt="glopyglerky"/><br /><sub><b>glopyglerky</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=glopyglerky" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=glopyglerky" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/iqbalmaulana03"><img src="https://avatars.githubusercontent.com/u/78488507?v=4?s=100" width="100px;" alt="Ahmad Maulana Iqbal"/><br /><sub><b>Ahmad Maulana Iqbal</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=iqbalmaulana03" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=iqbalmaulana03" title="Tests">⚠️</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/hongzexin"><img src="https://avatars.githubusercontent.com/u/136784169?v=4?s=100" width="100px;" alt="Jason HONG"/><br /><sub><b>Jason HONG</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=hongzexin" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=hongzexin" title="Tests">⚠️</a> <a href="https://github.com/Soju06/codex-lb/commits?author=hongzexin" title="Documentation">📖</a></td>
+    </tr>
   </tbody>
 </table>
 

app/core/config/settings.py

@@ -249,6 +249,7 @@ class Settings(BaseSettings):
     database_sqlite_pre_migrate_backup_enabled: bool = True
     database_sqlite_pre_migrate_backup_max_files: int = Field(default=5, ge=1)
     database_sqlite_startup_check_mode: Literal["quick", "full", "off"] = "quick"
+    database_postgres_schema: str | None = None
     database_alembic_auto_remap_enabled: bool = True
     database_migration_lock_timeout_seconds: float = Field(default=300.0, gt=0)
     upstream_base_url: str = "https://chatgpt.com/backend-api"
@@ -491,6 +492,20 @@ def _expand_database_url(cls, value: str) -> str:
                     return f"{prefix}{Path(path).expanduser()}"
         return value
 
+    @field_validator("database_postgres_schema", mode="before")
+    @classmethod
+    def _normalize_database_postgres_schema(cls, value: OptionalStringInput) -> str | None:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            stripped = value.strip()
+            if not stripped:
+                return None
+            if len(stripped.encode("utf-8")) > 63:
+                raise ValueError("database_postgres_schema must be at most 63 UTF-8 bytes")
+            return stripped
+        raise TypeError("database_postgres_schema must be a string")
+
     @field_validator("encryption_key_file", mode="before")
     @classmethod
     def _expand_encryption_key_file(cls, value: str | Path) -> Path:

app/db/alembic/env.py

@@ -3,10 +3,16 @@
 from logging.config import fileConfig
 
 from alembic import context
-from sqlalchemy import engine_from_config, pool
+from sqlalchemy import engine_from_config, make_url, pool
 
 from app.core.config.settings import get_settings
-from app.db.migration_url import to_sync_database_url
+from app.db.migration_url import (
+    apply_postgres_migration_search_path,
+    ensure_postgres_schema_exists,
+    normalize_postgres_schema,
+    postgres_migration_search_path,
+    to_sync_database_url,
+)
 from app.db.models import Base
 
 config = context.config
@@ -27,6 +33,9 @@ def _sync_database_url() -> str:
 def run_migrations_offline() -> None:
     url = _sync_database_url()
     config.set_main_option("sqlalchemy.url", url)
+    schema = None
+    if make_url(url).get_backend_name() == "postgresql":
+        schema = normalize_postgres_schema(get_settings().database_postgres_schema)
 
     context.configure(
         url=url,
@@ -35,8 +44,13 @@ def run_migrations_offline() -> None:
         dialect_opts={"paramstyle": "named"},
         compare_type=True,
         render_as_batch=url.startswith("sqlite"),
+        version_table_schema=schema,
     )
 
+    search_path = postgres_migration_search_path(schema)
+    if search_path is not None:
+        context.execute(f"SET search_path TO {search_path}")
+
     with context.begin_transaction():
         context.run_migrations()
 
@@ -53,11 +67,19 @@ def run_migrations_online() -> None:
     )
 
     with connectable.connect() as connection:
+        schema = normalize_postgres_schema(get_settings().database_postgres_schema)
+        ensure_postgres_schema_exists(connection, schema)
+        apply_postgres_migration_search_path(connection, schema)
+        if connection.dialect.name == "postgresql" and connection.in_transaction():
+            # Schema creation and set_config autobegin a transaction. Commit it
+            # before Alembic owns transaction boundaries and autocommit blocks.
+            connection.commit()
         context.configure(
             connection=connection,
             target_metadata=target_metadata,
             compare_type=True,
             render_as_batch=connection.dialect.name == "sqlite",
+            version_table_schema=schema if connection.dialect.name == "postgresql" else None,
         )
 
         with context.begin_transaction():

app/db/alembic/versions/20260310_000000_fix_postgresql_enum_value_casing.py

@@ -43,7 +43,8 @@ def _enum_value_exists(bind: sa.engine.Connection, enum_type_name: str, enum_val
         sa.text(
             "SELECT 1 FROM pg_enum e "
             "JOIN pg_type t ON e.enumtypid = t.oid "
-            "WHERE t.typname = :type_name AND e.enumlabel = :value"
+            "WHERE pg_type_is_visible(t.oid) "
+            "AND t.typname = :type_name AND e.enumlabel = :value"
         ),
         {"type_name": enum_type_name, "value": enum_value},
     )

app/db/alembic/versions/20260604_000000_add_reauth_required_account_status.py

@@ -73,7 +73,8 @@ def _enum_value_exists(enum_type_name: str, enum_value: str) -> bool:
         sa.text(
             "SELECT 1 FROM pg_enum e "
             "JOIN pg_type t ON e.enumtypid = t.oid "
-            "WHERE t.typname = :type_name AND e.enumlabel = :value"
+            "WHERE pg_type_is_visible(t.oid) "
+            "AND t.typname = :type_name AND e.enumlabel = :value"
         ),
         {"type_name": enum_type_name, "value": enum_value},
     )

app/db/migrate.py

@@ -22,7 +22,13 @@
 from app.core.config.settings import get_settings
 from app.db.alembic.revision_ids import LEGACY_MIGRATION_TO_NEW_REVISION, OLD_TO_NEW_REVISION_MAP, REVISION_ID_PATTERN
 from app.db.migration_lock import migration_lock
-from app.db.migration_url import to_sync_database_url
+from app.db.migration_url import (
+    apply_postgres_migration_search_path,
+    ensure_postgres_schema_exists,
+    normalize_postgres_schema,
+    postgres_qualified_name,
+    to_sync_database_url,
+)
 from app.db.models import Base
 
 logger = logging.getLogger(__name__)
@@ -178,6 +184,8 @@ def _sync_connection(sync_database_url: str) -> Iterator[Connection]:
     engine = create_engine(sync_database_url, future=True)
     try:
         with engine.connect() as connection:
+            ensure_postgres_schema_exists(connection, get_settings().database_postgres_schema)
+            apply_postgres_migration_search_path(connection, get_settings().database_postgres_schema)
             yield connection
     finally:
         engine.dispose()
@@ -188,25 +196,40 @@ def _sync_transaction(sync_database_url: str) -> Iterator[Connection]:
     engine = create_engine(sync_database_url, future=True)
     try:
         with engine.begin() as connection:
+            ensure_postgres_schema_exists(connection, get_settings().database_postgres_schema)
+            apply_postgres_migration_search_path(connection, get_settings().database_postgres_schema)
             yield connection
     finally:
         engine.dispose()
 
 
 def _read_table_names(connection: Connection) -> set[str]:
     inspector = inspect(connection)
-    return set(inspector.get_table_names())
+    schema = _configured_postgres_schema(connection)
+    return set(inspector.get_table_names(schema=schema))
+
+
+def _configured_postgres_schema(connection: Connection) -> str | None:
+    if connection.dialect.name != "postgresql":
+        return None
+    return normalize_postgres_schema(get_settings().database_postgres_schema)
+
+
+def _migration_table_name(connection: Connection, table_name: str) -> str:
+    return postgres_qualified_name(table_name, _configured_postgres_schema(connection))
 
 
 def _read_legacy_migration_names(connection: Connection) -> set[str]:
-    result = connection.execute(text(f"SELECT name FROM {_LEGACY_MIGRATIONS_TABLE}"))
+    table_name = _migration_table_name(connection, _LEGACY_MIGRATIONS_TABLE)
+    result = connection.execute(text(f"SELECT name FROM {table_name}"))
     names = {str(row[0]) for row in result.fetchall() if row and row[0] is not None}
     return names
 
 
 def _read_current_revisions_from_connection(connection: Connection) -> tuple[str, ...]:
+    table_name = _migration_table_name(connection, _ALEMBIC_VERSION_TABLE)
     try:
-        rows = connection.execute(text(f"SELECT {_ALEMBIC_VERSION_COLUMN} FROM {_ALEMBIC_VERSION_TABLE}")).fetchall()
+        rows = connection.execute(text(f"SELECT {_ALEMBIC_VERSION_COLUMN} FROM {table_name}")).fetchall()
     except (sa_exc.ProgrammingError, sa_exc.OperationalError) as exc:
         # PostgreSQL can still raise UndefinedTable here on a fresh database if
         # the alembic_version table is absent when startup migration state is
@@ -355,15 +378,17 @@ def _ensure_alembic_version_table_capacity_for_connection(connection: Connection
         return
 
     inspector = inspect(connection)
-    if not inspector.has_table(_ALEMBIC_VERSION_TABLE):
+    schema = _configured_postgres_schema(connection)
+    table_name = postgres_qualified_name(_ALEMBIC_VERSION_TABLE, schema)
+    if not inspector.has_table(_ALEMBIC_VERSION_TABLE, schema=schema):
         # IF NOT EXISTS is defense-in-depth against concurrent out-of-band
         # `alembic upgrade` invocations that bypass run_upgrade's migration
         # lock; the product paths are already serialized by migration_lock.
         connection.execute(
             text(
                 " ".join(
                     (
-                        f"CREATE TABLE IF NOT EXISTS {_ALEMBIC_VERSION_TABLE} (",
+                        f"CREATE TABLE IF NOT EXISTS {table_name} (",
                         f"{_ALEMBIC_VERSION_COLUMN} VARCHAR({required_length}) NOT NULL,",
                         f"PRIMARY KEY ({_ALEMBIC_VERSION_COLUMN})",
                         ")",
@@ -373,7 +398,7 @@ def _ensure_alembic_version_table_capacity_for_connection(connection: Connection
         )
         return
 
-    columns = inspector.get_columns(_ALEMBIC_VERSION_TABLE)
+    columns = inspector.get_columns(_ALEMBIC_VERSION_TABLE, schema=schema)
     version_num_column = next((column for column in columns if column.get("name") == _ALEMBIC_VERSION_COLUMN), None)
     if version_num_column is None:
         raise MigrationBootstrapError(
@@ -385,10 +410,7 @@ def _ensure_alembic_version_table_capacity_for_connection(connection: Connection
         return
 
     connection.execute(
-        text(
-            f"ALTER TABLE {_ALEMBIC_VERSION_TABLE} "
-            f"ALTER COLUMN {_ALEMBIC_VERSION_COLUMN} TYPE VARCHAR({required_length})"
-        )
+        text(f"ALTER TABLE {table_name} ALTER COLUMN {_ALEMBIC_VERSION_COLUMN} TYPE VARCHAR({required_length})")
     )
 
 
@@ -478,13 +500,11 @@ def _remap_legacy_alembic_revisions(config: Config) -> tuple[str, ...]:
         if remapped == current_revisions:
             return ()
 
-        connection.execute(text(f"DELETE FROM {_ALEMBIC_VERSION_TABLE}"))
+        table_name = _migration_table_name(connection, _ALEMBIC_VERSION_TABLE)
+        connection.execute(text(f"DELETE FROM {table_name}"))
         for revision in remapped:
             connection.execute(
-                text(
-                    f"INSERT INTO {_ALEMBIC_VERSION_TABLE} ({_ALEMBIC_VERSION_COLUMN}) "
-                    f"VALUES (:{_ALEMBIC_VERSION_COLUMN})"
-                ),
+                text(f"INSERT INTO {table_name} ({_ALEMBIC_VERSION_COLUMN}) VALUES (:{_ALEMBIC_VERSION_COLUMN})"),
                 {_ALEMBIC_VERSION_COLUMN: revision},
             )
 

app/db/migration_url.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from sqlalchemy.engine import make_url
+from sqlalchemy import text
+from sqlalchemy.engine import Connection, make_url
+from sqlalchemy.schema import CreateSchema
 
 from app.db.sqlite_utils import normalize_sqlite_url
 
@@ -16,3 +18,78 @@ def to_sync_database_url(database_url: str) -> str:
         parsed = parsed.set(drivername="postgresql+psycopg")
 
     return parsed.render_as_string(hide_password=False)
+
+
+def normalize_postgres_schema(schema: str | None) -> str | None:
+    if schema is None:
+        return None
+    normalized = schema.strip()
+    if not normalized:
+        return None
+    return normalized
+
+
+def quote_postgres_identifier(identifier: str) -> str:
+    return f'"{identifier.replace(chr(34), chr(34) * 2)}"'
+
+
+def postgres_search_path(schema: str | None) -> str | None:
+    normalized = normalize_postgres_schema(schema)
+    if normalized is None:
+        return None
+    quoted = quote_postgres_identifier(normalized)
+    if normalized == "public":
+        return quoted
+    return f"{quoted},public"
+
+
+def postgres_migration_search_path(schema: str | None) -> str | None:
+    normalized = normalize_postgres_schema(schema)
+    if normalized is None:
+        return None
+    return quote_postgres_identifier(normalized)
+
+
+def postgres_qualified_name(identifier: str, schema: str | None) -> str:
+    normalized = normalize_postgres_schema(schema)
+    quoted_identifier = quote_postgres_identifier(identifier)
+    if normalized is None:
+        return quoted_identifier
+    return f"{quote_postgres_identifier(normalized)}.{quoted_identifier}"
+
+
+def ensure_postgres_schema_exists(connection: Connection, schema: str | None) -> None:
+    if schema is None or connection.dialect.name != "postgresql":
+        return
+
+    normalized = normalize_postgres_schema(schema)
+    if normalized is None:
+        return
+    if normalized == "public":
+        return
+
+    connection.execute(CreateSchema(normalized, if_not_exists=True))
+
+
+def apply_postgres_search_path(connection: Connection, schema: str | None) -> None:
+    search_path = postgres_search_path(schema)
+    if search_path is None or connection.dialect.name != "postgresql":
+        return
+    connection.execute(
+        text("SELECT set_config('search_path', :search_path, false)"),
+        {"search_path": search_path},
+    )
+
+
+def apply_postgres_migration_search_path(connection: Connection, schema: str | None) -> None:
+    normalized = normalize_postgres_schema(schema)
+    search_path = postgres_migration_search_path(normalized)
+    if search_path is None or connection.dialect.name != "postgresql":
+        return
+    connection.execute(
+        text("SELECT set_config('search_path', :search_path, false)"),
+        {"search_path": search_path},
+    )
+    # Reflection caches PostgreSQL's default schema before this deployment-
+    # specific search path is applied. Keep it aligned with unqualified DDL.
+    connection.dialect.default_schema_name = normalized

app/db/session.py

@@ -17,6 +17,7 @@
 from sqlalchemy.pool import NullPool
 
 from app.core.config.settings import get_settings
+from app.db.migration_url import postgres_search_path
 from app.db.sqlite_utils import (
     SqliteIntegrityCheckMode,
     check_sqlite_integrity,
@@ -76,7 +77,11 @@ def _postgres_async_connect_args(url: str) -> dict[str, object] | None:
     # bridge-session cleanup stop running, and account/stream lease expiry is
     # mis-evaluated. Forcing UTC keeps stored timestamps correct regardless of
     # the container time zone.
-    connect_args: dict[str, object] = {"server_settings": {"timezone": "UTC"}}
+    server_settings: dict[str, str] = {"timezone": "UTC"}
+    search_path = postgres_search_path(_settings.database_postgres_schema)
+    if search_path is not None:
+        server_settings["search_path"] = search_path
+    connect_args: dict[str, object] = {"server_settings": server_settings}
     if os.environ.get("CODEX_LB_TEST_DATABASE_URL"):
         connect_args["prepared_statement_cache_size"] = 0
     return connect_args

app/modules/proxy/durable_bridge_repository.py

@@ -1786,7 +1786,7 @@ async def missing_durable_bridge_tables(session: AsyncSession) -> tuple[str, ...
         result = await session.execute(
             text(
                 "SELECT table_name FROM information_schema.tables "
-                "WHERE table_schema = 'public' "
+                "WHERE table_schema = ANY (current_schemas(false)) "
                 "AND table_name IN ("
                 "'http_bridge_sessions', 'http_bridge_session_aliases', 'http_bridge_retry_circuits', "
                 "'http_bridge_recovery_attempts'"