generated from kgod/ai-review-template
157 lines
6.2 KiB
Python
157 lines
6.2 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlalchemy import create_engine, select
|
|
from sqlalchemy.dialects.postgresql import insert as postgres_insert
|
|
|
|
from .schema import build_session_tables
|
|
|
|
|
|
TABLES = ("sessions", "turns", "blocks", "resumes", "resume_imports", "optimization_runs")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class MigrationReport:
|
|
source_counts: dict[str, int]
|
|
target_counts: dict[str, int]
|
|
source_checksum: str
|
|
target_checksum: str
|
|
|
|
|
|
def migrate_sqlite_to_postgres(
|
|
source_path: str | Path,
|
|
target_url: str,
|
|
*,
|
|
schema: str = "resume_agent",
|
|
dry_run: bool = False,
|
|
conflict_policy: str = "error",
|
|
) -> MigrationReport:
|
|
"""Copy SQLite rows into PostgreSQL without deleting target-only records."""
|
|
if conflict_policy not in {"error", "update"}:
|
|
raise ValueError("conflict_policy must be 'error' or 'update'")
|
|
source_rows = _read_source(Path(source_path))
|
|
source_counts = {name: len(rows) for name, rows in source_rows.items()}
|
|
source_checksum = _checksum(source_rows)
|
|
if dry_run:
|
|
return MigrationReport(source_counts, {}, source_checksum, "")
|
|
|
|
engine = create_engine(target_url)
|
|
metadata, tables = build_session_tables(schema)
|
|
try:
|
|
with engine.begin() as connection:
|
|
connection.exec_driver_sql(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')
|
|
metadata.create_all(connection)
|
|
for name in TABLES:
|
|
rows = source_rows[name]
|
|
if not rows:
|
|
continue
|
|
table = tables[name]
|
|
statement = postgres_insert(table).values(rows)
|
|
if conflict_policy == "update":
|
|
statement = statement.on_conflict_do_update(
|
|
index_elements=[table.c.id],
|
|
set_={column.name: statement.excluded[column.name]
|
|
for column in table.columns if column.name != "id"},
|
|
)
|
|
connection.execute(statement)
|
|
target_rows = _read_target(connection, tables, source_rows)
|
|
target_checksum = _checksum(target_rows)
|
|
if target_checksum != source_checksum:
|
|
raise RuntimeError("PostgreSQL verification checksum does not match SQLite source")
|
|
finally:
|
|
engine.dispose()
|
|
target_counts = {name: len(rows) for name, rows in target_rows.items()}
|
|
target_checksum = _checksum(target_rows)
|
|
return MigrationReport(source_counts, target_counts, source_checksum, target_checksum)
|
|
|
|
|
|
def _read_source(path: Path) -> dict[str, list[dict[str, Any]]]:
|
|
connection = sqlite3.connect(f"file:{path.resolve().as_posix()}?mode=ro", uri=True)
|
|
connection.row_factory = sqlite3.Row
|
|
try:
|
|
rows: dict[str, list[dict[str, Any]]] = {}
|
|
existing = {row[0] for row in connection.execute(
|
|
"SELECT name FROM sqlite_master WHERE type = 'table'"
|
|
)}
|
|
for name in TABLES:
|
|
if name not in existing:
|
|
rows[name] = []
|
|
continue
|
|
result = [dict(row) for row in connection.execute(f"SELECT * FROM {name}")]
|
|
rows[name] = [_normalize_source(name, row) for row in result]
|
|
return rows
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
def _normalize_source(name: str, row: dict[str, Any]) -> dict[str, Any]:
|
|
result = dict(row)
|
|
if name == "sessions":
|
|
result["profile"] = json.loads(result.pop("profile_json"))
|
|
if result["profile"].get("job_type") == "other":
|
|
result["profile"]["job_type"] = "internship"
|
|
elif name == "blocks":
|
|
result["data"] = json.loads(result.pop("data_json"))
|
|
elif name == "resumes":
|
|
result["content"] = json.loads(result.pop("content_json"))
|
|
elif name == "resume_imports":
|
|
raw_document = result.pop("document_json")
|
|
result["document"] = json.loads(raw_document) if raw_document else None
|
|
result["field_reviews"] = json.loads(result.pop("field_reviews_json"))
|
|
elif name == "optimization_runs":
|
|
result["state"] = json.loads(result.pop("state_json"))
|
|
raw_proposal = result.pop("proposal_json")
|
|
result["proposal"] = json.loads(raw_proposal) if raw_proposal else None
|
|
|
|
return result
|
|
|
|
def _read_target(
|
|
connection: Any, tables: dict[str, Any], source_rows: dict[str, list[dict[str, Any]]]
|
|
) -> dict[str, list[dict[str, Any]]]:
|
|
result: dict[str, list[dict[str, Any]]] = {}
|
|
for name in TABLES:
|
|
table = tables[name]
|
|
source_ids = [row["id"] for row in source_rows[name]]
|
|
if not source_ids:
|
|
result[name] = []
|
|
continue
|
|
rows = connection.execute(select(table).where(
|
|
table.c.id.in_(source_ids)
|
|
).order_by(table.c.id)).mappings().all()
|
|
result[name] = [dict(row) for row in rows]
|
|
return result
|
|
|
|
|
|
def _checksum(rows_by_table: dict[str, list[dict[str, Any]]]) -> str:
|
|
stable: dict[str, list[dict[str, Any]]] = {}
|
|
for name in TABLES:
|
|
rows = rows_by_table.get(name, [])
|
|
stable[name] = [_checksum_row(name, row) for row in rows]
|
|
stable[name].sort(key=lambda row: str(row["id"]))
|
|
serialized = json.dumps(stable, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
return hashlib.sha256(serialized.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _checksum_row(name: str, row: dict[str, Any]) -> dict[str, Any]:
|
|
keys = {
|
|
"sessions": ("id", "stage", "revision", "profile", "draft_id", "resume_id"),
|
|
"turns": ("id", "session_id", "sequence", "role", "content", "composer_mode"),
|
|
"blocks": ("id", "session_id", "turn_id", "block_index", "type", "lifecycle", "data", "version"),
|
|
"resumes": ("id", "session_id", "idempotency_key", "revision", "content"),
|
|
"resume_imports": (
|
|
"id", "session_id", "file_name", "mime_type", "size_bytes", "sha256",
|
|
"object_key", "status", "document", "field_reviews", "error_code",
|
|
),
|
|
"optimization_runs": (
|
|
"id", "session_id", "entry_id", "mode", "status", "source_revision",
|
|
"state", "proposal",
|
|
),
|
|
}[name]
|
|
return {key: row[key] for key in keys}
|