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:
hobart 2025-05-26 11:33:21 +09:00
parent b203a72c61
commit ea0f54064e
4 changed files with 30 additions and 5 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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,

View File

@ -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