alembic/pipeline/lib/build-fingerprint-index.py

#!/usr/bin/env python3
"""
build-fingerprint-index.py — build/refresh a Chromaprint fingerprint index of
the beets library at $ALEMBIC_CONFIG_DIR/beets/fingerprints.db.

Reads paths from `beet ls -f '$id|||$path'` (beets now runs in-process in this
same container, so no container-to-host path translation is needed — the path
beets reports is already valid in this container's own filesystem view).
Skips entries whose (path, mtime, size) is unchanged. Parallel fpcalc workers
keep this under ~10 min for a 3k-track library.

Used by import-dj-collection.py to dedup incoming DJ files against the existing
library via acoustid.compare_fingerprints().

Usage:
    build-fingerprint-index.py [--workers N] [--limit N]
"""

import argparse
import multiprocessing
import os
import sqlite3
import subprocess
import sys
import time

ALEMBIC_CONFIG_DIR = os.environ.get("ALEMBIC_CONFIG_DIR", "/config")
DB_PATH = f"{ALEMBIC_CONFIG_DIR}/beets/fingerprints.db"
DEFAULT_WORKERS = 8


def init_db(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS fingerprints (
            beets_id INTEGER PRIMARY KEY,
            path TEXT NOT NULL,
            mtime REAL NOT NULL,
            size INTEGER NOT NULL,
            duration INTEGER NOT NULL,
            fingerprint TEXT NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_fp_path ON fingerprints(path);
        """
    )
    conn.commit()


def enumerate_library() -> list[tuple[int, str]]:
    """Return [(beets_id, path), ...] for every track in the library."""
    cmd = ["beet", "ls", "-f", "$id|||$path"]
    out = subprocess.run(cmd, capture_output=True, text=True, check=True).stdout
    items: list[tuple[int, str]] = []
    for line in out.splitlines():
        line = line.rstrip("\r")
        if "|||" not in line:
            continue
        id_str, path = line.split("|||", 1)
        try:
            beets_id = int(id_str)
        except ValueError:
            continue
        items.append((beets_id, path))
    return items


def fpcalc(host_path: str) -> tuple[int, str] | None:
    """Run fpcalc -raw on a file. Returns (duration_seconds, fingerprint_csv)."""
    try:
        proc = subprocess.run(
            ["fpcalc", "-raw", host_path],
            capture_output=True,
            text=True,
            timeout=120,
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None
    if proc.returncode != 0:
        return None
    duration = 0
    fingerprint = ""
    for line in proc.stdout.splitlines():
        if line.startswith("DURATION="):
            try:
                duration = int(line.split("=", 1)[1])
            except ValueError:
                pass
        elif line.startswith("FINGERPRINT="):
            fingerprint = line.split("=", 1)[1]
    if not fingerprint:
        return None
    return duration, fingerprint


def worker(job: tuple[int, str, float, int]) -> tuple[int, str, float, int, int, str] | None:
    beets_id, host_path, mtime, size = job
    result = fpcalc(host_path)
    if result is None:
        return None
    duration, fingerprint = result
    return beets_id, host_path, mtime, size, duration, fingerprint


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--workers", type=int, default=DEFAULT_WORKERS)
    ap.add_argument("--limit", type=int, default=0, help="process at most N new files (0=all)")
    ap.add_argument("--verbose", action="store_true")
    args = ap.parse_args()

    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    init_db(conn)

    print(f"[index] enumerating library via beets...", flush=True)
    library = enumerate_library()
    print(f"[index] library has {len(library)} tracks", flush=True)

    existing: dict[int, tuple[str, float, int]] = {}
    for row in conn.execute("SELECT beets_id, path, mtime, size FROM fingerprints"):
        existing[row[0]] = (row[1], row[2], row[3])

    jobs: list[tuple[int, str, float, int]] = []
    missing_count = 0
    for beets_id, host_path in library:
        try:
            st = os.stat(host_path)
        except FileNotFoundError:
            missing_count += 1
            continue
        cur = existing.get(beets_id)
        if cur is not None and cur[0] == host_path and abs(cur[1] - st.st_mtime) < 1 and cur[2] == st.st_size:
            continue
        jobs.append((beets_id, host_path, st.st_mtime, st.st_size))

    print(f"[index] {len(jobs)} new/changed tracks to fingerprint "
          f"({len(library) - len(jobs) - missing_count} already current, {missing_count} missing on disk)", flush=True)
    if args.limit > 0:
        jobs = jobs[:args.limit]
        print(f"[index] --limit {args.limit} applied; processing {len(jobs)}", flush=True)
    if not jobs:
        conn.close()
        return 0

    purge_ids = set(existing.keys()) - {b for b, *_ in library}
    if purge_ids:
        conn.executemany("DELETE FROM fingerprints WHERE beets_id = ?", [(i,) for i in purge_ids])
        conn.commit()
        print(f"[index] purged {len(purge_ids)} stale rows", flush=True)

    start = time.time()
    done = 0
    failed = 0
    batch: list[tuple] = []
    with multiprocessing.Pool(args.workers) as pool:
        for result in pool.imap_unordered(worker, jobs, chunksize=4):
            done += 1
            if result is None:
                failed += 1
            else:
                batch.append(result)
            if len(batch) >= 50:
                conn.executemany(
                    "INSERT OR REPLACE INTO fingerprints "
                    "(beets_id, path, mtime, size, duration, fingerprint) VALUES (?,?,?,?,?,?)",
                    batch,
                )
                conn.commit()
                batch.clear()
            if done % 100 == 0 or args.verbose:
                elapsed = time.time() - start
                rate = done / elapsed if elapsed > 0 else 0
                print(f"[index] {done}/{len(jobs)} ({rate:.1f}/s, {failed} failed)", flush=True)
    if batch:
        conn.executemany(
            "INSERT OR REPLACE INTO fingerprints "
            "(beets_id, path, mtime, size, duration, fingerprint) VALUES (?,?,?,?,?,?)",
            batch,
        )
        conn.commit()

    total = conn.execute("SELECT COUNT(*) FROM fingerprints").fetchone()[0]
    conn.close()
    elapsed = time.time() - start
    print(f"[index] done in {elapsed:.0f}s. wrote {done - failed} rows, {failed} failed. "
          f"index now has {total} entries.", flush=True)
    return 0 if failed == 0 else 1


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