mirror of
https://github.com/open-webui/open-webui
synced 2025-06-11 00:49:44 +00:00
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.
This commit is contained in:
parent
b203a72c61
commit
ea0f54064e
@ -100,7 +100,12 @@ else:
|
|||||||
SessionLocal = sessionmaker(
|
SessionLocal = sessionmaker(
|
||||||
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
||||||
)
|
)
|
||||||
|
# SQLite를 사용하는 경우 스키마를 None으로 설정
|
||||||
|
if "sqlite" in SQLALCHEMY_DATABASE_URL:
|
||||||
|
metadata_obj = MetaData(schema=None)
|
||||||
|
else:
|
||||||
metadata_obj = MetaData(schema=DATABASE_SCHEMA)
|
metadata_obj = MetaData(schema=DATABASE_SCHEMA)
|
||||||
|
|
||||||
Base = declarative_base(metadata=metadata_obj)
|
Base = declarative_base(metadata=metadata_obj)
|
||||||
Session = scoped_session(SessionLocal)
|
Session = scoped_session(SessionLocal)
|
||||||
|
|
||||||
|
@ -52,12 +52,30 @@ def register_connection(db_url):
|
|||||||
|
|
||||||
# Get the connection details
|
# Get the connection details
|
||||||
connection = parse(db_url, unquote_user=True, unquote_password=True)
|
connection = parse(db_url, unquote_user=True, unquote_password=True)
|
||||||
|
|
||||||
if DATABASE_SCHEMA:
|
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:
|
if "options" not in connection:
|
||||||
connection["options"] = f"-c search_path={DATABASE_SCHEMA},public"
|
connection["options"] = f"-c search_path={DATABASE_SCHEMA},public"
|
||||||
else:
|
else:
|
||||||
connection["options"] += f" -c search_path={DATABASE_SCHEMA},public"
|
connection["options"] += f" -c search_path={DATABASE_SCHEMA},public"
|
||||||
|
|
||||||
|
# Close the temporary connection
|
||||||
|
db.close()
|
||||||
|
|
||||||
# Use our custom database class that supports reconnection
|
# Use our custom database class that supports reconnection
|
||||||
db = ReconnectingPostgresqlDatabase(**connection)
|
db = ReconnectingPostgresqlDatabase(**connection)
|
||||||
db.connect(reuse_if_open=True)
|
db.connect(reuse_if_open=True)
|
||||||
|
@ -77,7 +77,8 @@ def run_migrations_online() -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
with connectable.connect() as connection:
|
with connectable.connect() as connection:
|
||||||
if target_metadata.schema:
|
if target_metadata.schema and connection.dialect.name == 'postgresql':
|
||||||
|
# PostgreSQL에서만 스키마 설정 적용
|
||||||
connection.execute(
|
connection.execute(
|
||||||
text(f"SET search_path TO {target_metadata.schema}, public")
|
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,
|
version_table_schema=target_metadata.schema,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
# For SQLite, ignore schema and use default behavior
|
||||||
context.configure(
|
context.configure(
|
||||||
connection=connection,
|
connection=connection,
|
||||||
target_metadata=target_metadata,
|
target_metadata=target_metadata,
|
||||||
|
@ -38,8 +38,8 @@ class TestWrappers:
|
|||||||
result = register_connection(db_url)
|
result = register_connection(db_url)
|
||||||
|
|
||||||
# Assertions
|
# 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)
|
||||||
mock_parse.assert_called_once_with(db_url, unquote_password=True)
|
mock_parse.assert_called_once_with(db_url, unquote_user=True, unquote_password=True)
|
||||||
|
|
||||||
# Check that options were added with schema
|
# Check that options were added with schema
|
||||||
expected_connection = {
|
expected_connection = {
|
||||||
@ -186,7 +186,7 @@ class TestWrappers:
|
|||||||
result = register_connection(db_url)
|
result = register_connection(db_url)
|
||||||
|
|
||||||
# Assertions
|
# 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
|
# Verify SQLite database properties are set
|
||||||
assert mock_sqlite_db.autoconnect is True
|
assert mock_sqlite_db.autoconnect is True
|
||||||
|
Loading…
Reference in New Issue
Block a user