Spaces:
Sleeping
Sleeping
File size: 3,814 Bytes
9381502 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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()
|