From 680fbaaa5fd2d205cf127c5f9de13b6fcb47536f Mon Sep 17 00:00:00 2001 From: lou Date: Mon, 10 Aug 2026 00:21:22 +0800 Subject: [PATCH] =?UTF-8?q?db=20=E5=B1=82=E6=8D=A2=20SQLite:=20PRoot=20?= =?UTF-8?q?=E5=AE=B9=E5=99=A8=20PG=20fsync=20=E5=8D=A1=E6=AD=BB,=20?= =?UTF-8?q?=E8=BD=BB=E9=87=8F=E6=96=87=E4=BB=B6=E5=BA=93=20(=E8=A1=A8?= =?UTF-8?q?=E7=BB=93=E6=9E=84=E5=85=BC=E5=AE=B9,=20=E5=8F=AF=E5=88=87?= =?UTF-8?q?=E5=9B=9E=20PG)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- db.py | 87 +++++++++++++++++++++++++---------------------- settings.py | 7 ++-- tests/test_api.py | 12 +++---- 3 files changed, 55 insertions(+), 51 deletions(-) diff --git a/db.py b/db.py index 5a18864..16fec61 100644 --- a/db.py +++ b/db.py @@ -4,46 +4,53 @@ 输入: 文件路径 + 元数据 输出: file_id -PostgreSQL 表: +SQLite 实现 (老板拍板: PRoot 容器装 PG/MySQL 服务进程都会 fsync 卡死, +轻量文件库最稳)。表结构与 PostgreSQL 版一致, 未来可平滑切回 PG。 + +表: transfers 任务主表 (状态机 uploading->assembling->decrypting->done/failed) transfer_chunks 每卷信息 + 接收时间 (null = 未收, 断点续传依据) files 明文文件记录 """ +import json +import sqlite3 import uuid -from typing import Any, cast +from pathlib import Path +from typing import Any -import psycopg -from psycopg.rows import dict_row -from psycopg.types.json import Jsonb - -from settings import DATABASE_URL +from settings import DB_PATH class ServerDB: - """服务端数据库访问 (psycopg3, 每次操作独立连接)""" + """服务端数据库访问 (sqlite3, 每次操作独立连接, WAL 并发安全)""" - def __init__(self, url: str = DATABASE_URL) -> None: - self.url = url + def __init__(self, db_path: str | Path = DB_PATH) -> None: + self.db_path = Path(db_path) self._init_schema() - def _conn(self) -> Any: - return psycopg.connect(self.url, row_factory=cast(Any, dict_row)) + def _conn(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.db_path) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=5000") + return conn def _init_schema(self) -> None: + self.db_path.parent.mkdir(parents=True, exist_ok=True) 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, + file_size INTEGER NOT NULL, chunk_count INTEGER NOT NULL, total_sha256 TEXT NOT NULL, - enc_params JSONB NOT NULL, + enc_params TEXT NOT NULL, status TEXT NOT NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) ) """ ) @@ -52,9 +59,9 @@ class ServerDB: CREATE TABLE IF NOT EXISTS transfer_chunks ( transfer_id TEXT NOT NULL REFERENCES transfers(transfer_id), idx INTEGER NOT NULL, - size BIGINT NOT NULL, + size INTEGER NOT NULL, sha256 TEXT NOT NULL, - received_at TIMESTAMPTZ, + received_at TEXT, PRIMARY KEY (transfer_id, idx) ) """ @@ -66,9 +73,9 @@ class ServerDB: transfer_id TEXT NOT NULL REFERENCES transfers(transfer_id), file_name TEXT NOT NULL, path TEXT NOT NULL, - size BIGINT NOT NULL, + size INTEGER NOT NULL, sha256 TEXT NOT NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT now() + created_at TEXT NOT NULL DEFAULT (datetime('now')) ) """ ) @@ -76,43 +83,43 @@ class ServerDB: # ---------- 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)", + "VALUES (?, ?, ?, ?, ?, ?, ?)", ( transfer_id, init_json["file_name"], init_json["file_size"], init_json["chunk_count"], init_json["total_sha256"], - Jsonb(enc), + json.dumps(init_json["enc"], ensure_ascii=False), "uploading", ), ) - with conn.cursor() as cur: - cur.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"] - ], - ) + conn.executemany( + "INSERT INTO transfer_chunks (transfer_id, idx, size, sha256) VALUES (?, ?, ?, ?)", + [ + (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,) + "SELECT * FROM transfers WHERE transfer_id = ?", (transfer_id,) ).fetchone() if not row: return None - return dict(row) + d = dict(row) + d["enc_params"] = json.loads(d["enc_params"]) + return d 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", + "UPDATE transfers SET status = ?, updated_at = datetime('now') WHERE transfer_id = ?", (status, transfer_id), ) @@ -122,7 +129,7 @@ class ServerDB: """卷规格 (size/sha256), 接收校验用""" with self._conn() as conn: row = conn.execute( - "SELECT size, sha256 FROM transfer_chunks WHERE transfer_id = %s AND idx = %s", + "SELECT size, sha256 FROM transfer_chunks WHERE transfer_id = ? AND idx = ?", (transfer_id, idx), ).fetchone() return dict(row) if row else None @@ -130,7 +137,7 @@ class ServerDB: 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", + "SELECT 1 FROM transfer_chunks WHERE transfer_id = ? AND idx = ? AND received_at IS NOT NULL", (transfer_id, idx), ).fetchone() return row is not None @@ -138,15 +145,15 @@ class ServerDB: 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", + "UPDATE transfer_chunks SET received_at = datetime('now') " + "WHERE transfer_id = ? AND idx = ?", (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", + "SELECT idx FROM transfer_chunks WHERE transfer_id = ? AND received_at IS NOT NULL", (transfer_id,), ).fetchall() return {r["idx"] for r in rows} @@ -165,7 +172,7 @@ class ServerDB: 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)", + "VALUES (?, ?, ?, ?, ?, ?)", (file_id, transfer_id, file_name, path, size, sha256), ) return file_id diff --git a/settings.py b/settings.py index 59e1985..fc2814a 100644 --- a/settings.py +++ b/settings.py @@ -5,11 +5,8 @@ from pathlib import Path # 应用绑定上下文 (必须与客户端一致) CONTEXT = os.environ.get("SZ_CONTEXT", "7z-encrypt:v1").encode() -# PostgreSQL 连接串 (本地默认走笔记本 PG socket) -DATABASE_URL = os.environ.get( - "DATABASE_URL", - "postgresql://lou@/7zencrypt?host=/home/lou/pgdata/socket", -) +# SQLite 数据库文件 (轻量部署: PRoot 容器装 PG/MySQL 会 fsync 卡死) +DB_PATH = Path(os.environ.get("SZ_DB_PATH", "data/app.db")) # 存储根目录 (明文落盘) STORAGE_ROOT = Path(os.environ.get("SZ_STORAGE_ROOT", "server/data/storage")) diff --git a/tests/test_api.py b/tests/test_api.py index bee32b5..9e6a89b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -17,7 +17,7 @@ import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # 测试用配置 (先于 api 导入) -os.environ["DATABASE_URL"] = "postgresql://lou@/7zencrypt_test?host=/home/lou/pgdata/socket" +os.environ["SZ_DB_PATH"] = "/tmp/sz-test.db" os.environ["SZ_TMP_ROOT"] = "/tmp/sz-test-tmp" os.environ["SZ_STORAGE_ROOT"] = "/tmp/sz-test-storage" os.environ["SZ_KEYRING"] = "/tmp/sz-test-keyring.json" @@ -57,10 +57,10 @@ def _prepare(keyring_path: str, data: bytes, chunk_size: int = 1 << 20): class ServerPipelineTest(unittest.TestCase): @classmethod def setUpClass(cls): - # 清空测试库 + # 清空测试库 (SQLite: 删文件重建, 代替 TRUNCATE) + for f in ("/tmp/sz-test.db", "/tmp/sz-test.db-wal", "/tmp/sz-test.db-shm"): + os.path.exists(f) and os.unlink(f) db = ServerDB() - with db._conn() as conn: - conn.execute("TRUNCATE transfers, transfer_chunks, files CASCADE") # 清空运行时目录 + 旧密钥文件 for d in ("/tmp/sz-test-tmp", "/tmp/sz-test-storage"): shutil.rmtree(d, ignore_errors=True) @@ -104,7 +104,7 @@ class ServerPipelineTest(unittest.TestCase): db = ServerDB() row = None with db._conn() as conn: - row = conn.execute("SELECT path, size FROM files WHERE file_id = %s", (file_id,)).fetchone() + row = conn.execute("SELECT path, size FROM files WHERE file_id = ?", (file_id,)).fetchone() self.assertIsNotNone(row) with open(row["path"], "rb") as f: restored = f.read() @@ -113,7 +113,7 @@ class ServerPipelineTest(unittest.TestCase): # transfer 状态 done with db._conn() as conn: - st = conn.execute("SELECT status FROM transfers WHERE transfer_id = %s", (tid,)).fetchone() + st = conn.execute("SELECT status FROM transfers WHERE transfer_id = ?", (tid,)).fetchone() self.assertEqual(st["status"], "done") def test_bad_chunk_rejected(self):