scripts/migrate_add_pid_source.py

"""
Schema migration: add pid_source column to items table, and backfill it for
existing rows. Idempotent.
"""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import inspect, text

from uraas.database import engine


def column_exists(table: str, column: str) -> bool:
    insp = inspect(engine)
    return column in {c["name"] for c in insp.get_columns(table)}


def main() -> int:
    print("Migration: adding pid_source to items")

    dialect = engine.dialect.name
    print(f"Dialect: {dialect}")

    if not column_exists("items", "pid_source"):
        with engine.begin() as conn:
            print("  -> ALTER TABLE items ADD COLUMN pid_source VARCHAR(20)")
            conn.execute(text("ALTER TABLE items ADD COLUMN pid_source VARCHAR(20)"))
    else:
        print("  pid_source already present, skipping ALTER")

    with engine.begin() as conn:
        r1 = conn.execute(
            text(
                "UPDATE items SET pid_source = 'handle' "
                "WHERE pid_source IS NULL "
                "AND source_repository LIKE '%IR (OAI-PMH)%' "
                "AND url LIKE '%/handle/%'"
            )
        )
        print(f"  -> backfilled pid_source='handle' on {r1.rowcount} rows")

        r2 = conn.execute(
            text(
                "UPDATE items SET pid_source = 'ark' "
                "WHERE pid_source IS NULL AND ark IS NOT NULL"
            )
        )
        print(f"  -> backfilled pid_source='ark' on {r2.rowcount} rows")

    print("Done.")
    return 0


if __name__ == "__main__":
    sys.exit(main())