oem_crawler/scripts/migrate_raw_payloads.py

#!/usr/bin/env python3
"""
Migration script to extract raw OEM payloads from SQLite into external compressed .json.gz files,
update vehicles with raw_payload_path, clear in-database raw JSON, and compact the database via VACUUM.
"""
from __future__ import annotations

import argparse
import gzip
import json
import logging
import os
import sqlite3
import sys
import time
from pathlib import Path
from typing import Any, Dict, Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("migrate_raw_payloads")


def format_bytes(num_bytes: int | float) -> str:
    """Formats bytes into human-readable string (KB, MB, GB)."""
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if abs(num_bytes) < 1024.0:
            return f"{num_bytes:.2f} {unit}"
        num_bytes /= 1024.0
    return f"{num_bytes:.2f} PB"


def sanitize_filename(name: str) -> str:
    """Sanitizes an OEM vehicle id for filesystem safety."""
    return "".join(c if c.isalnum() or c in ("-", "_") else "_" for c in str(name))


def migrate_database(
    db_path: Path,
    raw_dir: Path,
    compression_level: int = 6,
    batch_size: int = 500,
    dry_run: bool = False,
    skip_vacuum: bool = False,
) -> None:
    if not db_path.exists():
        logger.error(f"Database file not found: {db_path}")
        sys.exit(1)

    initial_db_size = db_path.stat().st_size
    logger.info(f"Target Database: {db_path} (Size: {format_bytes(initial_db_size)})")
    logger.info(f"Storage Directory: {raw_dir.resolve()}")
    logger.info(f"Dry Run Mode: {dry_run}")

    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()

    # 1. Ensure raw_payload_path column exists
    cursor.execute("PRAGMA table_info(vehicles);")
    columns = [row["name"] for row in cursor.fetchall()]
    if "raw_payload_path" not in columns:
        if dry_run:
            logger.info("[Dry Run] Would execute: ALTER TABLE vehicles ADD COLUMN raw_payload_path VARCHAR(512);")
        else:
            logger.info("Adding 'raw_payload_path' column to vehicles table...")
            cursor.execute("ALTER TABLE vehicles ADD COLUMN raw_payload_path VARCHAR(512);")
            conn.commit()

    # 2. Count vehicles requiring migration
    cursor.execute("""
        SELECT COUNT(*) as count 
        FROM vehicles 
        WHERE raw_payload IS NOT NULL AND length(raw_payload) > 2;
    """)
    total_vehicles = cursor.fetchone()["count"]
    logger.info(f"Found {total_vehicles:,} vehicles with raw payload to extract.")

    cursor.execute("""
        SELECT COUNT(*) as count 
        FROM vehicle_options 
        WHERE raw_data IS NOT NULL AND raw_data != '{}';
    """)
    total_options = cursor.fetchone()["count"]
    logger.info(f"Found {total_options:,} vehicle options with redundant raw_data.")

    if total_vehicles == 0 and total_options == 0:
        logger.info("No vehicles or options require migration.")
        if not skip_vacuum and not dry_run:
            logger.info("Running VACUUM to ensure database is compact...")
            cursor.execute("VACUUM;")
        conn.close()
        return

    if dry_run:
        logger.info("[Dry Run] Migration simulation complete. No changes written.")
        conn.close()
        return

    # 3. Process vehicles in batches
    raw_dir.mkdir(parents=True, exist_ok=True)
    migrated_count = 0
    total_uncompressed_bytes = 0
    total_compressed_bytes = 0
    start_time = time.time()

    cursor.execute("""
        SELECT id, brand, oem_vehicle_id, created_at, last_seen_at, raw_payload, raw_payload_path
        FROM vehicles 
        WHERE raw_payload IS NOT NULL AND length(raw_payload) > 2;
    """)

    updates_batch = []
    
    while True:
        rows = cursor.fetchmany(batch_size)
        if not rows:
            break

        for row in rows:
            v_id = row["id"]
            brand = (row["brand"] or "unknown").strip().lower()
            oem_id = row["oem_vehicle_id"] or v_id
            created_at = str(row["last_seen_at"] or row["created_at"] or "")
            date_str = created_at[:10] if len(created_at) >= 10 else "archive"

            raw_str = row["raw_payload"]
            if not raw_str or raw_str == "{}":
                continue

            try:
                payload = json.loads(raw_str) if isinstance(raw_str, str) else raw_str
            except Exception as e:
                logger.warning(f"Could not parse JSON for vehicle {v_id}: {e}")
                continue

            # Compute relative path
            safe_id = sanitize_filename(oem_id)
            rel_path = f"{brand}/{date_str}/{safe_id}.json.gz"
            dest_file = raw_dir / rel_path
            dest_file.parent.mkdir(parents=True, exist_ok=True)

            # Compress and write
            json_bytes = json.dumps(payload, ensure_ascii=False).encode("utf-8")
            compressed_data = gzip.compress(json_bytes, compresslevel=compression_level)
            
            with open(dest_file, "wb") as f:
                f.write(compressed_data)

            total_uncompressed_bytes += len(json_bytes)
            total_compressed_bytes += len(compressed_data)
            migrated_count += 1

            updates_batch.append((rel_path, v_id))

        # Commit DB updates for batch
        if updates_batch:
            update_cursor = conn.cursor()
            update_cursor.executemany(
                "UPDATE vehicles SET raw_payload_path = ?, raw_payload = '{}' WHERE id = ?;",
                updates_batch,
            )
            conn.commit()
            updates_batch.clear()

        elapsed = time.time() - start_time
        rate = migrated_count / elapsed if elapsed > 0 else 0
        logger.info(
            f"Progress: {migrated_count:,}/{total_vehicles:,} vehicles ({migrated_count/total_vehicles*100:.1f}%) "
            f"| Rate: {rate:.0f} veh/sec"
        )

    # 4. Clean vehicle_options raw_data
    logger.info("Clearing redundant raw_data in vehicle_options...")
    cursor.execute("UPDATE vehicle_options SET raw_data = '{}' WHERE raw_data != '{}';")
    conn.commit()

    # 5. Compact the database with VACUUM
    if not skip_vacuum:
        logger.info("Running VACUUM to reclaim free space (this may take a few moments)...")
        vac_start = time.time()
        cursor.execute("VACUUM;")
        logger.info(f"VACUUM completed in {time.time() - vac_start:.1f}s.")

    conn.close()

    final_db_size = db_path.stat().st_size
    reclaimed_bytes = initial_db_size - final_db_size

    # Summary
    logger.info("\n" + "=" * 60)
    logger.info("MIGRATION COMPLETED SUCCESSFULLY")
    logger.info("=" * 60)
    logger.info(f"Vehicles Migrated:         {migrated_count:,}")
    logger.info(f"Options Cleaned:          {total_options:,}")
    logger.info(f"Uncompressed Payload Size:{format_bytes(total_uncompressed_bytes)}")
    logger.info(f"Compressed File Size:     {format_bytes(total_compressed_bytes)} "
                f"({total_compressed_bytes / max(total_uncompressed_bytes, 1) * 100:.1f}%)")
    logger.info(f"Original DB Size:         {format_bytes(initial_db_size)}")
    logger.info(f"New DB Size:              {format_bytes(final_db_size)}")
    logger.info(f"Space Reclaimed:          {format_bytes(reclaimed_bytes)} "
                f"({reclaimed_bytes / max(initial_db_size, 1) * 100:.1f}% reduction)")
    logger.info("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="Extract raw JSON payloads from SQLite to external compressed files")
    parser.add_argument(
        "--db-path",
        type=Path,
        default=Path("data/crawler_stock.db"),
        help="Path to SQLite database file (default: data/crawler_stock.db)",
    )
    parser.add_argument(
        "--raw-dir",
        type=Path,
        default=Path("data/raw"),
        help="Destination directory for compressed JSON files (default: data/raw)",
    )
    parser.add_argument(
        "--compression-level",
        type=int,
        default=6,
        help="Gzip compression level (1-9, default: 6)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=500,
        help="Batch size for processing and DB updates (default: 500)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Inspect database without writing files or modifying database",
    )
    parser.add_argument(
        "--skip-vacuum",
        action="store_true",
        help="Skip database VACUUM at end of migration",
    )
    args = parser.parse_args()

    migrate_database(
        db_path=args.db_path,
        raw_dir=args.raw_dir,
        compression_level=args.compression_level,
        batch_size=args.batch_size,
        dry_run=args.dry_run,
        skip_vacuum=args.skip_vacuum,
    )


if __name__ == "__main__":
    main()