assistant/apps/server/db.py

316 lines
18 KiB
Python

from __future__ import annotations
import hashlib
import json
import os
import re
import sqlite3
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator
def now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="milliseconds")
def identifier(prefix: str) -> str:
return f"{prefix}_{uuid.uuid4().hex}"
def encoded(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
def digest(value: Any) -> str:
return hashlib.sha256(encoded(value).encode("utf-8")).hexdigest()
SCHEMA = """
CREATE TABLE IF NOT EXISTS schema_versions(version INTEGER PRIMARY KEY, applied_at TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS operators(
id TEXT PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS auth_sessions(
token_hash TEXT PRIMARY KEY, operator_id TEXT REFERENCES operators(id), csrf TEXT NOT NULL,
expires_at INTEGER NOT NULL, created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS auth_attempts(key TEXT PRIMARY KEY, failures INTEGER NOT NULL, window_start INTEGER NOT NULL);
CREATE TABLE IF NOT EXISTS learning_tasks(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id), name TEXT NOT NULL,
platform TEXT NOT NULL, entry_url TEXT NOT NULL, goal TEXT NOT NULL, account_key TEXT NOT NULL,
parameters_json TEXT NOT NULL, status TEXT NOT NULL, identity_label TEXT, execution_profile TEXT,
created_at TEXT NOT NULL, updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS observations(
id TEXT PRIMARY KEY, task_id TEXT NOT NULL REFERENCES learning_tasks(id), data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS observations_task ON observations(task_id,created_at);
CREATE TABLE IF NOT EXISTS learning_steps(
id TEXT PRIMARY KEY, task_id TEXT NOT NULL REFERENCES learning_tasks(id), position INTEGER NOT NULL,
data_json TEXT NOT NULL, created_at TEXT NOT NULL, UNIQUE(task_id,position)
);
CREATE TABLE IF NOT EXISTS candidates(
id TEXT PRIMARY KEY, task_id TEXT NOT NULL REFERENCES learning_tasks(id), name TEXT NOT NULL,
version INTEGER NOT NULL, digest TEXT NOT NULL, manifest_json TEXT NOT NULL, created_at TEXT NOT NULL,
UNIQUE(task_id,version)
);
CREATE TABLE IF NOT EXISTS validations(
id TEXT PRIMARY KEY, candidate_id TEXT NOT NULL REFERENCES candidates(id), task_id TEXT NOT NULL REFERENCES learning_tasks(id),
kind TEXT NOT NULL, status TEXT NOT NULL, data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS validations_candidate ON validations(candidate_id,created_at);
CREATE TABLE IF NOT EXISTS quality_observations(
id TEXT PRIMARY KEY, validation_id TEXT NOT NULL REFERENCES validations(id),
candidate_id TEXT NOT NULL REFERENCES candidates(id), task_id TEXT NOT NULL REFERENCES learning_tasks(id),
data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS quality_observations_candidate ON quality_observations(candidate_id,created_at);
CREATE TABLE IF NOT EXISTS discovery_runs(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id),
candidate_id TEXT NOT NULL REFERENCES candidates(id), validation_id TEXT REFERENCES validations(id),
data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS discovery_runs_owner ON discovery_runs(owner_id,created_at);
CREATE TABLE IF NOT EXISTS reading_runs(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id),
discovery_run_id TEXT NOT NULL REFERENCES discovery_runs(id), candidate_id TEXT NOT NULL REFERENCES candidates(id),
validation_id TEXT REFERENCES validations(id), data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS reading_runs_owner ON reading_runs(owner_id,created_at);
CREATE TABLE IF NOT EXISTS analysis_runs(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id),
reading_run_id TEXT NOT NULL REFERENCES reading_runs(id), data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS analysis_runs_reading ON analysis_runs(reading_run_id,created_at);
CREATE TABLE IF NOT EXISTS business_briefs(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id),
kind TEXT NOT NULL CHECK(kind IN ('parse','confirmed')), parent_id TEXT REFERENCES business_briefs(id), cache_key TEXT,
data_json TEXT NOT NULL, created_at TEXT NOT NULL,
CHECK((kind='parse' AND parent_id IS NULL) OR (kind='confirmed' AND parent_id IS NOT NULL))
);
CREATE INDEX IF NOT EXISTS business_briefs_owner ON business_briefs(owner_id,created_at);
CREATE INDEX IF NOT EXISTS business_briefs_cache ON business_briefs(owner_id,cache_key);
CREATE TABLE IF NOT EXISTS analysis_reviews(
id TEXT PRIMARY KEY, owner_id TEXT NOT NULL REFERENCES operators(id),
analysis_run_id TEXT NOT NULL REFERENCES analysis_runs(id), source_ref TEXT NOT NULL,
previous_id TEXT REFERENCES analysis_reviews(id), data_json TEXT NOT NULL, created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS analysis_reviews_analysis ON analysis_reviews(analysis_run_id,created_at);
CREATE UNIQUE INDEX IF NOT EXISTS analysis_reviews_first ON analysis_reviews(owner_id,analysis_run_id,source_ref) WHERE previous_id IS NULL;
CREATE UNIQUE INDEX IF NOT EXISTS analysis_reviews_successor ON analysis_reviews(previous_id) WHERE previous_id IS NOT NULL;
CREATE TABLE IF NOT EXISTS artifacts(
id TEXT PRIMARY KEY, task_id TEXT NOT NULL REFERENCES learning_tasks(id), relative_path TEXT NOT NULL,
sha256 TEXT NOT NULL, bytes INTEGER NOT NULL, created_at TEXT NOT NULL
);
CREATE TRIGGER IF NOT EXISTS immutable_candidate_manifest BEFORE UPDATE ON candidates
BEGIN SELECT RAISE(ABORT,'candidate versions are immutable'); END;
CREATE TRIGGER IF NOT EXISTS immutable_validation BEFORE UPDATE ON validations
BEGIN SELECT RAISE(ABORT,'validation records are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_candidate_delete BEFORE DELETE ON candidates
BEGIN SELECT RAISE(ABORT,'candidate versions are immutable'); END;
CREATE TRIGGER IF NOT EXISTS immutable_validation_delete BEFORE DELETE ON validations
BEGIN SELECT RAISE(ABORT,'validation records are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_observation_update BEFORE UPDATE ON observations
BEGIN SELECT RAISE(ABORT,'source observations are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_observation_delete BEFORE DELETE ON observations
BEGIN SELECT RAISE(ABORT,'source observations are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_step_update BEFORE UPDATE ON learning_steps
BEGIN SELECT RAISE(ABORT,'learned steps are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_step_delete BEFORE DELETE ON learning_steps
BEGIN SELECT RAISE(ABORT,'learned steps are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_artifact_update BEFORE UPDATE ON artifacts
BEGIN SELECT RAISE(ABORT,'evidence descriptors are immutable'); END;
CREATE TRIGGER IF NOT EXISTS immutable_artifact_delete BEFORE DELETE ON artifacts
BEGIN SELECT RAISE(ABORT,'evidence descriptors are immutable'); END;
CREATE TRIGGER IF NOT EXISTS immutable_quality_observation_update BEFORE UPDATE ON quality_observations
BEGIN SELECT RAISE(ABORT,'quality observations are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_quality_observation_delete BEFORE DELETE ON quality_observations
BEGIN SELECT RAISE(ABORT,'quality observations are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_discovery_run_update BEFORE UPDATE ON discovery_runs
BEGIN SELECT RAISE(ABORT,'discovery runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_discovery_run_delete BEFORE DELETE ON discovery_runs
BEGIN SELECT RAISE(ABORT,'discovery runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_reading_run_update BEFORE UPDATE ON reading_runs
BEGIN SELECT RAISE(ABORT,'reading runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_reading_run_delete BEFORE DELETE ON reading_runs
BEGIN SELECT RAISE(ABORT,'reading runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_analysis_run_update BEFORE UPDATE ON analysis_runs
BEGIN SELECT RAISE(ABORT,'analysis runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_analysis_run_delete BEFORE DELETE ON analysis_runs
BEGIN SELECT RAISE(ABORT,'analysis runs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_business_brief_update BEFORE UPDATE ON business_briefs
BEGIN SELECT RAISE(ABORT,'business briefs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_business_brief_delete BEFORE DELETE ON business_briefs
BEGIN SELECT RAISE(ABORT,'business briefs are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_analysis_review_update BEFORE UPDATE ON analysis_reviews
BEGIN SELECT RAISE(ABORT,'analysis reviews are append only'); END;
CREATE TRIGGER IF NOT EXISTS immutable_analysis_review_delete BEFORE DELETE ON analysis_reviews
BEGIN SELECT RAISE(ABORT,'analysis reviews are append only'); END;
"""
class SQLiteConnection(sqlite3.Connection):
def execute(self, sql: str, parameters=(), /):
# SQLite treats native $1 names as named parameters; bind by name rather
# than the sequence compatibility path removed in Python 3.14. Query text
# remains unchanged, and legacy qmark calls in SQLite-only tools still work.
if isinstance(parameters, (tuple, list)) and re.search(r"\$[1-9][0-9]*\b", sql):
parameters = {str(index): value for index, value in enumerate(parameters, 1)}
return super().execute(sql, parameters)
class Store:
backend = "sqlite"
def __init__(self, home: Path):
self.home = home.resolve()
self.home.mkdir(parents=True, exist_ok=True)
self.path = self.home / "learning.sqlite"
with self.connect() as conn:
conn.executescript(SCHEMA)
if "execution_profile" not in {row["name"] for row in conn.execute("PRAGMA table_info(learning_tasks)")}:
conn.execute("ALTER TABLE learning_tasks ADD COLUMN execution_profile TEXT")
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(1,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(2,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(4,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(5,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(6,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(7,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(8,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(9,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(10,$1) ON CONFLICT(version) DO NOTHING", (now(),))
def order_column(self, alias: str = "") -> str:
if alias not in ("", "c", "v"):
raise ValueError("unsupported internal table alias")
return (alias + "." if alias else "") + ("ordinal" if self.backend == "postgresql" else "rowid")
@contextmanager
def connect(self, immediate: bool = False) -> Iterator[sqlite3.Connection]:
conn = sqlite3.connect(self.path, timeout=15, factory=SQLiteConnection)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA foreign_keys=ON")
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=15000")
try:
if immediate:
conn.execute("BEGIN IMMEDIATE")
yield conn
conn.commit()
except BaseException:
conn.rollback()
raise
finally:
conn.close()
def one(self, query: str, args: tuple = ()) -> dict | None:
with self.connect() as conn:
row = conn.execute(query, args).fetchone()
return dict(row) if row else None
def rows(self, query: str, args: tuple = ()) -> list[dict]:
with self.connect() as conn:
return [dict(row) for row in conn.execute(query, args).fetchall()]
def execute(self, query: str, args: tuple = ()) -> None:
with self.connect() as conn:
conn.execute(query, args)
class StoreError(Exception):
"""Public database diagnostics never include a DSN or a driver exception."""
def __init__(self, code: str, message: str, status: int = 503):
self.code, self.message, self.status = code, message, status
super().__init__(message)
class StoreConflict(StoreError):
def __init__(self):
super().__init__("database_conflict", "记录冲突或不可变事实禁止修改,请重新读取后核对", 409)
class PostgresStore(Store):
backend = "postgresql"
def __init__(self, home: Path, database_url: str, *, schema: str = "assistant_learning"):
if len(schema) > 63 or not re.fullmatch(r"assistant_learning(?:_[a-z0-9_]+)?", schema):
raise StoreError("database_schema_invalid", "学习数据库命名空间无效")
if not database_url.startswith(("postgresql://", "postgres://")):
raise StoreError("database_configuration_invalid", "学习数据库地址需使用 PostgreSQL URL")
try:
import psycopg
from psycopg import sql
from psycopg.rows import dict_row
except ImportError:
raise StoreError("database_driver_unavailable", "缺少 PostgreSQL 驱动,请安装学习后台依赖") from None
self.home = home.resolve()
self.home.mkdir(parents=True, exist_ok=True)
self.path = None
self.schema = schema
self._database_url = database_url
self._driver, self._sql, self._row_factory = psycopg, sql, dict_row
self._write_lock = int.from_bytes(hashlib.sha256((schema + ":write").encode()).digest()[:8], "big", signed=True)
with self.connect(immediate=True) as conn:
conn.execute(sql.SQL("CREATE SCHEMA IF NOT EXISTS {}" ).format(sql.Identifier(schema)))
conn.execute(sql.SQL("REVOKE ALL ON SCHEMA {} FROM PUBLIC").format(sql.Identifier(schema)))
conn.execute((Path(__file__).with_name("postgres.sql")).read_text(encoding="utf-8"))
conn.execute("ALTER TABLE learning_tasks ADD COLUMN IF NOT EXISTS execution_profile TEXT")
for table in ("observations", "learning_steps", "candidates", "validations", "artifacts", "quality_observations", "discovery_runs", "reading_runs", "analysis_runs", "business_briefs", "analysis_reviews"):
conn.execute(sql.SQL("CREATE OR REPLACE TRIGGER immutable_learning_fact BEFORE UPDATE OR DELETE ON {} FOR EACH ROW EXECUTE FUNCTION reject_learning_fact_mutation()").format(sql.Identifier(table)))
conn.execute(sql.SQL("CREATE OR REPLACE TRIGGER immutable_learning_truncate BEFORE TRUNCATE ON {} FOR EACH STATEMENT EXECUTE FUNCTION reject_learning_fact_mutation()").format(sql.Identifier(table)))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(3,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(4,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(5,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(6,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(7,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(8,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(9,$1) ON CONFLICT(version) DO NOTHING", (now(),))
conn.execute("INSERT INTO schema_versions(version,applied_at) VALUES(10,$1) ON CONFLICT(version) DO NOTHING", (now(),))
@contextmanager
def connect(self, immediate: bool = False) -> Iterator[Any]:
connection = None
try:
connection = self._driver.connect(self._database_url, cursor_factory=self._driver.RawCursor,
row_factory=self._row_factory, connect_timeout=5)
# This schema is fixed by the host, never supplied by an HTTP caller.
# Both drivers bind native $1/$2 parameters; no query text rewriting.
connection.execute(self._sql.SQL("SET search_path TO {}, pg_catalog").format(self._sql.Identifier(self.schema)))
connection.execute("SET lock_timeout TO '15s'")
connection.execute("SET statement_timeout TO '30s'")
if immediate:
# Same write serialization guarantee as SQLite BEGIN IMMEDIATE.
connection.execute("SELECT pg_advisory_xact_lock($1)", (self._write_lock,))
yield connection
connection.commit()
except self._driver.IntegrityError:
if connection is not None:
self._rollback(connection)
raise StoreConflict() from None
except self._driver.Error:
if connection is not None:
self._rollback(connection)
raise StoreError("database_unavailable", "学习数据库操作未完成,请检查本机 PostgreSQL 服务并重新读取记录") from None
except BaseException:
if connection is not None:
self._rollback(connection)
raise
finally:
if connection is not None:
connection.close()
def _rollback(self, connection) -> None:
try:
connection.rollback()
except self._driver.Error:
pass
def configured_store(home: Path, database_url: str | None = None, *, schema: str | None = None) -> Store:
selected = os.environ.get("ASSISTANT_DATABASE_URL", "") if database_url is None else database_url
namespace = os.environ.get("ASSISTANT_DATABASE_SCHEMA", "assistant_learning") if schema is None else schema
return PostgresStore(home, selected, schema=namespace) if selected else Store(home)