From 54fe1d7f3b7979bfcfe1db000499be23feb96a91 Mon Sep 17 00:00:00 2001
From: hobart <hobart@supertone.ai>
Date: Fri, 23 May 2025 15:09:03 +0900
Subject: [PATCH] fix: improve PostgreSQL migration environment configuration

- Added proper search_path configuration for PostgreSQL
- Improved version table schema handling
---
 backend/open_webui/internal/wrappers.py |  7 ++++-
 backend/open_webui/migrations/env.py    | 42 +++++++++++++++++++------
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/backend/open_webui/internal/wrappers.py b/backend/open_webui/internal/wrappers.py
index 5cf352930..2f06de785 100644
--- a/backend/open_webui/internal/wrappers.py
+++ b/backend/open_webui/internal/wrappers.py
@@ -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)
diff --git a/backend/open_webui/migrations/env.py b/backend/open_webui/migrations/env.py
index 128881647..95b491286 100644
--- a/backend/open_webui/migrations/env.py
+++ b/backend/open_webui/migrations/env.py
@@ -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():