copilot-swe-agent[bot]
f6fcaaeccc
Add a guided database configuration wizard and a data migration tool that allows users to: - Build database connection strings through a step-by-step UI - Test database connections before applying - Preview and execute data migrations from SQLite to PostgreSQL/MySQL - Copy to clipboard for easy .env file updates New files: - app/utils/db_wizard.py — connection string builder, parser, and tester - app/utils/db_migrate.py — table-by-table data migration utility - app/api/database.py — REST API endpoints for wizard operations - app/views/db_wizard.py — view route for the wizard page - frontend/templates/db_wizard.html — multi-tab wizard UI - tests/test_db_wizard.py — unit tests for db_wizard utilities - tests/test_db_migrate.py — unit tests for db_migrate utilities - tests/test_db_wizard_api.py — integration tests for API and views Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
249 lines
8.8 KiB
Python
249 lines
8.8 KiB
Python
"""
|
|
Database migration utility for transferring data between databases.
|
|
|
|
Copies all table rows from a *source* SQLAlchemy database to a *target*
|
|
database. This is designed for the common scenario of migrating from the
|
|
built-in SQLite database to an external PostgreSQL / MySQL instance.
|
|
|
|
The utility:
|
|
1. Creates the schema in the target via ``Base.metadata.create_all``.
|
|
2. Copies rows table-by-table in dependency order.
|
|
3. Stamps the Alembic version in the target to ``head``.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from sqlalchemy import MetaData, create_engine, inspect, text
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.engine.url import make_url
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Tables to skip during migration (Alembic manages its own state).
|
|
_SKIP_TABLES = {"alembic_version"}
|
|
|
|
# Ordered list — parent tables first to respect foreign-key constraints.
|
|
_TABLE_ORDER = [
|
|
"documents",
|
|
"files",
|
|
"file_processing_steps",
|
|
"processing_logs",
|
|
"application_settings",
|
|
"settings_audit_log",
|
|
"saved_searches",
|
|
"webhook_configs",
|
|
]
|
|
|
|
|
|
def _make_engine(url: str) -> Engine:
|
|
"""Create a SQLAlchemy engine from *url* with sensible defaults."""
|
|
parsed = make_url(url)
|
|
connect_args: dict[str, Any] = {}
|
|
if parsed.get_backend_name() == "sqlite":
|
|
connect_args["check_same_thread"] = False
|
|
return create_engine(url, connect_args=connect_args)
|
|
|
|
|
|
def _ordered_tables(inspector: Any) -> list[str]:
|
|
"""Return table names in safe insertion order.
|
|
|
|
Tables listed in ``_TABLE_ORDER`` come first (in that order); any
|
|
remaining tables are appended alphabetically.
|
|
"""
|
|
existing = set(inspector.get_table_names())
|
|
ordered: list[str] = []
|
|
for name in _TABLE_ORDER:
|
|
if name in existing and name not in _SKIP_TABLES:
|
|
ordered.append(name)
|
|
for name in sorted(existing):
|
|
if name not in ordered and name not in _SKIP_TABLES:
|
|
ordered.append(name)
|
|
return ordered
|
|
|
|
|
|
def preview_migration(source_url: str) -> dict[str, Any]:
|
|
"""Preview what a migration would do without actually copying data.
|
|
|
|
Args:
|
|
source_url: Connection string for the source database.
|
|
|
|
Returns:
|
|
Dict with ``tables`` (list of dicts with ``name`` and ``row_count``)
|
|
and ``total_rows``.
|
|
"""
|
|
try:
|
|
src_engine = _make_engine(source_url)
|
|
src_inspector = inspect(src_engine)
|
|
tables = _ordered_tables(src_inspector)
|
|
|
|
result: list[dict[str, Any]] = []
|
|
total = 0
|
|
with src_engine.connect() as conn:
|
|
for table_name in tables:
|
|
row = conn.execute(text(f'SELECT COUNT(*) FROM "{table_name}"')).fetchone() # noqa: S608
|
|
count = row[0] if row else 0
|
|
result.append({"name": table_name, "row_count": count})
|
|
total += count
|
|
|
|
src_engine.dispose()
|
|
return {"tables": result, "total_rows": total, "success": True}
|
|
except Exception as exc:
|
|
logger.error(f"Migration preview failed: {exc}")
|
|
return {"success": False, "error": str(exc), "tables": [], "total_rows": 0}
|
|
|
|
|
|
def migrate_data(
|
|
source_url: str,
|
|
target_url: str,
|
|
*,
|
|
batch_size: int = 500,
|
|
progress_callback: Any | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Copy all data from *source_url* to *target_url*.
|
|
|
|
The target schema is created automatically from the application models.
|
|
Alembic is stamped to ``head`` in the target after a successful copy.
|
|
|
|
Args:
|
|
source_url: SQLAlchemy connection string for the source DB.
|
|
target_url: SQLAlchemy connection string for the target DB.
|
|
batch_size: Number of rows to insert per batch.
|
|
progress_callback: Optional ``callable(table_name, copied, total)``
|
|
invoked after each batch.
|
|
|
|
Returns:
|
|
Dict with ``success`` (bool), ``tables_copied`` (int),
|
|
``rows_copied`` (int), and ``errors`` (list of str).
|
|
"""
|
|
errors: list[str] = []
|
|
tables_copied = 0
|
|
rows_copied = 0
|
|
|
|
try:
|
|
src_engine = _make_engine(source_url)
|
|
tgt_engine = _make_engine(target_url)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 1. Create schema in target from application models
|
|
# ------------------------------------------------------------------
|
|
from app.database import Base # local import to avoid circular deps
|
|
|
|
Base.metadata.create_all(bind=tgt_engine)
|
|
logger.info("Target schema created from application models.")
|
|
|
|
# ------------------------------------------------------------------
|
|
# 2. Reflect source schema & determine copy order
|
|
# ------------------------------------------------------------------
|
|
src_meta = MetaData()
|
|
src_meta.reflect(bind=src_engine)
|
|
|
|
src_inspector = inspect(src_engine)
|
|
table_names = _ordered_tables(src_inspector)
|
|
|
|
SrcSession = sessionmaker(bind=src_engine)
|
|
TgtSession = sessionmaker(bind=tgt_engine)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 3. Copy data table-by-table
|
|
# ------------------------------------------------------------------
|
|
for table_name in table_names:
|
|
try:
|
|
src_session = SrcSession()
|
|
tgt_session = TgtSession()
|
|
|
|
src_table = src_meta.tables.get(table_name)
|
|
if src_table is None:
|
|
continue
|
|
|
|
# Read all rows from source
|
|
rows = src_session.execute(src_table.select()).fetchall()
|
|
column_names = [c.name for c in src_table.columns]
|
|
|
|
if not rows:
|
|
logger.info(f"Skipping empty table: {table_name}")
|
|
tables_copied += 1
|
|
src_session.close()
|
|
tgt_session.close()
|
|
continue
|
|
|
|
# Reflect the target table to insert into
|
|
tgt_meta = MetaData()
|
|
tgt_meta.reflect(bind=tgt_engine, only=[table_name])
|
|
tgt_table = tgt_meta.tables.get(table_name)
|
|
if tgt_table is None:
|
|
errors.append(f"Target table {table_name} not found after schema creation")
|
|
src_session.close()
|
|
tgt_session.close()
|
|
continue
|
|
|
|
# Batch insert
|
|
total_for_table = len(rows)
|
|
for i in range(0, total_for_table, batch_size):
|
|
batch = rows[i : i + batch_size]
|
|
insert_data = [dict(zip(column_names, row, strict=False)) for row in batch]
|
|
tgt_session.execute(tgt_table.insert(), insert_data)
|
|
tgt_session.commit()
|
|
|
|
rows_copied += len(batch)
|
|
if progress_callback:
|
|
progress_callback(table_name, min(i + batch_size, total_for_table), total_for_table)
|
|
|
|
tables_copied += 1
|
|
logger.info(f"Copied {total_for_table} rows from {table_name}")
|
|
src_session.close()
|
|
tgt_session.close()
|
|
|
|
except Exception as exc:
|
|
msg = f"Error copying table {table_name}: {exc}"
|
|
logger.error(msg)
|
|
errors.append(msg)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 4. Stamp Alembic to head in the target
|
|
# ------------------------------------------------------------------
|
|
try:
|
|
_stamp_alembic_head(tgt_engine)
|
|
logger.info("Alembic version stamped to head in target database.")
|
|
except Exception as exc:
|
|
msg = f"Failed to stamp Alembic version: {exc}"
|
|
logger.error(msg)
|
|
errors.append(msg)
|
|
|
|
src_engine.dispose()
|
|
tgt_engine.dispose()
|
|
|
|
return {
|
|
"success": len(errors) == 0,
|
|
"tables_copied": tables_copied,
|
|
"rows_copied": rows_copied,
|
|
"errors": errors,
|
|
}
|
|
|
|
except Exception as exc:
|
|
logger.error(f"Migration failed: {exc}")
|
|
return {
|
|
"success": False,
|
|
"tables_copied": tables_copied,
|
|
"rows_copied": rows_copied,
|
|
"errors": errors + [str(exc)],
|
|
}
|
|
|
|
|
|
def _stamp_alembic_head(engine: Engine) -> None:
|
|
"""Stamp the Alembic version table to ``head`` in the given engine."""
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
|
|
migrations_dir = str(Path(__file__).resolve().parent.parent.parent / "migrations")
|
|
alembic_cfg = Config()
|
|
alembic_cfg.set_main_option("script_location", migrations_dir)
|
|
alembic_cfg.set_main_option("sqlalchemy.url", "")
|
|
|
|
with engine.begin() as connection:
|
|
alembic_cfg.attributes["connection"] = connection
|
|
command.stamp(alembic_cfg, "head")
|