fix: improve PostgreSQL migration environment configuration

- Added proper search_path configuration for PostgreSQL
- Improved version table schema handling
This commit is contained in:
hobart 2025-05-23 15:09:03 +09:00
parent a2f12db8d9
commit 54fe1d7f3b
2 changed files with 38 additions and 11 deletions

View File

@ -1,7 +1,7 @@
import logging
from contextvars import ContextVar
from open_webui.env import SRC_LOG_LEVELS
from open_webui.env import SRC_LOG_LEVELS, DATABASE_SCHEMA
from peewee import *
from peewee import InterfaceError as PeeWeeInterfaceError
from peewee import PostgresqlDatabase
@ -52,6 +52,11 @@ def register_connection(db_url):
# Get the connection details
connection = parse(db_url, unquote_user=True, unquote_password=True)
if DATABASE_SCHEMA:
if "options" not in connection:
connection["options"] = f"-c search_path={DATABASE_SCHEMA},public"
else:
connection["options"] += f" -c search_path={DATABASE_SCHEMA},public"
# Use our custom database class that supports reconnection
db = ReconnectingPostgresqlDatabase(**connection)

View File

@ -1,9 +1,9 @@
from logging.config import fileConfig
from alembic import context
from open_webui.models.auths import Auth
from open_webui.internal.db import Base
from open_webui.env import DATABASE_URL
from sqlalchemy import engine_from_config, pool
from sqlalchemy import engine_from_config, pool, text
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
@ -18,7 +18,7 @@ if config.config_file_name is not None:
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = Auth.metadata
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
@ -44,12 +44,20 @@ def run_migrations_offline() -> None:
"""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
# Configure for offline mode with proper schema support
configure_kwargs = {
"url": url,
"target_metadata": target_metadata,
"literal_binds": True,
"dialect_opts": {"paramstyle": "named"},
}
# Add schema configuration if target_metadata.schema is set
if target_metadata.schema:
configure_kwargs["version_table_schema"] = target_metadata.schema
context.configure(**configure_kwargs)
with context.begin_transaction():
context.run_migrations()
@ -69,10 +77,24 @@ def run_migrations_online() -> None:
)
with connectable.connect() as connection:
context.configure(connection=connection, target_metadata=target_metadata)
if target_metadata.schema:
connection.execute(
text(f"SET search_path TO {target_metadata.schema}, public")
)
context.configure(
connection=connection,
target_metadata=target_metadata,
version_table_schema=target_metadata.schema,
)
else:
context.configure(
connection=connection,
target_metadata=target_metadata,
)
with context.begin_transaction():
context.run_migrations()
connection.commit()
if context.is_offline_mode():