439 lines
17 KiB
Python
439 lines
17 KiB
Python
"""中心服务端数据层（psycopg 3）。
|
||
|
||
和桌面端数据层的区别：这里不按 school_id 过滤，因为中心要看所有学校的东西；
|
||
学校隔离体现在发布范围（scope）和上报凭证上。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import secrets
|
||
import time
|
||
from contextlib import contextmanager
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import psycopg
|
||
from psycopg.rows import dict_row
|
||
|
||
from server.config import settings
|
||
from shared import knowledge_match
|
||
|
||
SCHEMA_PATH = Path(__file__).parent / "schema.sql"
|
||
VERSION_KEY = "publish_version"
|
||
|
||
|
||
# ---------- 连接 ----------
|
||
|
||
|
||
def _connect(dbname: str | None = None) -> psycopg.Connection:
|
||
return psycopg.connect(
|
||
host=settings.pg_host,
|
||
port=settings.pg_port,
|
||
user=settings.pg_user,
|
||
password=settings.pg_password,
|
||
dbname=dbname or settings.pg_database,
|
||
connect_timeout=5,
|
||
row_factory=dict_row,
|
||
options="-c timezone=Asia/Shanghai",
|
||
)
|
||
|
||
|
||
@contextmanager
|
||
def get_conn():
|
||
conn = _connect()
|
||
try:
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
conn.rollback()
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _fetchone(conn, sql: str, params: tuple | list = ()) -> dict | None:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, params)
|
||
return cur.fetchone()
|
||
|
||
|
||
def _fetchall(conn, sql: str, params: tuple | list = ()) -> list[dict]:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, params)
|
||
return list(cur.fetchall())
|
||
|
||
|
||
def _execute(conn, sql: str, params: tuple | list = ()) -> int:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, params)
|
||
return cur.rowcount
|
||
|
||
|
||
def _insert_id(conn, sql: str, params: tuple | list = ()) -> int:
|
||
with conn.cursor() as cur:
|
||
cur.execute(sql, params)
|
||
row = cur.fetchone()
|
||
return int(row["id"]) if row else 0
|
||
|
||
|
||
# ---------- 初始化 ----------
|
||
|
||
|
||
def init_db(retries: int = 40, delay: float = 1.5) -> None:
|
||
last_err: Exception | None = None
|
||
for attempt in range(1, retries + 1):
|
||
try:
|
||
_ensure_database()
|
||
with get_conn() as conn:
|
||
_apply_schema(conn)
|
||
return
|
||
except psycopg.Error as e:
|
||
last_err = e
|
||
if attempt == 1:
|
||
print(f"[hub-db] 等待 PostgreSQL {settings.pg_host}:{settings.pg_port} 就绪 ... ({e})", flush=True)
|
||
time.sleep(delay)
|
||
raise RuntimeError(
|
||
f"无法连接 PostgreSQL {settings.pg_host}:{settings.pg_port}。请先执行 docker compose up -d。原始错误：{last_err}"
|
||
)
|
||
|
||
|
||
def _ensure_database() -> None:
|
||
"""中心库单独一个 database，不存在就建（CREATE DATABASE 得在维护库里执行且不能在事务里）。"""
|
||
conn = _connect(settings.pg_maintenance_database)
|
||
try:
|
||
conn.autocommit = True
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (settings.pg_database,))
|
||
if cur.fetchone() is None:
|
||
cur.execute(f'CREATE DATABASE "{settings.pg_database}"')
|
||
print(f"[hub-db] 已创建中心库 {settings.pg_database}", flush=True)
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _strip_comment(line: str) -> str:
|
||
"""去掉行尾 -- 注释。语句要拼成一行执行，留着注释会把后面的列定义一起注释掉。"""
|
||
in_quote = False
|
||
for i, ch in enumerate(line):
|
||
if ch == "'":
|
||
in_quote = not in_quote
|
||
elif ch == "-" and not in_quote and line[i : i + 2] == "--":
|
||
return line[:i]
|
||
return line
|
||
|
||
|
||
def _apply_schema(conn) -> None:
|
||
sql = SCHEMA_PATH.read_text(encoding="utf-8")
|
||
for stmt in sql.split(";"):
|
||
body = " ".join(_strip_comment(line) for line in stmt.splitlines()).strip()
|
||
if body:
|
||
_execute(conn, body)
|
||
|
||
|
||
def _fmt(value: Any) -> Any:
|
||
return value.strftime("%Y-%m-%d %H:%M:%S") if isinstance(value, datetime) else value
|
||
|
||
|
||
def _row(row: dict) -> dict:
|
||
return {k: _fmt(v) for k, v in row.items()}
|
||
|
||
|
||
# ---------- 学校与凭证 ----------
|
||
|
||
|
||
def hash_token(token: str) -> str:
|
||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def register_school(school_id: str, name: str = "", note: str = "") -> dict:
|
||
"""新建或重置一所学校的上报凭证。明文 token 只在这里返回一次，库里只存哈希。"""
|
||
token = f"sa_{secrets.token_urlsafe(24)}"
|
||
with get_conn() as conn:
|
||
_execute(
|
||
conn,
|
||
"""INSERT INTO schools (id, name, token_hash, note) VALUES (%s, %s, %s, %s)
|
||
ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, token_hash = EXCLUDED.token_hash,
|
||
note = EXCLUDED.note""",
|
||
(school_id, name or school_id, hash_token(token), note),
|
||
)
|
||
return {"school_id": school_id, "name": name or school_id, "token": token}
|
||
|
||
|
||
def authenticate(school_id: str, token: str) -> bool:
|
||
if not school_id or not token:
|
||
return False
|
||
with get_conn() as conn:
|
||
row = _fetchone(conn, "SELECT token_hash FROM schools WHERE id = %s", (school_id,))
|
||
if not row or not secrets.compare_digest(row["token_hash"], hash_token(token)):
|
||
return False
|
||
_execute(conn, "UPDATE schools SET last_seen_at = NOW() WHERE id = %s", (school_id,))
|
||
return True
|
||
|
||
|
||
def list_schools() -> list[dict]:
|
||
with get_conn() as conn:
|
||
rows = _fetchall(
|
||
conn,
|
||
"""SELECT s.*,
|
||
(SELECT count(*) FROM hub_entry_sources es WHERE es.school_id = s.id) AS reported_entries
|
||
FROM schools s ORDER BY s.created_at""",
|
||
)
|
||
return [_row(r) for r in rows]
|
||
|
||
|
||
def delete_school(school_id: str) -> None:
|
||
with get_conn() as conn:
|
||
_execute(conn, "DELETE FROM schools WHERE id = %s", (school_id,))
|
||
|
||
|
||
# ---------- 上报（桌面端 → 中心） ----------
|
||
|
||
|
||
def _scope_rows(conn, knowledge_id: str, business_id: str) -> list[dict]:
|
||
return _fetchall(
|
||
conn,
|
||
"""SELECT id, "key", "value", category, status, observations, school_count, reviewed_by
|
||
FROM hub_entries WHERE knowledge_id = %s AND business_id = %s""",
|
||
(knowledge_id, business_id),
|
||
)
|
||
|
||
|
||
def receive_drafts(school_id: str, entries: list[dict]) -> dict:
|
||
"""收下一所学校报来的草稿，按内容合并进中心队列。
|
||
|
||
已被驳回的条目不会因为再次上报而复活（那所学校可能还没拉到驳回结果）。
|
||
"""
|
||
stats = {"created": 0, "merged": 0, "skipped": 0, "revived": 0}
|
||
if not entries:
|
||
return stats
|
||
with get_conn() as conn:
|
||
for raw in entries:
|
||
key = str(raw.get("key") or "").strip()[:128]
|
||
value = str(raw.get("value") or "").strip()
|
||
if not key or not value:
|
||
stats["skipped"] += 1
|
||
continue
|
||
knowledge_id = str(raw.get("knowledge_id") or "").strip()
|
||
business_id = str(raw.get("business_id") or "").strip()
|
||
category = knowledge_match.normalise_category(raw.get("category"))
|
||
observations = max(1, int(raw.get("observations") or 1))
|
||
source_task_id = str(raw.get("source_task_id") or "")[:64]
|
||
|
||
rows = _scope_rows(conn, knowledge_id, business_id)
|
||
match = next((r for r in rows if knowledge_match.same_knowledge(key, value, r["key"], r["value"])), None)
|
||
|
||
if match is None:
|
||
entry_id = _insert_id(
|
||
conn,
|
||
"""INSERT INTO hub_entries
|
||
(knowledge_id, business_id, "key", "value", category, observations, school_count, origin_school_id)
|
||
VALUES (%s, %s, %s, %s, %s, %s, 1, %s) RETURNING id""",
|
||
(knowledge_id, business_id, key, value, category, observations, school_id),
|
||
)
|
||
stats["created"] += 1
|
||
else:
|
||
if match["status"] == "rejected" and (match.get("reviewed_by") or "human") != "auto":
|
||
# 人工驳回是终局，不让它因为再次上报而复活；
|
||
# 自动驳回只是暂缓，再被观察到就退回队列重新排队，避免机器误判永久吃掉一条知识
|
||
stats["skipped"] += 1
|
||
continue
|
||
if match["status"] == "rejected":
|
||
_execute(
|
||
conn,
|
||
"""UPDATE hub_entries SET status = 'draft', reviewed_by = NULL, reviewed_at = NULL,
|
||
version = 0, suggestion = NULL, suggestion_note = '' WHERE id = %s""",
|
||
(match["id"],),
|
||
)
|
||
stats["revived"] += 1
|
||
entry_id = match["id"]
|
||
richer = value if len(value) > len(match["value"] or "") else match["value"]
|
||
category = knowledge_match.pick_category(category, match["category"])
|
||
_execute(
|
||
conn,
|
||
"""UPDATE hub_entries SET "value" = %s, category = %s, updated_at = NOW() WHERE id = %s""",
|
||
(richer, category, entry_id),
|
||
)
|
||
stats["merged"] += 1
|
||
|
||
confirmed = max(0, int(raw.get("confirmed") or 0))
|
||
contradicted = max(0, int(raw.get("contradicted") or 0))
|
||
_execute(
|
||
conn,
|
||
"""INSERT INTO hub_entry_sources (entry_id, school_id, source_task_id, observations, confirmed, contradicted)
|
||
VALUES (%s, %s, %s, %s, %s, %s)
|
||
ON CONFLICT (entry_id, school_id) DO UPDATE
|
||
SET observations = GREATEST(hub_entry_sources.observations, EXCLUDED.observations),
|
||
confirmed = GREATEST(hub_entry_sources.confirmed, EXCLUDED.confirmed),
|
||
contradicted = GREATEST(hub_entry_sources.contradicted, EXCLUDED.contradicted),
|
||
source_task_id = EXCLUDED.source_task_id, reported_at = NOW()""",
|
||
(entry_id, school_id, source_task_id, observations, confirmed, contradicted),
|
||
)
|
||
_execute(
|
||
conn,
|
||
"""UPDATE hub_entries SET
|
||
observations = (SELECT COALESCE(sum(observations), 1) FROM hub_entry_sources WHERE entry_id = %s),
|
||
school_count = (SELECT count(*) FROM hub_entry_sources WHERE entry_id = %s),
|
||
confirmed = (SELECT COALESCE(sum(confirmed), 0) FROM hub_entry_sources WHERE entry_id = %s),
|
||
contradicted = (SELECT COALESCE(sum(contradicted), 0) FROM hub_entry_sources WHERE entry_id = %s)
|
||
WHERE id = %s""",
|
||
(entry_id, entry_id, entry_id, entry_id, entry_id),
|
||
)
|
||
_execute(conn, "UPDATE schools SET last_report_at = NOW() WHERE id = %s", (school_id,))
|
||
_execute(
|
||
conn,
|
||
"INSERT INTO sync_log (school_id, direction, count, detail) VALUES (%s, 'report', %s, %s)",
|
||
(school_id, len(entries), f"新建 {stats['created']}、合并 {stats['merged']}、跳过 {stats['skipped']}"),
|
||
)
|
||
return stats
|
||
|
||
|
||
# ---------- 发布（中心 → 桌面端） ----------
|
||
|
||
|
||
def _next_version(conn) -> int:
|
||
row = _fetchone(conn, 'SELECT "value" FROM hub_meta WHERE "key" = %s', (VERSION_KEY,))
|
||
version = int(row["value"]) + 1 if row and str(row["value"]).isdigit() else 1
|
||
_execute(
|
||
conn,
|
||
"""INSERT INTO hub_meta ("key", "value") VALUES (%s, %s)
|
||
ON CONFLICT ("key") DO UPDATE SET "value" = EXCLUDED."value" """,
|
||
(VERSION_KEY, str(version)),
|
||
)
|
||
return version
|
||
|
||
|
||
def published_for(school_id: str, since: int = 0) -> dict:
|
||
"""某所学校该拿到的已发布条目：发布范围是 all 的，或范围是 origin 且由它自己报的。"""
|
||
with get_conn() as conn:
|
||
rows = _fetchall(
|
||
conn,
|
||
"""SELECT e.id, e.knowledge_id, e.business_id, e."key", e."value", e.category, e.version, e.status
|
||
FROM hub_entries e
|
||
WHERE e.version > %s AND e.status IN ('verified', 'rejected')
|
||
AND (e.scope = 'all'
|
||
OR EXISTS (SELECT 1 FROM hub_entry_sources s WHERE s.entry_id = e.id AND s.school_id = %s))
|
||
ORDER BY e.version LIMIT 500""",
|
||
(int(since), school_id),
|
||
)
|
||
latest = _fetchone(conn, 'SELECT "value" FROM hub_meta WHERE "key" = %s', (VERSION_KEY,))
|
||
if rows:
|
||
_execute(conn, "UPDATE schools SET last_pull_at = NOW() WHERE id = %s", (school_id,))
|
||
_execute(
|
||
conn,
|
||
"INSERT INTO sync_log (school_id, direction, count, detail) VALUES (%s, 'pull', %s, %s)",
|
||
(school_id, len(rows), f"从版本 {since} 起"),
|
||
)
|
||
version = int(latest["value"]) if latest and str(latest["value"]).isdigit() else 0
|
||
return {"version": version, "entries": [_row(r) for r in rows]}
|
||
|
||
|
||
# ---------- 审核台 ----------
|
||
|
||
|
||
def list_entries(
|
||
status: str | None = None,
|
||
knowledge_id: str | None = None,
|
||
business_id: str | None = None,
|
||
category: str | None = None,
|
||
school_id: str | None = None,
|
||
query: str | None = None,
|
||
limit: int = 200,
|
||
) -> list[dict]:
|
||
where: list[str] = []
|
||
params: list[Any] = []
|
||
if status:
|
||
where.append("e.status = %s")
|
||
params.append(status)
|
||
if knowledge_id:
|
||
where.append("e.knowledge_id = %s")
|
||
params.append(knowledge_id)
|
||
if business_id:
|
||
where.append("e.business_id = %s")
|
||
params.append(business_id)
|
||
if category:
|
||
where.append("e.category = %s")
|
||
params.append(category)
|
||
if school_id:
|
||
where.append("EXISTS (SELECT 1 FROM hub_entry_sources s WHERE s.entry_id = e.id AND s.school_id = %s)")
|
||
params.append(school_id)
|
||
if query:
|
||
where.append('(e."key" ILIKE %s OR e."value" ILIKE %s)')
|
||
params += [f"%{query}%", f"%{query}%"]
|
||
clause = f"WHERE {' AND '.join(where)}" if where else ""
|
||
params.append(int(limit))
|
||
with get_conn() as conn:
|
||
rows = _fetchall(
|
||
conn,
|
||
f"""SELECT e.*,
|
||
(SELECT string_agg(s.school_id, ',' ORDER BY s.school_id)
|
||
FROM hub_entry_sources s WHERE s.entry_id = e.id) AS schools
|
||
FROM hub_entries e {clause}
|
||
ORDER BY e.school_count DESC, e.observations DESC, e.updated_at DESC LIMIT %s""",
|
||
params,
|
||
)
|
||
return [_row(r) for r in rows]
|
||
|
||
|
||
def get_entry(entry_id: int) -> dict | None:
|
||
with get_conn() as conn:
|
||
row = _fetchone(conn, "SELECT * FROM hub_entries WHERE id = %s", (entry_id,))
|
||
return _row(row) if row else None
|
||
|
||
|
||
def review_entry(entry_id: int, action: str, value: str | None = None, scope: str | None = None,
|
||
by: str = "human") -> dict | None:
|
||
"""采纳 / 驳回 / 退回草稿。采纳和驳回都要分配新版本号，桌面端才拉得到这次变更。"""
|
||
status = {"verify": "verified", "reject": "rejected", "reset": "draft"}.get(action)
|
||
if not status:
|
||
raise ValueError(f"未知的审核动作: {action}")
|
||
with get_conn() as conn:
|
||
sets = ['status = %s', "updated_at = NOW()"]
|
||
params: list[Any] = [status]
|
||
if value is not None and value.strip():
|
||
sets.append('"value" = %s')
|
||
params.append(value.strip())
|
||
if scope in ("origin", "all"):
|
||
sets.append("scope = %s")
|
||
params.append(scope)
|
||
if status == "draft":
|
||
sets += ["reviewed_by = NULL", "reviewed_at = NULL", "version = 0"]
|
||
else:
|
||
sets += ["reviewed_by = %s", "reviewed_at = NOW()", "version = %s"]
|
||
params += [by, _next_version(conn)]
|
||
params.append(entry_id)
|
||
_execute(conn, f"UPDATE hub_entries SET {', '.join(sets)} WHERE id = %s", params)
|
||
row = _fetchone(conn, "SELECT * FROM hub_entries WHERE id = %s", (entry_id,))
|
||
return _row(row) if row else None
|
||
|
||
|
||
def save_suggestion(entry_id: int, suggestion: str, note: str) -> None:
|
||
with get_conn() as conn:
|
||
_execute(
|
||
conn,
|
||
"UPDATE hub_entries SET suggestion = %s, suggestion_note = %s, suggested_at = NOW() WHERE id = %s",
|
||
(suggestion, note[:2000], entry_id),
|
||
)
|
||
|
||
|
||
def stats() -> dict:
|
||
with get_conn() as conn:
|
||
by_status = _fetchall(conn, "SELECT status, count(*) AS n FROM hub_entries GROUP BY status")
|
||
by_category = _fetchall(
|
||
conn, "SELECT category, count(*) AS n FROM hub_entries WHERE status = 'draft' GROUP BY category")
|
||
schools = _fetchone(conn, "SELECT count(*) AS n FROM schools")
|
||
recent = _fetchall(conn, "SELECT * FROM sync_log ORDER BY at DESC LIMIT 20")
|
||
facets = _fetchall(
|
||
conn,
|
||
"""SELECT knowledge_id, business_id, count(*) FILTER (WHERE status = 'draft') AS drafts, count(*) AS total
|
||
FROM hub_entries GROUP BY knowledge_id, business_id ORDER BY drafts DESC""",
|
||
)
|
||
return {
|
||
"by_status": {r["status"]: r["n"] for r in by_status},
|
||
"draft_by_category": {r["category"]: r["n"] for r in by_category},
|
||
"schools": schools["n"] if schools else 0,
|
||
"recent_sync": [_row(r) for r in recent],
|
||
"facets": [_row(r) for r in facets],
|
||
}
|