snapshot

scan_vids.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "psycopg[binary]",
+# ]
+# ///
+import argparse
+import contextlib
+import hashlib
+import json
+import logging
+import subprocess
+import typing
+from pathlib import Path
+
+import psycopg
+import psycopg.rows
+
+import ffmpeg
+
+
+def hash_partial(f: Path) -> str:
+    sha1 = hashlib.sha1()
+    chunk_size = 1024 * 1024 * 10  # 10MB chunk size
+
+    total_read = 0
+    with f.open("rb") as file:
+        while chunk := file.read(chunk_size):
+            total_read += len(chunk)
+            sha1.update(chunk)
+            break  # Only reads the first 10MB
+
+    return f"sha1:{total_read}:{sha1.hexdigest()}"
+
+
+@contextlib.contextmanager
+def connect_db() -> typing.Generator[psycopg.Connection, typing.Any, typing.Any]:
+    with psycopg.connect("postgres://abdus:abdus@db.abdus.dev:5444/smut?sslmode=disable") as conn:
+        conn.row_factory = psycopg.rows.dict_row
+        yield conn
+
+
+def video_exists(con: psycopg.Connection, file_path: str) -> bool:
+    stmt = con.execute(
+        "UPDATE videos SET last_seen_at = NOW() WHERE file_path = %s RETURNING true AS exists",
+        [file_path],
+    )
+    row = stmt.fetchone()
+    return row["exists"]
+
+
+def delete_video(con: psycopg.Connection, file_path: str) -> bool:
+    con.execute("DELETE FROM videos WHERE file_path = %s", [file_path])
+
+
+def save_video(
+    con: psycopg.Connection,
+    category: str,
+    file_path: str,
+    ffprobe: dict,
+    hash_partial: str,
+):
+    sql = """INSERT INTO videos(category, file_path, ffprobe, hash_partial) VALUES (%s, %s, %s, %s) ON CONFLICT(file_path) SET last_seen_at = NOW();"""
+    con.execute(
+        sql,
+        [
+            category,
+            file_path,
+            json.dumps(ffprobe),
+            hash_partial,
+        ],
+    )
+
+
+def save_error(con: psycopg.Connection, file_path: str, stderr: str):
+    sql = """INSERT INTO scan_state(file_path, stderr) VALUES (%s, %s) ON conflict do nothing"""
+    con.execute(
+        sql,
+        [
+            file_path,
+            stderr,
+        ],
+    )
+
+
+def scan_videos(con: psycopg.Connection, video_paths: list[Path]) -> None:
+    for i, f in enumerate(video_paths):
+        if "/_picks/" in str(f):
+            continue
+        progress = f"[{i + 1}/{len(video_paths)}]"
+        if f.is_dir():
+            continue
+
+        file_path = str(f.absolute())
+
+        if video_exists(con, file_path=file_path):
+            logging.debug(f"{progress} video already exists {f=}")
+            continue
+
+        logging.debug(f"probing {f=}")
+        try:
+            probed = ffmpeg.ffprobe(f)
+        except subprocess.CalledProcessError as e:
+            save_error(con, file_path=file_path, stderr=e.stderr)
+            continue
+
+        file_hash = hash_partial(f)
+
+        save_video(
+            category=category_for_file(f),
+            con=con,
+            file_path=file_path,
+            ffprobe=probed.__dict__,
+            hash_partial=file_hash,
+        )
+        logging.info(f"{progress} saved {f=}")
+        con.commit()
+
+
+def category_for_file(file_path: Path) -> str:
+    if str(file_path).startswith("/Volumes/"):
+        return f"downloaded.{file_path.parts[1].lower()}"
+    if str(file_path).startswith("/mnt/box"):
+        return "box"
+    raise ValueError(f"unknown category for {file_path}")
+
+
+def find_copies(con: psycopg.Connection, video_path: Path) -> list[str]:
+    probed = ffmpeg.ffprobe(video_path)
+
+    sql = """
+SELECT
+    file_path
+FROM videos
+WHERE
+    file_path != %(file_path)s
+    AND ( 
+        hash_partial = %(hash_partial)s
+        OR (abs(duration_sec - %(duration_sec)s) < 0.1 AND abs(size_bytes - %(size_bytes)s) < 3000000)
+    );
+   """
+
+    stmt: psycopg.Cursor = con.execute(
+        sql,
+        {
+            "file_path": str(video_path.absolute()),
+            "hash_partial": hash_partial(video_path),
+            "duration_sec": probed.duration_sec,
+            "size_bytes": probed.size_bytes,
+        },
+    )
+    rows = stmt.fetchall()
+
+    return [it["file_path"] for it in rows]
+
+
+def parse_args():
+    arger = argparse.ArgumentParser()
+    arger.add_argument("video_paths", nargs="+", type=Path, help="Path to the videos")
+    arger.add_argument("--debug", action="store_true", help="Path to the videos")
+
+    return arger.parse_args()
+
+
+def main():
+    logging.basicConfig(format=logging.BASIC_FORMAT, level=logging.INFO)
+    args = parse_args()
+    if args.debug:
+        logging.getLogger().setLevel(logging.DEBUG)
+
+    video_paths: list[Path] = args.video_paths
+    if not video_paths:
+        logging.info("no files to scan")
+        return
+
+    with connect_db() as con:
+        scan_videos(con=con, video_paths=video_paths)
+
+
+if __name__ == "__main__":
+    main()