yingshiyun/analysis/db.py

"""analysis.db（SQLite）建表 + CRUD，唯一直接写 SQL 的地方。"""
import sqlite3
from contextlib import contextmanager

from analysis.config import DB_PATH

SCHEMA = """
CREATE TABLE IF NOT EXISTS videos (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    device_serial TEXT NOT NULL,
    capture_date TEXT NOT NULL,
    file_path TEXT NOT NULL UNIQUE,
    file_name TEXT NOT NULL,
    file_size_bytes INTEGER,
    duration_sec REAL,
    fps REAL,
    width INTEGER,
    height INTEGER,
    start_time TEXT,
    status TEXT NOT NULL DEFAULT 'pending',
    motion_windows_json TEXT,
    error_message TEXT,
    analyzed_at TEXT,
    file_md5 TEXT,
    duplicate_of INTEGER REFERENCES videos(id)
);

CREATE TABLE IF NOT EXISTS events (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    video_id INTEGER NOT NULL REFERENCES videos(id),
    event_type TEXT NOT NULL,
    label TEXT,
    start_offset_sec REAL,
    end_offset_sec REAL,
    timestamp_abs TEXT,
    confidence REAL,
    track_id INTEGER,
    thumbnail_path TEXT,
    extra_json TEXT,
    created_at TEXT
);

CREATE INDEX IF NOT EXISTS idx_events_type_time ON events(event_type, timestamp_abs);
"""


@contextmanager
def get_connection():
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    # WAL 模式允许读写并发（analyze.py 写 + webapp 读不互锁），
    # busy_timeout 避免 "database is locked" 直接报错
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=30000")
    conn.execute("PRAGMA foreign_keys=ON")
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()


def init_db():
    with get_connection() as conn:
        conn.executescript(SCHEMA)
        _migrate_add_dedup_columns(conn)


def _migrate_add_dedup_columns(conn):
    """老库没有 file_md5/duplicate_of 字段，CREATE TABLE IF NOT EXISTS 对已存在的表不生效，
    这里用 ALTER TABLE 补齐（新建库这两列已经在 SCHEMA 里，ALTER 会因列已存在而报错，捕获忽略）。
    """
    existing_columns = {row["name"] for row in conn.execute("PRAGMA table_info(videos)")}
    if "file_md5" not in existing_columns:
        conn.execute("ALTER TABLE videos ADD COLUMN file_md5 TEXT")
    if "duplicate_of" not in existing_columns:
        conn.execute("ALTER TABLE videos ADD COLUMN duplicate_of INTEGER REFERENCES videos(id)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_videos_md5 ON videos(file_md5)")


def video_exists(file_path):
    with get_connection() as conn:
        row = conn.execute(
            "SELECT 1 FROM videos WHERE file_path = ?", (str(file_path),)
        ).fetchone()
        return row is not None


def insert_video(video):
    """插入一条 video 记录，file_path 已存在则忽略（保证重复扫描幂等）。返回该记录的 id。"""
    with get_connection() as conn:
        conn.execute(
            """
            INSERT OR IGNORE INTO videos (
                device_serial, capture_date, file_path, file_name,
                file_size_bytes, duration_sec, fps, width, height,
                start_time, status, error_message, file_md5, duplicate_of
            ) VALUES (
                :device_serial, :capture_date, :file_path, :file_name,
                :file_size_bytes, :duration_sec, :fps, :width, :height,
                :start_time, :status, :error_message, :file_md5, :duplicate_of
            )
            """,
            video,
        )
        row = conn.execute(
            "SELECT id FROM videos WHERE file_path = ?", (video["file_path"],)
        ).fetchone()
        return row["id"]


def get_video_id_by_md5(file_md5):
    with get_connection() as conn:
        row = conn.execute(
            "SELECT id FROM videos WHERE file_md5 = ?", (file_md5,)
        ).fetchone()
        return row["id"] if row else None


def get_videos_missing_md5():
    with get_connection() as conn:
        rows = conn.execute(
            "SELECT id, file_path FROM videos WHERE file_md5 IS NULL"
        ).fetchall()
        return [dict(r) for r in rows]


def update_video_md5(video_id, file_md5):
    with get_connection() as conn:
        conn.execute(
            "UPDATE videos SET file_md5 = ? WHERE id = ?", (file_md5, video_id)
        )


def count_videos():
    with get_connection() as conn:
        row = conn.execute("SELECT COUNT(*) AS c FROM videos").fetchone()
        return row["c"]


def get_all_videos(status=None):
    with get_connection() as conn:
        if status is None:
            rows = conn.execute("SELECT * FROM videos ORDER BY capture_date, file_name").fetchall()
        else:
            rows = conn.execute(
                "SELECT * FROM videos WHERE status = ? ORDER BY capture_date, file_name", (status,)
            ).fetchall()
        return [dict(r) for r in rows]


def update_video_status(video_id, status, error_message=None):
    with get_connection() as conn:
        conn.execute(
            "UPDATE videos SET status = ?, error_message = ? WHERE id = ?",
            (status, error_message, video_id),
        )


def update_start_time(video_id, start_time):
    with get_connection() as conn:
        conn.execute(
            "UPDATE videos SET start_time = ? WHERE id = ?",
            (start_time, video_id),
        )


def update_motion_windows(video_id, motion_windows_json, status, analyzed_at):
    with get_connection() as conn:
        conn.execute(
            "UPDATE videos SET motion_windows_json = ?, status = ?, analyzed_at = ? WHERE id = ?",
            (motion_windows_json, status, analyzed_at, video_id),
        )


def insert_event(event):
    with get_connection() as conn:
        conn.execute(
            """
            INSERT INTO events (
                video_id, event_type, label, start_offset_sec, end_offset_sec,
                timestamp_abs, confidence, track_id, thumbnail_path, extra_json, created_at
            ) VALUES (
                :video_id, :event_type, :label, :start_offset_sec, :end_offset_sec,
                :timestamp_abs, :confidence, :track_id, :thumbnail_path, :extra_json, :created_at
            )
            """,
            event,
        )


def get_video_by_id(video_id):
    with get_connection() as conn:
        row = conn.execute("SELECT * FROM videos WHERE id = ?", (video_id,)).fetchone()
        return dict(row) if row else None


def update_event_thumbnail(event_id, thumbnail_path):
    with get_connection() as conn:
        conn.execute(
            "UPDATE events SET thumbnail_path = ? WHERE id = ?",
            (thumbnail_path, event_id),
        )


def update_event_extra(event_id, extra_json):
    with get_connection() as conn:
        conn.execute(
            "UPDATE events SET extra_json = ? WHERE id = ?",
            (extra_json, event_id),
        )


def get_events(event_type=None):
    with get_connection() as conn:
        if event_type is None:
            rows = conn.execute("SELECT * FROM events ORDER BY timestamp_abs").fetchall()
        else:
            rows = conn.execute(
                "SELECT * FROM events WHERE event_type = ? ORDER BY timestamp_abs", (event_type,)
            ).fetchall()
        return [dict(r) for r in rows]