236 lines
12 KiB
Python
236 lines
12 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,
|
|
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 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;
|
|
"""
|
|
|
|
|
|
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)
|
|
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(),))
|
|
|
|
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"))
|
|
for table in ("observations", "learning_steps", "candidates", "validations", "artifacts"):
|
|
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(),))
|
|
|
|
@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)
|