From ea0f54064e2cd90e75dfe8772faf50e8bd33b79e Mon Sep 17 00:00:00 2001
From: hobart <hobart@supertone.ai>
Date: Mon, 26 May 2025 11:33:21 +0900
Subject: [PATCH] fix: enhance database schema handling for SQLite and
 PostgreSQL

- Updated schema configuration to set None for SQLite databases.
- Added schema creation logic in the connection wrapper for PostgreSQL.
- Adjusted migration environment to ignore schema for SQLite.
- Fixed unquote parameters in connection tests for consistency.
---
 backend/open_webui/internal/db.py              |  7 ++++++-
 backend/open_webui/internal/wrappers.py        | 18 ++++++++++++++++++
 backend/open_webui/migrations/env.py           |  4 +++-
 .../test/apps/webui/internal/test_wrappers.py  |  6 +++---
 4 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/backend/open_webui/internal/db.py b/backend/open_webui/internal/db.py
index 840f571cc..8c5129708 100644
--- a/backend/open_webui/internal/db.py
+++ b/backend/open_webui/internal/db.py
@@ -100,7 +100,12 @@ else:
 SessionLocal = sessionmaker(
     autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
 )
-metadata_obj = MetaData(schema=DATABASE_SCHEMA)
+# SQLite를 사용하는 경우 스키마를 None으로 설정
+if "sqlite" in SQLALCHEMY_DATABASE_URL:
+    metadata_obj = MetaData(schema=None)
+else:
+    metadata_obj = MetaData(schema=DATABASE_SCHEMA)
+    
 Base = declarative_base(metadata=metadata_obj)
 Session = scoped_session(SessionLocal)
 
diff --git a/backend/open_webui/internal/wrappers.py b/backend/open_webui/internal/wrappers.py
index 2f06de785..496df3973 100644
--- a/backend/open_webui/internal/wrappers.py
+++ b/backend/open_webui/internal/wrappers.py
@@ -52,11 +52,29 @@ def register_connection(db_url):
 
         # Get the connection details
         connection = parse(db_url, unquote_user=True, unquote_password=True)
+        
         if DATABASE_SCHEMA:
+            # Create schema if it doesn't exist
+            try:
+                cursor = db.cursor()
+                log.info(f"Creating schema '{DATABASE_SCHEMA}' if it doesn't exist...")
+                cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {DATABASE_SCHEMA}")
+                db.commit()
+                cursor.close()
+                log.info(f"Schema '{DATABASE_SCHEMA}' is ready")
+            except Exception as e:
+                log.error(f"Error creating schema '{DATABASE_SCHEMA}': {e}")
+                # Continue with the connection even if schema creation fails
+            
+            log.info(f"Setting search path to {DATABASE_SCHEMA},public")
+            # Add schema to search path options
             if "options" not in connection:
                 connection["options"] = f"-c search_path={DATABASE_SCHEMA},public"
             else:
                 connection["options"] += f" -c search_path={DATABASE_SCHEMA},public"
+            
+            # Close the temporary connection
+            db.close()
 
         # 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 95b491286..36f57424b 100644
--- a/backend/open_webui/migrations/env.py
+++ b/backend/open_webui/migrations/env.py
@@ -77,7 +77,8 @@ def run_migrations_online() -> None:
     )
 
     with connectable.connect() as connection:
-        if target_metadata.schema:
+        if target_metadata.schema and connection.dialect.name == 'postgresql':
+            # PostgreSQL에서만 스키마 설정 적용
             connection.execute(
                 text(f"SET search_path TO {target_metadata.schema}, public")
             )
@@ -87,6 +88,7 @@ def run_migrations_online() -> None:
                 version_table_schema=target_metadata.schema,
             )
         else:
+            # For SQLite, ignore schema and use default behavior
             context.configure(
                 connection=connection,
                 target_metadata=target_metadata,
diff --git a/backend/open_webui/test/apps/webui/internal/test_wrappers.py b/backend/open_webui/test/apps/webui/internal/test_wrappers.py
index 3fbb50e57..0e99bbc1f 100644
--- a/backend/open_webui/test/apps/webui/internal/test_wrappers.py
+++ b/backend/open_webui/test/apps/webui/internal/test_wrappers.py
@@ -38,8 +38,8 @@ class TestWrappers:
                 result = register_connection(db_url)
                 
                 # Assertions
-                mock_connect.assert_called_once_with(db_url, unquote_password=True)
-                mock_parse.assert_called_once_with(db_url, unquote_password=True)
+                mock_connect.assert_called_once_with(db_url, unquote_user=True, unquote_password=True)
+                mock_parse.assert_called_once_with(db_url, unquote_user=True, unquote_password=True)
                 
                 # Check that options were added with schema
                 expected_connection = {
@@ -186,7 +186,7 @@ class TestWrappers:
                 result = register_connection(db_url)
                 
                 # Assertions
-                mock_connect.assert_called_once_with(db_url, unquote_password=True)
+                mock_connect.assert_called_once_with(db_url, unquote_user=True, unquote_password=True)
                 
                 # Verify SQLite database properties are set
                 assert mock_sqlite_db.autoconnect is True