banban/talkingq-url/scripts/import_device_imei_mapping.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import pymysql
from pymysql.cursors import DictCursor

from config import settings


SERIAL_NUMBER_COLUMNS = ("serial_number", "device_sn")

CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS device_imei_mapping (
    id INT NOT NULL AUTO_INCREMENT,
    imei VARCHAR(64) NOT NULL,
    device_id VARCHAR(64) NOT NULL,
    serial_number VARCHAR(64) NOT NULL,
    status VARCHAR(32) NOT NULL DEFAULT 'pending',
    activated_at DATETIME NULL,
    created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
    updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    PRIMARY KEY (id),
    UNIQUE INDEX imei_UNIQUE (imei ASC),
    UNIQUE INDEX device_id_UNIQUE (device_id ASC),
    INDEX idx_device_imei_mapping_status (status ASC)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Import IMEI to device identity mappings from a CSV file."
    )
    parser.add_argument("csv_path", help="CSV file with imei, device_id, serial_number columns.")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate the CSV without writing to the database.",
    )
    return parser.parse_args()


def normalize_row(row: dict[str, str], row_number: int) -> dict[str, str]:
    serial_number = ""
    for column in SERIAL_NUMBER_COLUMNS:
        value = str(row.get(column) or "").strip()
        if value:
            serial_number = value
            break

    normalized = {
        "imei": str(row.get("imei") or "").strip(),
        "device_id": str(row.get("device_id") or "").strip(),
        "serial_number": serial_number,
    }
    missing = [key for key, value in normalized.items() if not value]
    if missing:
        raise ValueError(f"row {row_number}: missing required columns: {', '.join(missing)}")
    return normalized


def read_rows(csv_path: Path) -> list[dict[str, str]]:
    with csv_path.open("r", encoding="utf-8-sig", newline="") as fp:
        reader = csv.DictReader(fp)
        fieldnames = set(reader.fieldnames or [])
        missing = [column for column in ("imei", "device_id") if column not in fieldnames]
        if not any(column in fieldnames for column in SERIAL_NUMBER_COLUMNS):
            missing.append("serial_number/device_sn")
        if missing:
            raise ValueError(f"CSV missing required columns: {', '.join(missing)}")

        rows: list[dict[str, str]] = []
        seen_imei: set[str] = set()
        seen_device_id: set[str] = set()
        for row_number, row in enumerate(reader, start=2):
            normalized = normalize_row(row, row_number)
            if normalized["imei"] in seen_imei:
                raise ValueError(f"row {row_number}: duplicate imei in CSV: {normalized['imei']}")
            if normalized["device_id"] in seen_device_id:
                raise ValueError(f"row {row_number}: duplicate device_id in CSV: {normalized['device_id']}")
            seen_imei.add(normalized["imei"])
            seen_device_id.add(normalized["device_id"])
            rows.append(normalized)
        return rows


def get_connection():
    return pymysql.connect(
        host=settings.db_host,
        port=settings.db_port,
        user=settings.db_user,
        password=settings.db_password,
        database=settings.db_name,
        charset="utf8mb4",
        cursorclass=DictCursor,
    )


def import_rows(rows: list[dict[str, str]]) -> int:
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(CREATE_TABLE_SQL)
            for row in rows:
                cur.execute(
                    """
                    SELECT imei, device_id, serial_number, status
                    FROM device_imei_mapping
                    WHERE imei = %s OR device_id = %s
                    """,
                    (row["imei"], row["device_id"]),
                )
                existing_rows = cur.fetchall()
                if len(existing_rows) > 1:
                    raise ValueError(
                        f"conflicting mapping: imei={row['imei']} device_id={row['device_id']}"
                    )

                if not existing_rows:
                    cur.execute(
                        """
                        INSERT INTO device_imei_mapping (imei, device_id, serial_number, status)
                        VALUES (%s, %s, %s, 'pending')
                        """,
                        (row["imei"], row["device_id"], row["serial_number"]),
                    )
                    continue

                existing = existing_rows[0]
                if existing["imei"] != row["imei"]:
                    raise ValueError(
                        f"device_id already mapped to another imei: {row['device_id']}"
                    )

                is_activated = str(existing["status"]) == "activated"
                changed_identity = (
                    existing["device_id"] != row["device_id"]
                    or existing["serial_number"] != row["serial_number"]
                )
                if is_activated and changed_identity:
                    raise ValueError(f"activated imei cannot be remapped: {row['imei']}")

                cur.execute(
                    """
                    UPDATE device_imei_mapping
                    SET device_id = %s,
                        serial_number = %s,
                        status = IF(status = 'activated', status, 'pending')
                    WHERE imei = %s
                    """,
                    (row["device_id"], row["serial_number"], row["imei"]),
                )
        conn.commit()
        return len(rows)
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def main() -> int:
    args = parse_args()
    csv_path = Path(args.csv_path).expanduser().resolve()
    try:
        rows = read_rows(csv_path)
        if args.dry_run:
            print(f"validated {len(rows)} rows from {csv_path}")
            return 0
        imported = import_rows(rows)
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    print(f"imported {imported} rows into device_imei_mapping")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())