Spaces:
Running
Running
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| from pathlib import Path | |
| from sqlalchemy import MetaData, create_engine, func, inspect, select, text | |
| ROOT_DIR = Path(__file__).resolve().parent.parent | |
| if str(ROOT_DIR) not in sys.path: | |
| sys.path.insert(0, str(ROOT_DIR)) | |
| from cert_study_app.config import DB_PATH | |
| from cert_study_app.db import Base | |
| from cert_study_app.models import Attempt, ConceptNote, IngestionJob, Question # noqa: F401 | |
| DEFAULT_TARGET_URL = "postgresql+pg8000://cert_study:cert_study@localhost:5432/cert_study" | |
| def sqlite_url(path: Path) -> str: | |
| return f"sqlite:///{path.expanduser().resolve().as_posix()}" | |
| def table_exists(engine, table_name: str) -> bool: | |
| return table_name in inspect(engine).get_table_names() | |
| def reset_postgres_sequence(conn, table_name: str, pk_name: str) -> None: | |
| sequence = conn.execute( | |
| text("SELECT pg_get_serial_sequence(:table_name, :pk_name)"), | |
| {"table_name": table_name, "pk_name": pk_name}, | |
| ).scalar() | |
| if not sequence: | |
| return | |
| max_id = conn.execute(text(f'SELECT COALESCE(MAX("{pk_name}"), 0) FROM "{table_name}"')).scalar() | |
| conn.execute(text("SELECT setval(:sequence, :value, :called)"), {"sequence": sequence, "value": max_id or 1, "called": bool(max_id)}) | |
| def migrate(source_path: Path, target_url: str, replace: bool = False) -> dict: | |
| source_engine = create_engine(sqlite_url(source_path), future=True) | |
| target_engine = create_engine(target_url, future=True) | |
| if replace: | |
| Base.metadata.drop_all(bind=target_engine) | |
| Base.metadata.create_all(bind=target_engine) | |
| source_meta = MetaData() | |
| source_meta.reflect(bind=source_engine) | |
| summary = {} | |
| with target_engine.begin() as target_conn: | |
| with source_engine.connect() as source_conn: | |
| for table in Base.metadata.sorted_tables: | |
| source_table = source_meta.tables.get(table.name) | |
| if source_table is None: | |
| continue | |
| source_columns = {column.name for column in source_table.columns} | |
| target_columns = [column.name for column in table.columns if column.name in source_columns] | |
| rows = [ | |
| {column: row._mapping[column] for column in target_columns} | |
| for row in source_conn.execute(select(*[source_table.c[column] for column in target_columns])) | |
| ] | |
| if rows: | |
| target_conn.execute(table.insert(), rows) | |
| summary[table.name] = len(rows) | |
| for table in Base.metadata.sorted_tables: | |
| pk_columns = list(table.primary_key.columns) | |
| if len(pk_columns) == 1 and target_url.startswith("postgres"): | |
| reset_postgres_sequence(target_conn, table.name, pk_columns[0].name) | |
| with target_engine.connect() as conn: | |
| for table in Base.metadata.sorted_tables: | |
| if table_exists(target_engine, table.name): | |
| summary[f"{table.name}_target_count"] = conn.execute(select(func.count()).select_from(table)).scalar() | |
| return summary | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Migrate cert-study SQLite data to PostgreSQL.") | |
| parser.add_argument("--source", default=str(DB_PATH), help="SQLite DB path. Default: data/questions.db") | |
| parser.add_argument("--target", default=os.getenv("DATABASE_URL") or DEFAULT_TARGET_URL, help="PostgreSQL SQLAlchemy URL.") | |
| parser.add_argument("--replace", action="store_true", help="Delete target table rows before copying.") | |
| args = parser.parse_args() | |
| summary = migrate(Path(args.source), args.target, replace=args.replace) | |
| for key, value in summary.items(): | |
| print(f"{key}: {value}") | |
| if __name__ == "__main__": | |
| main() | |