lzwlab_cn/transfer/storage.py

"""文件元数据与磁盘缓存的存储层（SQLite + 文件目录，仅标准库）。"""
import os
import sqlite3
import threading
from datetime import datetime, timezone

_DB_LOCK = threading.Lock()

SCHEMA = """
CREATE TABLE IF NOT EXISTS files (
    id          TEXT PRIMARY KEY,
    name        TEXT NOT NULL,
    stored_path TEXT NOT NULL,
    size        INTEGER NOT NULL,
    created_at  INTEGER NOT NULL,
    expires_at  INTEGER NOT NULL,
    downloads   INTEGER NOT NULL DEFAULT 0
);
"""


def connect(db_path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(db_path, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=30000")
    return conn


def init_db(db_path: str) -> None:
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            conn.executescript(SCHEMA)
            conn.commit()
        finally:
            conn.close()


def iso(ts: int) -> str:
    return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()


def file_to_dict(row: sqlite3.Row, now: int) -> dict:
    return {
        "id": row["id"],
        "name": row["name"],
        "size": row["size"],
        "created_at": iso(row["created_at"]),
        "expires_at": iso(row["expires_at"]),
        "remaining_seconds": max(0, row["expires_at"] - now),
        "downloads": row["downloads"],
    }


def insert_file(db_path: str, file_id: str, name: str, stored_path: str,
                size: int, expires_at: int, created_at: int) -> None:
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            conn.execute(
                "INSERT INTO files (id, name, stored_path, size, created_at, expires_at, downloads) "
                "VALUES (?, ?, ?, ?, ?, ?, 0)",
                (file_id, name, stored_path, size, created_at, expires_at),
            )
            conn.commit()
        finally:
            conn.close()


def get_file(db_path: str, file_id: str):
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            return conn.execute("SELECT * FROM files WHERE id = ?", (file_id,)).fetchone()
        finally:
            conn.close()


def list_files(db_path: str, now: int) -> list[dict]:
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            rows = conn.execute(
                "SELECT * FROM files WHERE expires_at > ? ORDER BY created_at DESC", (now,)
            ).fetchall()
            return [file_to_dict(r, now) for r in rows]
        finally:
            conn.close()


def increment_downloads(db_path: str, file_id: str) -> None:
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            conn.execute("UPDATE files SET downloads = downloads + 1 WHERE id = ?", (file_id,))
            conn.commit()
        finally:
            conn.close()


def delete_file(db_path: str, files_dir: str, file_id: str) -> bool:
    """删除数据库记录并移除磁盘文件，返回是否删除过数据库记录。"""
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            row = conn.execute("SELECT id, stored_path FROM files WHERE id = ?", (file_id,)).fetchone()
            if row is None:
                return False
            conn.execute("DELETE FROM files WHERE id = ?", (file_id,))
            conn.commit()
        finally:
            conn.close()
        try:
            os.remove(row["stored_path"])
        except FileNotFoundError:
            pass
        except OSError:
            pass
        return True


def cleanup_expired(db_path: str, files_dir: str, now: int | None = None) -> int:
    now = now if now is not None else int(__import__("time").time())
    removed = 0
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            rows = conn.execute(
                "SELECT id, stored_path FROM files WHERE expires_at <= ?", (now,)
            ).fetchall()
            for row in rows:
                try:
                    os.remove(row["stored_path"])
                except FileNotFoundError:
                    pass
                except OSError:
                    continue  # 磁盘文件暂时删不掉时保留记录，下次重试
                conn.execute("DELETE FROM files WHERE id = ?", (row["id"],))
                removed += 1
            conn.commit()
        finally:
            conn.close()
    return removed


def cleanup_orphans(db_path: str, files_dir: str, now: int | None = None,
                    min_age_seconds: int = 172800) -> int:
    """清理磁盘上不在数据库且超过 48 小时未修改的孤儿文件。"""
    import time as _time
    now = now if now is not None else int(_time.time())
    removed = 0
    with _DB_LOCK:
        conn = connect(db_path)
        try:
            ids = {r["id"] for r in conn.execute("SELECT id FROM files").fetchall()}
        finally:
            conn.close()
    if not os.path.isdir(files_dir):
        return 0
    for name in os.listdir(files_dir):
        if name in ids:
            continue
        path = os.path.join(files_dir, name)
        if not os.path.isfile(path):
            continue
        try:
            if now - os.path.getmtime(path) >= min_age_seconds:
                os.remove(path)
                removed += 1
        except OSError:
            pass
    return removed