Autoparts-DB/scripts/warm_vehicle_cache.py

#!/usr/bin/env python3
"""Warm Redis cache for vehicle info.

Runs in batches over all parts in the catalog, populating
nexus:vehicle:{part_id} keys in Redis. This eliminates the
DISTINCT ON + 4 JOINs query on vehicle_parts (2B rows) for
cached parts.

Usage:
    python3 warm_vehicle_cache.py
    python3 warm_vehicle_cache.py --dsn "postgresql://user:pass@localhost/db"
    python3 warm_vehicle_cache.py --batch-size 10000 --ttl 7200
"""

import argparse
import os
import subprocess
import sys
import time
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'pos'))

import psycopg2
import redis

DEFAULT_DSN = os.environ.get('MASTER_DB_URL', 'postgresql://postgres@/nexus_autoparts')
REDIS_URL = os.environ.get('REDIS_URL', 'redis://localhost:6379/0')
DEFAULT_BATCH_SIZE = 5000
DEFAULT_TTL = 3600


def log(msg):
    print(f"[{datetime.now().isoformat(timespec='seconds')}] {msg}", flush=True)


def _connect(dsn):
    """Connect to PostgreSQL; raise on failure."""
    return psycopg2.connect(dsn)


def _ensure_connection(dsn):
    """Try to connect. On peer-auth failure, re-run with sudo -u postgres."""
    try:
        return _connect(dsn)
    except psycopg2.OperationalError as exc:
        err = str(exc).lower()
        if 'peer' in err or 'authentication' in err:
            if os.geteuid() == 0:
                # Already root — can't sudo to postgres usefully; give clear message
                log("ERROR: PostgreSQL peer authentication failed.")
                log("       Run as postgres OS user:")
                log("         sudo -u postgres python3 " + __file__)
                log("       Or set MASTER_DB_URL with TCP host+password:")
                log("         export MASTER_DB_URL=postgresql://user:pass@localhost/nexus_autoparts")
                sys.exit(1)
            log("Peer auth failed. Re-running with sudo -u postgres ...")
            cmd = ['sudo', '-u', 'postgres', sys.executable, __file__]
            # Forward original env + CLI args
            env = os.environ.copy()
            env['MASTER_DB_URL'] = dsn
            for i, arg in enumerate(sys.argv[1:], start=1):
                if arg in ('--dsn', '-d') and i < len(sys.argv) - 1:
                    env['MASTER_DB_URL'] = sys.argv[i + 1]
            ret = subprocess.call(cmd, env=env)
            sys.exit(ret)
        raise


def main():
    parser = argparse.ArgumentParser(description='Warm Redis cache for vehicle info')
    parser.add_argument('--dsn', '-d', default=DEFAULT_DSN,
                        help='PostgreSQL DSN (default: MASTER_DB_URL env or peer auth)')
    parser.add_argument('--batch-size', '-b', type=int, default=DEFAULT_BATCH_SIZE,
                        help=f'Batch size (default: {DEFAULT_BATCH_SIZE})')
    parser.add_argument('--ttl', '-t', type=int, default=DEFAULT_TTL,
                        help=f'Redis TTL in seconds (default: {DEFAULT_TTL})')
    parser.add_argument('--dry-run', action='store_true',
                        help='Do not write to Redis, just log what would be done')
    args = parser.parse_args()

    log("Connecting to master DB and Redis...")
    conn = _ensure_connection(args.dsn)
    cur = conn.cursor()
    r = redis.from_url(REDIS_URL, decode_responses=True)
    r.ping()
    log("Connected.")

    # Get all part_ids
    cur.execute("SELECT id_part FROM parts WHERE oem_part_number IS NOT NULL ORDER BY id_part")
    all_ids = [row[0] for row in cur.fetchall()]
    total = len(all_ids)
    log(f"Total parts to warm: {total}")

    if total == 0:
        log("No parts found. Exiting.")
        return

    processed = 0
    cached = 0
    start = time.time()

    for i in range(0, total, args.batch_size):
        batch = all_ids[i:i + args.batch_size]
        cur.execute("""
            SELECT part_id, name_brand, name_model, year_car
            FROM part_vehicle_preview
            WHERE part_id = ANY(%s)
        """, (batch,))

        rows = cur.fetchall()
        if not args.dry_run:
            pipe = r.pipeline()
            for row in rows:
                info = f"{row[1]} {row[2]} {row[3]}"
                pipe.setex(f'nexus:vehicle:{row[0]}', args.ttl, info)
            pipe.execute()

        batch_cached = len(rows)
        processed += len(batch)
        cached += batch_cached
        elapsed = time.time() - start
        rate = processed / elapsed if elapsed > 0 else 0
        log(f"[{processed}/{total}] cached={batch_cached} ({rate:.0f}/s)")

    cur.close()
    conn.close()
    elapsed = time.time() - start
    log(f"Done. Cached {cached} vehicle entries in {elapsed:.0f}s")


if __name__ == '__main__':
    main()