117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
"""
|
|
Shared helpers for all Phase 1 SQLite → Postgres migration scripts.
|
|
|
|
Usage in each script:
|
|
from migration.utils import open_sqlite, get_pg, log_run, parse_dt, parse_json
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import aiosqlite
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
|
|
from config import settings
|
|
|
|
|
|
# ── SQLite ────────────────────────────────────────────────────────────────────
|
|
|
|
async def open_sqlite() -> aiosqlite.Connection:
|
|
"""Open the SQLite database (read-only; no writes during migration)."""
|
|
db_path = Path(settings.sqlite_db_path)
|
|
if not db_path.exists():
|
|
print(f"ERROR: SQLite database not found at {db_path.resolve()}", file=sys.stderr)
|
|
sys.exit(1)
|
|
conn = await aiosqlite.connect(str(db_path))
|
|
conn.row_factory = aiosqlite.Row
|
|
return conn
|
|
|
|
|
|
# ── Postgres ──────────────────────────────────────────────────────────────────
|
|
|
|
def _make_pg_session() -> async_sessionmaker:
|
|
engine = create_async_engine(settings.database_url, pool_size=5, echo=False)
|
|
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
|
|
AsyncPgSession = _make_pg_session()
|
|
|
|
|
|
# ── Type helpers ──────────────────────────────────────────────────────────────
|
|
|
|
def parse_dt(value: str | None) -> datetime | None:
|
|
"""Parse a SQLite TEXT timestamp → timezone-aware datetime (UTC)."""
|
|
if not value:
|
|
return None
|
|
for fmt in (
|
|
"%Y-%m-%dT%H:%M:%S.%f",
|
|
"%Y-%m-%dT%H:%M:%S",
|
|
"%Y-%m-%d %H:%M:%S.%f",
|
|
"%Y-%m-%d %H:%M:%S",
|
|
"%Y-%m-%d",
|
|
):
|
|
try:
|
|
dt = datetime.strptime(value, fmt)
|
|
return dt.replace(tzinfo=timezone.utc)
|
|
except ValueError:
|
|
continue
|
|
# ISO format with offset — let fromisoformat handle it
|
|
try:
|
|
dt = datetime.fromisoformat(value)
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
except ValueError:
|
|
pass
|
|
print(f"WARNING: could not parse timestamp {value!r} — using now()", file=sys.stderr)
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def parse_json(value: str | None, default=None):
|
|
"""Parse a SQLite TEXT JSON column → Python object."""
|
|
if value is None:
|
|
return default
|
|
try:
|
|
return json.loads(value)
|
|
except (json.JSONDecodeError, TypeError):
|
|
return default
|
|
|
|
|
|
# ── Migration run log ─────────────────────────────────────────────────────────
|
|
|
|
async def log_run(
|
|
script_name: str,
|
|
source_rows: int,
|
|
dest_rows: int,
|
|
success: bool = True,
|
|
notes: str | None = None,
|
|
) -> None:
|
|
"""Insert a row into _migration_runs recording this script's execution."""
|
|
async with AsyncPgSession() as session:
|
|
await session.execute(
|
|
text("""
|
|
INSERT INTO _migration_runs
|
|
(script_name, ran_at, source_rows, dest_rows, success, notes)
|
|
VALUES
|
|
(:script_name, now(), :source_rows, :dest_rows, :success, :notes)
|
|
"""),
|
|
{
|
|
"script_name": script_name,
|
|
"source_rows": source_rows,
|
|
"dest_rows": dest_rows,
|
|
"success": "ok" if success else "error",
|
|
"notes": notes,
|
|
},
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
# ── Count helper ──────────────────────────────────────────────────────────────
|
|
|
|
async def pg_count(session: AsyncSession, table: str) -> int:
|
|
row = await session.execute(text(f"SELECT COUNT(*) FROM {table}"))
|
|
return row.scalar()
|