battery-tracker-app/migrate_to_mariadb.py

"""
Migrate data from SQLite to MariaDB using SQLAlchemy ORM only — no raw SQL
except for the AUTO_INCREMENT reset which requires a DDL statement.

Usage:
    MARIADB_URL='mysql+pymysql://user:pass@host/dbname?charset=utf8mb4' \\
        python migrate_to_mariadb.py

Or pass the URL as a CLI argument:
    python migrate_to_mariadb.py 'mysql+pymysql://user:pass@host/dbname?charset=utf8mb4'
"""

import sys

from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

import config
from models import Base, Battery, Device


def migrate(mariadb_url: str):
    print("=== Battery Tracker: SQLite → MariaDB Migration ===\n")

    # -- Engines --
    sqlite_engine = create_engine(config.SQLALCHEMY_DATABASE_URI)
    mariadb_engine = create_engine(mariadb_url, pool_pre_ping=True)

    SrcSession = sessionmaker(bind=sqlite_engine)
    DstSession = sessionmaker(bind=mariadb_engine)

    src = SrcSession()
    dst = DstSession()

    try:
        # -- Create tables on MariaDB --
        print("Creating tables on MariaDB (if not exist)…")
        Base.metadata.create_all(mariadb_engine)

        # -- Read source data --
        src_devices   = src.query(Device).all()
        src_batteries = src.query(Battery).all()
        print(f"Source: {len(src_devices)} devices, {len(src_batteries)} batteries\n")

        # -- Migrate Devices first (batteries have FK → device) --
        print("Migrating devices…")
        for d in src_devices:
            new_d = Device(
                id=d.id,
                name=d.name,
                battery_slots=d.battery_slots,
                notes=d.notes,
            )
            dst.add(new_d)
        dst.flush()

        # -- Migrate Batteries --
        print("Migrating batteries…")
        for b in src_batteries:
            new_b = Battery(
                id=b.id,
                label=b.label,
                brand=b.brand,
                status=b.status,
                device_id=b.device_id,
                notes=b.notes,
            )
            dst.add(new_b)
        dst.flush()

        dst.commit()
        print("Commit successful.\n")

        # -- Reset AUTO_INCREMENT so new rows don't collide --
        # This DDL is required; there is no ORM-level equivalent.
        with mariadb_engine.connect() as conn:
            max_device_id   = max((d.id for d in src_devices),   default=0)
            max_battery_id  = max((b.id for b in src_batteries),  default=0)
            conn.execute(text("ALTER TABLE device SET AUTO_INCREMENT = :v"),
                         {"v": max_device_id + 1})
            conn.execute(text("ALTER TABLE battery SET AUTO_INCREMENT = :v"),
                         {"v": max_battery_id + 1})
            conn.commit()
        print("AUTO_INCREMENT counters reset.\n")

        # -- Verify counts --
        dst_device_count   = dst.query(Device).count()
        dst_battery_count  = dst.query(Battery).count()
        src_device_count   = len(src_devices)
        src_battery_count  = len(src_batteries)

        print("=== Verification ===")
        print(f"{'Table':<12} {'SQLite':>8} {'MariaDB':>9} {'OK?':>5}")
        print("-" * 38)

        device_ok  = src_device_count  == dst_device_count
        battery_ok = src_battery_count == dst_battery_count

        print(f"{'device':<12} {src_device_count:>8} {dst_device_count:>9} {'OK' if device_ok else 'MISMATCH':>5}")
        print(f"{'battery':<12} {src_battery_count:>8} {dst_battery_count:>9} {'OK' if battery_ok else 'MISMATCH':>5}")

        if not (device_ok and battery_ok):
            print("\nERROR: Row count mismatch. Do not decommission SQLite.")
            sys.exit(1)

        print("\nMigration complete. All row counts match.")

    except Exception as exc:
        dst.rollback()
        print(f"\nERROR: {exc}")
        raise
    finally:
        src.close()
        dst.close()


if __name__ == "__main__":
    if len(sys.argv) > 1:
        url = sys.argv[1]
    else:
        import os
        url = os.environ.get("MARIADB_URL")

    if not url:
        print("ERROR: Provide MariaDB URL via MARIADB_URL env var or as a CLI argument.")
        print("  Example: python migrate_to_mariadb.py 'mysql+pymysql://user:pass@host/db?charset=utf8mb4'")
        sys.exit(1)

    migrate(url)