src/storage.py

"""存储层：数据目录与 SQLite 封装。"""

from __future__ import annotations

import os
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# 注意：日志仍然复用主网关 logger 名称，方便统一过滤
import logging

logger = logging.getLogger("gateway")


# ── 路径与基本配置 ─────────────────────────────────────────────────────────────
DATA_DIR = Path(os.environ.get("DATA_DIR", "/data"))
UPLOAD_DIR = DATA_DIR / "uploads"
JOB_DIR = DATA_DIR / "jobs"
DB_PATH = DATA_DIR / "gateway.db"


_db_lock = threading.Lock()
_db_conn: sqlite3.Connection | None = None


def now_iso() -> str:
    """返回当前 UTC 时间的 ISO 字符串。"""
    return datetime.now(timezone.utc).isoformat()


def ensure_data_dirs() -> None:
    """确保数据目录存在。"""
    UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
    JOB_DIR.mkdir(parents=True, exist_ok=True)


def init_db() -> None:
    """初始化 SQLite 数据库与基础表结构。"""
    global _db_conn

    ensure_data_dirs()
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    with conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                username TEXT NOT NULL,
                filename TEXT NOT NULL,
                input_path TEXT NOT NULL,
                output_dir TEXT NOT NULL,
                status TEXT NOT NULL,
                progress REAL NOT NULL DEFAULT 0,
                message TEXT,
                error TEXT,
                model TEXT NOT NULL,
                lang_in TEXT NOT NULL,
                lang_out TEXT NOT NULL,
                cancel_requested INTEGER NOT NULL DEFAULT 0,
                mono_pdf_path TEXT,
                dual_pdf_path TEXT,
                glossary_path TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                started_at TEXT,
                finished_at TEXT
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS usage_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL,
                job_id TEXT,
                model TEXT NOT NULL,
                prompt_tokens INTEGER NOT NULL,
                completion_tokens INTEGER NOT NULL,
                total_tokens INTEGER NOT NULL,
                cost_usd REAL NOT NULL,
                created_at TEXT NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_jobs_user_time
            ON jobs(username, created_at DESC)
            """
        )
        conn.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_usage_user_time
            ON usage_records(username, created_at DESC)
            """
        )

    _db_conn = conn
    logger.info("Database initialized at %s", DB_PATH)


def close_db() -> None:
    """关闭数据库连接，用于应用关闭阶段。"""
    global _db_conn

    if _db_conn is not None:
        _db_conn.close()
        _db_conn = None


def db_execute(sql: str, params: tuple[Any, ...] = ()) -> None:
    """执行写操作 SQL。"""
    if _db_conn is None:
        raise RuntimeError("DB is not initialized")
    with _db_lock, _db_conn:
        _db_conn.execute(sql, params)


def db_fetchone(sql: str, params: tuple[Any, ...] = ()) -> sqlite3.Row | None:
    """执行查询并返回单行。"""
    if _db_conn is None:
        raise RuntimeError("DB is not initialized")
    with _db_lock:
        return _db_conn.execute(sql, params).fetchone()


def db_fetchall(sql: str, params: tuple[Any, ...] = ()) -> list[sqlite3.Row]:
    """执行查询并返回多行。"""
    if _db_conn is None:
        raise RuntimeError("DB is not initialized")
    with _db_lock:
        return _db_conn.execute(sql, params).fetchall()