7z-encrypt-server/db.py

"""
入库层 (db)
模块: 服务端 / 入库层
输入: 文件路径 + 元数据
输出: file_id

PostgreSQL 表:
    transfers       任务主表 (状态机 uploading->assembling->decrypting->done/failed)
    transfer_chunks 每卷信息 + 接收时间 (null = 未收, 断点续传依据)
    files           明文文件记录
"""

import uuid
from typing import Any, cast

import psycopg
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from settings import DATABASE_URL


class ServerDB:
    """服务端数据库访问 (psycopg3, 每次操作独立连接)"""

    def __init__(self, url: str = DATABASE_URL) -> None:
        self.url = url
        self._init_schema()

    def _conn(self):
        return psycopg.connect(self.url, row_factory=cast(Any, dict_row))

    def _init_schema(self) -> None:
        with self._conn() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS transfers (
                    transfer_id  TEXT PRIMARY KEY,
                    file_name    TEXT NOT NULL,
                    file_size    BIGINT NOT NULL,
                    chunk_count  INTEGER NOT NULL,
                    total_sha256 TEXT NOT NULL,
                    enc_params   JSONB NOT NULL,
                    status       TEXT NOT NULL,
                    created_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
                    updated_at   TIMESTAMPTZ NOT NULL DEFAULT now()
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS transfer_chunks (
                    transfer_id TEXT NOT NULL REFERENCES transfers(transfer_id),
                    idx         INTEGER NOT NULL,
                    size        BIGINT NOT NULL,
                    sha256      TEXT NOT NULL,
                    received_at TIMESTAMPTZ,
                    PRIMARY KEY (transfer_id, idx)
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS files (
                    file_id     TEXT PRIMARY KEY,
                    transfer_id TEXT NOT NULL REFERENCES transfers(transfer_id),
                    file_name   TEXT NOT NULL,
                    path        TEXT NOT NULL,
                    size        BIGINT NOT NULL,
                    sha256      TEXT NOT NULL,
                    created_at  TIMESTAMPTZ NOT NULL DEFAULT now()
                )
                """
            )

    # ---------- transfers ----------

    def create_transfer(self, transfer_id: str, init_json: dict[str, Any]) -> None:
        enc = init_json["enc"]
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO transfers (transfer_id, file_name, file_size, chunk_count, total_sha256, enc_params, status) "
                "VALUES (%s, %s, %s, %s, %s, %s, %s)",
                (
                    transfer_id,
                    init_json["file_name"],
                    init_json["file_size"],
                    init_json["chunk_count"],
                    init_json["total_sha256"],
                    Jsonb(enc),
                    "uploading",
                ),
            )
            conn.executemany(
                "INSERT INTO transfer_chunks (transfer_id, idx, size, sha256) VALUES (%s, %s, %s, %s)",
                [
                    (transfer_id, ch["index"], ch["size"], ch["sha256"])
                    for ch in init_json["chunks"]
                ],
            )

    def get_transfer(self, transfer_id: str) -> dict[str, Any] | None:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM transfers WHERE transfer_id = %s", (transfer_id,)
            ).fetchone()
        if not row:
            return None
        return dict(row)

    def set_status(self, transfer_id: str, status: str) -> None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE transfers SET status = %s, updated_at = now() WHERE transfer_id = %s",
                (status, transfer_id),
            )

    # ---------- chunks ----------

    def get_chunk_spec(self, transfer_id: str, idx: int) -> dict[str, Any] | None:
        """卷规格 (size/sha256), 接收校验用"""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT size, sha256 FROM transfer_chunks WHERE transfer_id = %s AND idx = %s",
                (transfer_id, idx),
            ).fetchone()
        return dict(row) if row else None

    def is_chunk_received(self, transfer_id: str, idx: int) -> bool:
        with self._conn() as conn:
            row = conn.execute(
                "SELECT 1 FROM transfer_chunks WHERE transfer_id = %s AND idx = %s AND received_at IS NOT NULL",
                (transfer_id, idx),
            ).fetchone()
        return row is not None

    def mark_chunk_received(self, transfer_id: str, idx: int) -> None:
        with self._conn() as conn:
            conn.execute(
                "UPDATE transfer_chunks SET received_at = now() "
                "WHERE transfer_id = %s AND idx = %s",
                (transfer_id, idx),
            )

    def get_received_chunks(self, transfer_id: str) -> set[int]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT idx FROM transfer_chunks WHERE transfer_id = %s AND received_at IS NOT NULL",
                (transfer_id,),
            ).fetchall()
        return {r["idx"] for r in rows}

    # ---------- files ----------

    def insert_file(
        self,
        transfer_id: str,
        file_name: str,
        path: str,
        size: int,
        sha256: str,
    ) -> str:
        file_id = uuid.uuid4().hex[:12]
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO files (file_id, transfer_id, file_name, path, size, sha256) "
                "VALUES (%s, %s, %s, %s, %s, %s)",
                (file_id, transfer_id, file_name, path, size, sha256),
            )
        return file_id