diff --git a/backend/app/auth/access.py b/backend/app/auth/access.py index 3c04db08..07eda877 100644 --- a/backend/app/auth/access.py +++ b/backend/app/auth/access.py @@ -62,6 +62,7 @@ def hidden_runs(principal: Any) -> Any: from app.db.models import Run from app.db.models_apps import App from app.db.models_blueprints import Blueprint + from app.db.portable import json_text if principal.mode == "off": return false() @@ -72,7 +73,7 @@ def hidden_runs(principal: Any) -> Any: kind = func.coalesce(Run.origin_kind, "") return or_( and_(kind == "personal_app", or_(Run.app_id.is_(None), Run.app_id.not_in(mine))), - and_(kind == "trial", func.coalesce(func.json_extract(Run.origin, "$.blueprint_id"), "").in_(others)), + and_(kind == "trial", func.coalesce(json_text(Run.origin, "blueprint_id"), "").in_(others)), ) diff --git a/backend/app/auth/keys.py b/backend/app/auth/keys.py index 1d4c5dc3..5f228c6c 100644 --- a/backend/app/auth/keys.py +++ b/backend/app/auth/keys.py @@ -39,7 +39,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Awaitable, Callable -from sqlalchemy import delete, func, select +from sqlalchemy import delete, select from starlette.requests import HTTPConnection from app.api.coded import CodedHTTPException @@ -535,12 +535,13 @@ async def started_before(principal: Any, idem_key: str) -> str | None: """拥有者按这个幂等键用密钥发起过的实例(来源里记着 idempotency_key):处理到一半服务停了、同一个键重新认领时交回它。""" from app.db.base import SessionLocal from app.db.models import Run + from app.db.portable import json_eq async with SessionLocal() as session: return await session.scalar(select(Run.id).where( Run.origin_kind == "app_entry", Run.started_by == principal.username, - func.json_extract(Run.origin, "$.idempotency_key") == idem_key, - func.json_extract(Run.origin, "$.via") == "key").order_by(Run.created_at.desc()).limit(1)) + json_eq(Run.origin, "idempotency_key", idem_key), + json_eq(Run.origin, "via", "key")).order_by(Run.created_at.desc()).limit(1)) # ---------------------------------------------------------------- 请求记录(A6) diff --git a/backend/app/blueprints/personal.py b/backend/app/blueprints/personal.py index 4c73a020..770d5864 100644 --- a/backend/app/blueprints/personal.py +++ b/backend/app/blueprints/personal.py @@ -416,9 +416,10 @@ async def refresh(session: AsyncSession, app: Any, *, principal: Any) -> tuple[A def _runs_of(app_id: str, graph_hash: str) -> Any: from app.db.models import Run + from app.db.portable import json_eq return select(Run).where(Run.app_id == app_id, Run.origin_kind == ORIGIN, - func.json_extract(Run.origin, "$.graph_hash") == graph_hash) + json_eq(Run.origin, "graph_hash", graph_hash)) async def latest_refresh(session: AsyncSession, app_id: str, graph_hash: str, *, active_only: bool = False) -> Any: diff --git a/backend/app/capabilities/registry.py b/backend/app/capabilities/registry.py index ea63d849..77a49d13 100644 --- a/backend/app/capabilities/registry.py +++ b/backend/app/capabilities/registry.py @@ -127,10 +127,9 @@ async def _discover(session: AsyncSession) -> dict[str, dict[str, Any]]: async def sync_all(session: AsyncSession) -> int: """把内置工具、数据源、自定义工具、MCP 同步进 capabilities。返回变动的行数。启动时和各对象增删改时调。""" - from sqlalchemy.dialects.sqlite import insert - from app.db.base import utcnow from app.db.models import Capability + from app.db.portable import insert_ignore found = await _discover(session) existing = {row.name: row for row in (await session.execute(select(Capability))).scalars()} @@ -169,8 +168,10 @@ async def sync_all(session: AsyncSession) -> int: row.synced_at = now changed += 1 if fresh: - # 两个请求同时同步(列表页和发起运行撞在一起):后到的那条插入什么都不做,不报主键冲突。不加锁(单进程 + SQLite) - done = await session.execute(insert(Capability).values(fresh).on_conflict_do_nothing(index_elements=["name"])) + # 两个请求同时同步(列表页和发起运行撞在一起):后到的那条插入什么都不做,不报主键冲突。不加锁:两种库都是同一条 + # INSERT … ON CONFLICT DO NOTHING,多个进程同时同步也一样 + done = await session.execute( + insert_ignore(session, Capability).values(fresh).on_conflict_do_nothing(index_elements=["name"])) changed += done.rowcount or 0 if changed or fresh: await session.commit() diff --git a/backend/app/chat/questions.py b/backend/app/chat/questions.py index 99d61240..54d56f50 100644 --- a/backend/app/chat/questions.py +++ b/backend/app/chat/questions.py @@ -98,6 +98,7 @@ async def recent_questions(session: AsyncSession, principal: Any, *, limit: int) """最近问过的问题,从新到旧,最多 limit 条(取法见模块开头)。principal 用来去掉这个人读不了的只给拥有者的运行。""" from app.auth.access import hidden_runs from app.db.models import Conversation, ConversationTurn, Run + from app.db.portable import json_is_string from app.engine.origin import counted_runs hidden = hidden_runs(principal) @@ -112,7 +113,7 @@ async def recent_questions(session: AsyncSession, principal: Any, *, limit: int) runs = (await session.execute( select(Run.id, Run.input, Run.workflow_name, Run.created_at, Run.origin_kind) .where(Run.run_class != "formal", func.coalesce(Run.origin_kind, "") != "chat", counted_runs(), not_(hidden), - func.json_type(Run.input, "$.question") == "text") + json_is_string(Run.input, "question")) .order_by(Run.created_at.desc()).limit(limit) )).all() out = [Question(text=question.strip(), at=at, run_id=run_id, turn_id=turn_id, workflow=workflow, origin="chat") diff --git a/backend/app/core/config.py b/backend/app/core/config.py index 955e7242..2ddb6629 100644 --- a/backend/app/core/config.py +++ b/backend/app/core/config.py @@ -47,6 +47,17 @@ class Settings(BaseSettings): # --- 存储 --- data_dir: Path = PROJECT_DIR / "data" + # 应用自己的库(工作流、运行、事件、审批……)。不设:数据目录里的 SQLite(agentlab.db,断点在 checkpoints.db), + # 单机单进程。设成 postgresql+asyncpg://用户:密码@主机:端口/库名(postgresql:// 也认)时整个应用库用 PostgreSQL, + # LangGraph 的断点也放进同一个库;checkpoint_url 另给就放那里。数据目录照旧要有:工件、上传表格、沙箱等仍是文件 + database_url: str = "" + checkpoint_url: str = "" + # PostgreSQL 的连接池(SQLite 不用这几项)。pool_size=0 表示不用连接池、每次新建连接(前面有 PgBouncer 时用, + # 测试里也这么用:每个测试一个事件循环,asyncpg 的连接不能跨循环) + db_pool_size: int = 10 + db_max_overflow: int = 20 + db_pool_timeout: int = 30 + db_pool_recycle: int = 1800 # --- 前端构建产物(生产部署用)--- # 设了就由服务端一并托管页面(见 app/web.py),页面和接口同源、只起一个进程。 @@ -195,6 +206,23 @@ class Settings(BaseSettings): return True return value + @field_validator("database_url", "checkpoint_url", mode="before") + @classmethod + def _valid_postgres_url(cls, value: object, info: Any) -> object: + # 写错了启动时就报出来,别等到第一次查库。只认 PostgreSQL:SQLite 就是不设 + if value is None or (isinstance(value, str) and not value.strip()): + return "" + if not isinstance(value, str): + return value + from app.db.urls import normalize_database_url, normalize_checkpoint_url + + name = f"AGENTLAB_{(info.field_name or '').upper()}" + normalize = normalize_database_url if info.field_name == "database_url" else normalize_checkpoint_url + try: + return normalize(value.strip()) + except ValueError as e: + raise ValueError(f"{name}:{e}") from None + @field_validator("web_dist", mode="before") @classmethod def _blank_is_unset(cls, value: object) -> object: @@ -223,13 +251,30 @@ class Settings(BaseSettings): @property def db_url(self) -> str: - return f"sqlite+aiosqlite:///{self.db_path}" + """应用库的 SQLAlchemy 地址:设了 database_url 用它(已规范成 postgresql+asyncpg://),否则数据目录里的 SQLite。""" + return self.database_url or f"sqlite+aiosqlite:///{self.db_path}" + + @property + def db_dialect(self) -> str: + """"sqlite" 或 "postgresql"。""" + return "postgresql" if self.database_url else "sqlite" @property def checkpoint_path(self) -> Path: - """LangGraph checkpointer 用独立的库文件,避免和业务表抢写锁。""" + """LangGraph checkpointer 用独立的库文件,避免和业务表抢写锁。只在 SQLite 时用。""" return self.data_dir / "checkpoints.db" + @property + def checkpoint_conninfo(self) -> str | None: + """PostgreSQL 断点库的 libpq 连接串(psycopg 用):checkpoint_url 优先,否则由 database_url 换算。SQLite 时 None。""" + if self.checkpoint_url: + return self.checkpoint_url + if self.database_url: + from app.db.urls import libpq_conninfo + + return libpq_conninfo(self.database_url) + return None + @property def workspace_dir(self) -> Path: """沙箱与文件工具的可写根目录,一切文件访问都被限制在这里面。""" diff --git a/backend/app/data/capability_review.py b/backend/app/data/capability_review.py index 3db355c0..9751f4e3 100644 --- a/backend/app/data/capability_review.py +++ b/backend/app/data/capability_review.py @@ -144,10 +144,18 @@ def normalize_fields(existing: Mapping[str, Any] | None, submitted: Any, *, at: async def run_counts(session: AsyncSession) -> dict[str, int]: """每项能力被几次运行用到:{能力名: 运行次数},没用过的不出现。按运行发起时记下的能力治理快照数 - (Run.options["capabilities"],一次运行里一项能力只记一次)。只用 SQLite 的 json_each,一次查询。""" - rows = await session.execute(text( - "SELECT j.key, COUNT(*) FROM runs, json_each(runs.options, '$.capabilities') AS j " - "WHERE json_valid(runs.options) GROUP BY j.key")) + (Run.options["capabilities"],一次运行里一项能力只记一次)。一次查询:两种库都用 json_each(PostgreSQL 的碰到不是 + 对象的值会报错,先按类型挑出来)。""" + from app.db.portable import dialect_of + + if dialect_of(session) == "postgresql": + sql = ("SELECT j.key, COUNT(*) FROM runs CROSS JOIN LATERAL json_each(CASE WHEN " + "json_typeof(runs.options -> 'capabilities') = 'object' THEN runs.options -> 'capabilities' " + "ELSE '{}'::json END) AS j GROUP BY j.key") + else: + sql = ("SELECT j.key, COUNT(*) FROM runs, json_each(runs.options, '$.capabilities') AS j " + "WHERE json_valid(runs.options) GROUP BY j.key") + rows = await session.execute(text(sql)) return {str(name): int(count) for name, count in rows.all()} diff --git a/backend/app/data/catalog_profile_jobs.py b/backend/app/data/catalog_profile_jobs.py index 37a4b097..810c3758 100644 --- a/backend/app/data/catalog_profile_jobs.py +++ b/backend/app/data/catalog_profile_jobs.py @@ -48,7 +48,7 @@ import logging import time from typing import Any -from sqlalchemy import case, delete, func, or_, select, update +from sqlalchemy import case, delete, func, not_, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession @@ -60,6 +60,7 @@ from app.data.catalog_profile import ( ) from app.data.engine import SnapshotTampered from app.db.base import SessionLocal, new_id, utcnow +from app.db.portable import json_flag from app.db.models import CatalogDraftJob, CatalogProfileJob, CatalogProfileTask, DataSource logger = logging.getLogger(__name__) @@ -345,7 +346,7 @@ _SHOWN = (case((CatalogProfileTask.started_at.is_(None), 1), else_=0), *_NEWEST) #: 只做了取值观察的那一项(observe 范围、maintenance 范围里只到期观察的表)在逐表一行的 summary 里带这个标记:它不是一次 #: 剖析,不算「上次剖析」——不然表清单的「上次剖析」、给助手的剖析事实(行数、日期范围)会被一次只记了行数的观察顶掉 OBSERVE_ONLY = "observe_only" -_PROFILED = func.coalesce(func.json_extract(CatalogProfileTask.summary, f"$.{OBSERVE_ONLY}"), 0) == 0 +_PROFILED = not_(json_flag(CatalogProfileTask.summary, OBSERVE_ONLY)) def _observe_only(checks: list[Any] | None) -> bool: diff --git a/backend/app/data/recipe_imports.py b/backend/app/data/recipe_imports.py index ee8d338c..e3d4598e 100644 --- a/backend/app/data/recipe_imports.py +++ b/backend/app/data/recipe_imports.py @@ -90,6 +90,7 @@ from app.data.recipe_types import ( from app.data.xlsx_scan import WorkbookScan from app.db.base import new_id, utcnow from app.db.models import DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe +from app.db.portable import unique_violation logger = logging.getLogger(__name__) @@ -2747,7 +2748,7 @@ async def commit_staging(session: AsyncSession, staging: ImportStaging, *, trial raise _publish_refused(e) from e except IntegrityError as e: await _back_to_drafting(session, staging_id) - if "data_sources.name" in str(e.orig if e.orig is not None else e): + if unique_violation(e, "data_sources", "name"): raise ImportRefused(409, "name_taken", f"名称「{name}」已被占用,请放弃这次导入后换一个名称") from e raise ImportRefused(500, "publish_failed", "导入失败:写入版本记录时出错,当前版本未受影响") from e except BaseException: diff --git a/backend/app/data/source_changes.py b/backend/app/data/source_changes.py index 5b625baf..51856bf0 100644 --- a/backend/app/data/source_changes.py +++ b/backend/app/data/source_changes.py @@ -37,6 +37,7 @@ from app.data import schema_diff as sd from app.data.names import name_key from app.db.base import utcnow from app.db.models_sources import SourceChange +from app.db.portable import json_eq logger = logging.getLogger(__name__) @@ -337,7 +338,7 @@ async def notify_after_job(source_id: str, job_id: str) -> int: n = int((await session.execute(select(func.count()).select_from(SourceChange).where( SourceChange.source_id == source_id, SourceChange.severity == "action", SourceChange.status == "open", - func.json_extract(SourceChange.detected_by, "$.job_id") == job_id))).scalar() or 0) + json_eq(SourceChange.detected_by, "job_id", job_id)))).scalar() or 0) if n: from app.notify import service diff --git a/backend/app/data/table_versions.py b/backend/app/data/table_versions.py index ccf212cf..39230169 100644 --- a/backend/app/data/table_versions.py +++ b/backend/app/data/table_versions.py @@ -540,10 +540,13 @@ _DB_SIDE_FILES = ("", "-wal", "-shm", "-journal") def _app_db_files() -> list[Path]: - """应用自己的库:元数据库、运行检查点库,以及实际连着的元数据库(测试里和 data_dir 下的不是同一个)。""" + """应用自己的库:元数据库、运行检查点库,以及实际连着的元数据库(测试里和 data_dir 下的不是同一个)。 + + 应用库是 PostgreSQL 时,数据目录里的 agentlab.db、checkpoints.db 照样算:导入 PostgreSQL 之前的那份还留在那里, + 里面是同一批元数据。连着的库这时是库名不是文件,不算进来。""" dbs = [settings.db_path, settings.checkpoint_path] try: - live = db_base.engine.url.database + live = db_base.engine.url.database if db_base.engine.dialect.name == "sqlite" else None except Exception: # noqa: BLE001 - 取不到就只按配置里的路径判断 live = None if live and live != ":memory:": diff --git a/backend/app/db/base.py b/backend/app/db/base.py index 24a40182..7e741218 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -1,16 +1,23 @@ from __future__ import annotations +import json +import logging +import math +import re import uuid from datetime import datetime, timezone from typing import Any, AsyncIterator -from sqlalchemy import JSON, DateTime, MetaData, event -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy import JSON, BigInteger, Column, DateTime, Integer, MetaData, event +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column +from sqlalchemy.pool import NullPool from sqlalchemy.types import TypeDecorator from app.core.config import settings +logger = logging.getLogger(__name__) + NAMING_CONVENTION = { "ix": "ix_%(column_0_label)s", "uq": "uq_%(table_name)s_%(column_0_name)s", @@ -57,11 +64,31 @@ class UTCDateTime(TypeDecorator): return value.astimezone(timezone.utc) +#: 应用库的 JSON 列:两种库上都是存原文的 JSON(PostgreSQL 的 json,不是 JSONB)。JSONB 会把对象里的键重新排序: +#: 应用里不少地方按存进去的键的先后出东西(快照的列、表的先后、复核结论各段的先后),换成 JSONB 读回来就变了样, +#: 测试当场抓到好几处。json 的代价是没有相等运算符:== 比较、DISTINCT、GROUP BY、ORDER BY 不能直接作用在整列上 +#: (应用里没有这种写法,test_db_portable 守着);按里面某个键筛选走 db/portable 的 json_eq 等,只把取出来的那一小段 +#: 转成 JSONB 比较 +PortableJSON = JSON() + + class Base(DeclarativeBase): metadata = MetaData(naming_convention=NAMING_CONVENTION) - # 图定义、节点配置、事件负载都是任意形状的 JSON,统一映射到 SQLite 的 JSON 列 - type_annotation_map = {dict[str, Any]: JSON, list[Any]: JSON, datetime: UTCDateTime} + # 图定义、节点配置、事件负载都是任意形状的 JSON,统一映射到 JSON 列(见 PortableJSON) + type_annotation_map = {dict[str, Any]: PortableJSON, list[Any]: PortableJSON, datetime: UTCDateTime} + + +@event.listens_for(Column, "after_parent_attach") +def _portable_integer(column: Column, table: Any) -> None: + """应用库的 Integer 列在 PostgreSQL 上换成 BIGINT(SQLite 照旧是 INTEGER,主键照旧是 rowid)。 + + SQLite 的 INTEGER 是 64 位,PostgreSQL 的只有 32 位:行数估计、字节数、累计的计数都可能过 21 亿。在列挂到表上 + 的那一刻换(定义模型时,任何查询之前):建表是 BIGINT(自增主键是 BIGSERIAL),和这一列比较、写进这一列的绑定 + 参数也按 BIGINT 转换。只管挂在应用 metadata 上、类型正好是 Integer 的列;业务数据源那边的表不受影响。""" + if getattr(table, "metadata", None) is Base.metadata and type(column.type) is Integer \ + and not column.type._variant_mapping: + column.type = Integer().with_variant(BigInteger(), "postgresql") class TimestampMixin: @@ -101,16 +128,10 @@ def _stop_aiosqlite_once() -> None: aiosqlite.Connection.stop = stop # type: ignore[method-assign] +# 应用库是 PostgreSQL 时也照样打:上传表格、手工登记的 SQLite 业务库(data/engine)还是走 aiosqlite 的连接池 _stop_aiosqlite_once() -# pool_pre_ping:取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时 -# (_close_connection 记完日志会把 CancelledError 原样抛出,跳过了把连接记录 -# 清空的那一行),已经关掉的连接会留在池里,下一个取到它的人拿到的是 -# "no active connection"。探活失败会被当成断开,换一条新的——SQLite 本地一次 -# SELECT 1 的代价,换池子不被污染。 -engine = create_async_engine(settings.db_url, echo=False, future=True, pool_pre_ping=True) -@event.listens_for(engine.sync_engine, "connect") def _set_sqlite_pragma(dbapi_conn, _record): # pragma: no cover - 驱动回调 """WAL + 外键约束。WAL 让读写不互相阻塞,事件写入和前端轮询能并行。""" cur = dbapi_conn.cursor() @@ -120,6 +141,92 @@ def _set_sqlite_pragma(dbapi_conn, _record): # pragma: no cover - 驱动回调 cur.close() +def _finite(value: Any) -> Any: + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, dict): + return {k: _finite(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_finite(v) for v in value] + return value + + +def _no_nul(value: Any) -> Any: + if isinstance(value, str): + return value.replace("\x00", "") + if isinstance(value, dict): + return {_no_nul(k): _no_nul(v) for k, v in value.items()} + if isinstance(value, list): + return [_no_nul(v) for v in value] + return value + + +def pg_json_dumps(value: Any) -> str: + """PostgreSQL 上写 JSON 列。它不收 NaN / Infinity(SQLite 照存);字符串里的 \\u0000 json 列收,但一转成 JSONB + (按键筛选时,db/portable)、一取成文本就报错,那一整条查询跟着失败。 + + 非有限的浮点写成 null(同浏览器的 JSON.stringify),字符串里的 NUL 字符去掉;平常的值一次 dumps 就过, + 只有碰上了才多走一遍。""" + try: + text = json.dumps(value, allow_nan=False) + except ValueError: + text = json.dumps(_finite(value), allow_nan=False) + if "\\u0000" in text: + text = json.dumps(_no_nul(json.loads(text)), allow_nan=False) + return text + + +#: 应用连 PostgreSQL 时报的 application_name。copy-to-postgres 按它认目标库上还有没有服务连着(它自己报的是 agentlab-copy) +APPLICATION_NAME = "agentlab" + + +def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME) -> AsyncEngine: + """应用库的引擎。SQLite:WAL、外键、忙等,单个文件;PostgreSQL:连接池、方言调整(db/postgres.py)。 + + copy-to-postgres 建目标库的引擎也走这里,两边的列类型、排序规则是同一套。 + """ + url = url or settings.db_url + if url.startswith("sqlite"): + # pool_pre_ping:取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时 + # (_close_connection 记完日志会把 CancelledError 原样抛出,跳过了把连接记录 + # 清空的那一行),已经关掉的连接会留在池里,下一个取到它的人拿到的是 + # "no active connection"。探活失败会被当成断开,换一条新的——SQLite 本地一次 + # SELECT 1 的代价,换池子不被污染。 + made = create_async_engine(url, echo=False, future=True, pool_pre_ping=True) + event.listen(made.sync_engine, "connect", _set_sqlite_pragma) + return made + from app.db import postgres + + pool: dict[str, Any] = ( + {"poolclass": NullPool} if settings.db_pool_size <= 0 else { + "pool_size": settings.db_pool_size, "max_overflow": max(settings.db_max_overflow, 0), + "pool_timeout": settings.db_pool_timeout, "pool_recycle": settings.db_pool_recycle}) + made = create_async_engine( + url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps, + # 会话时区定成 UTC:时间列是不带时区的 UTC(UTCDateTime),库里 now() 之类的也按 UTC 算 + connect_args={"server_settings": {"application_name": application_name, "timezone": "UTC"}}, + **pool) + postgres.install(made.sync_engine.dialect) + return made + + +def describe_database() -> str: + """启动日志里的一句:用的是哪种库、在哪(不写密码)。""" + from app.db.urls import describe + + if settings.checkpoint_url: + points = f"断点在 PostgreSQL {describe(settings.checkpoint_url)}" + elif settings.db_dialect == "sqlite": + points = f"断点 {settings.checkpoint_path}" + else: + points = "断点在同一个库" + if settings.db_dialect == "sqlite": + return f"SQLite({settings.db_path},{points})" + return f"PostgreSQL({describe(settings.db_url)},{points})" + + +engine = make_engine() + SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) @@ -129,8 +236,10 @@ async def get_session() -> AsyncIterator[AsyncSession]: yield session -# 轻量迁移:create_all 只建新表,不给已有表加列。SQLite 的 ADD COLUMN -# 便宜且带默认值回填,够用;等真需要改列类型再考虑 alembic。 +# 轻量迁移:create_all 只建新表,不给已有表加列。ADD COLUMN 便宜且带默认值回填,够用;等真需要改列类型再考虑 alembic。 +# 第三项照 SQLite 写(类型和缺省值的字面量)。PostgreSQL 上类型取模型里这一列的(按应用库的方言调整编译:String(n) +# 是 VARCHAR、Integer 是 BIGINT、JSON 是 json、时间是 TIMESTAMP),缺省值取这里写的,布尔的 0 / 1 换成 false / true +# (_column_ddl)。所以这里加的每一列模型里都得有:test_db_portable 守着 _COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [ ("runs", "run_class", "TEXT DEFAULT 'exploratory'"), ("runs", "version", "INTEGER"), @@ -274,44 +383,92 @@ _INDEX_MIGRATIONS: list[tuple[str, str, str | tuple[str, ...]]] = [ #: 迁移补建的唯一索引:(索引名, 表, 列, 部分索引的条件, 建之前的去重语句)。只在索引还不存在时先去重再建——存量库里可能 #: 已经有并发留下的重复行,直接建会失败。去重的规矩写在各自的语句里:多出的钩子行删掉(只留最早的那一行,它的密钥照旧 -#: 能用),多出的助手对话归档(不删,历史还在) +#: 能用),多出的助手对话归档(不删,历史还在)。条件和语句里的 {false} / {true} 按方言换成布尔的字面量(SQLite 是 0 / 1) _UNIQUE_INDEX_MIGRATIONS: list[tuple[str, str, tuple[str, ...], str | None, str]] = [ ("ux_app_hooks_app_name", "app_hooks", ("app_id", "name"), None, "DELETE FROM app_hooks WHERE id NOT IN (SELECT id FROM (SELECT id, ROW_NUMBER() OVER " - "(PARTITION BY app_id, name ORDER BY created_at, id) AS rn FROM app_hooks) WHERE rn = 1)"), - ("ux_conversations_app_owner", "conversations", ("app_id", "created_by"), "kind = 'app' AND archived = 0", - "UPDATE conversations SET archived = 1 WHERE kind = 'app' AND archived = 0 AND created_by IS NOT NULL " + "(PARTITION BY app_id, name ORDER BY created_at, id) AS rn FROM app_hooks) AS ranked WHERE rn = 1)"), + ("ux_conversations_app_owner", "conversations", ("app_id", "created_by"), "kind = 'app' AND archived = {false}", + "UPDATE conversations SET archived = {true} WHERE kind = 'app' AND archived = {false} AND created_by IS NOT NULL " "AND id NOT IN (SELECT id FROM (SELECT id, ROW_NUMBER() OVER (PARTITION BY app_id, created_by " - "ORDER BY last_active_at DESC, id) AS rn FROM conversations WHERE kind = 'app' AND archived = 0 " - "AND created_by IS NOT NULL) WHERE rn = 1)"), + "ORDER BY last_active_at DESC, id) AS rn FROM conversations WHERE kind = 'app' AND archived = {false} " + "AND created_by IS NOT NULL) AS ranked WHERE rn = 1)"), ] +#: 布尔字面量:SQLite 的布尔列存 0 / 1 +_BOOL_LITERALS = {"sqlite": {"false": "0", "true": "1"}, "postgresql": {"false": "false", "true": "true"}} + +_DEFAULT = re.compile(r"\s+DEFAULT\s+(.+)$", re.I | re.S) + + +def _column_ddl(dialect: Any, table: str, column: str, ddl: str) -> str: + """ADD COLUMN 里列名后面那一段。SQLite 原样用迁移表里写的;PostgreSQL 类型取模型里这一列的,缺省值取迁移表里的。""" + if dialect.name == "sqlite": + return ddl + from sqlalchemy import Boolean + + from app.db.postgres import AppTypeCompiler + + col = Base.metadata.tables[table].c[column] + out = AppTypeCompiler(dialect).process(col.type, type_expression=col) + if match := _DEFAULT.search(ddl): + default = match.group(1).strip() + if isinstance(col.type, Boolean): + default = {"0": "false", "1": "true"}.get(default, default) + out += f" DEFAULT {default}" + return out + def _migrate(conn) -> None: # pragma: no cover - 同步回调 - from sqlalchemy import text + """给老库补列、补索引。已有的列和索引按 inspector 认(两种库同一套),再跑一遍什么都不做。""" + from sqlalchemy import inspect, text + dialect = conn.dialect + literals = _BOOL_LITERALS.get(dialect.name, _BOOL_LITERALS["postgresql"]) + inspector = inspect(conn) + tables = set(inspector.get_table_names()) + columns = {t: {c["name"] for c in inspector.get_columns(t)} for t in tables} + # IF NOT EXISTS:PostgreSQL 上多个进程同时启动,各自的检查和补列之间别人可能已经补上了(init_db 另外拿着锁) + add = "ADD COLUMN IF NOT EXISTS" if dialect.name == "postgresql" else "ADD COLUMN" for table, column, ddl in _COLUMN_MIGRATIONS: - cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))} + cols = columns.get(table) if cols and column not in cols: - conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}")) - for name, table, columns in _INDEX_MIGRATIONS: - wanted = (columns,) if isinstance(columns, str) else columns - cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))} + conn.execute(text(f"ALTER TABLE {table} {add} {column} {_column_ddl(dialect, table, column, ddl)}")) + cols.add(column) + for name, table, wanted in _INDEX_MIGRATIONS: + wanted = (wanted,) if isinstance(wanted, str) else wanted + cols = columns.get(table) or set() if all(c in cols for c in wanted): conn.execute(text(f"CREATE INDEX IF NOT EXISTS {name} ON {table} ({', '.join(wanted)})")) - for name, table, columns, where, dedupe in _UNIQUE_INDEX_MIGRATIONS: - cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))} - exists = conn.execute(text("SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = :n"), {"n": name}).first() - if exists or not all(c in cols for c in columns): + for name, table, wanted, where, dedupe in _UNIQUE_INDEX_MIGRATIONS: + cols = columns.get(table) or set() + if not all(c in cols for c in wanted): continue - conn.execute(text(dedupe)) - conn.execute(text(f"CREATE UNIQUE INDEX {name} ON {table} ({', '.join(columns)})" - + (f" WHERE {where}" if where else ""))) + if any(ix["name"] == name for ix in inspect(conn).get_indexes(table)): + continue + conn.execute(text(dedupe.format(**literals))) + conn.execute(text(f"CREATE UNIQUE INDEX {name} ON {table} ({', '.join(wanted)})" + + (f" WHERE {where.format(**literals)}" if where else ""))) + + +#: init_db 在 PostgreSQL 上拿的会话级咨询锁的号:几个进程同时启动时,建表和迁移一个一个来(同一个号别处别用) +_INIT_LOCK = 0x4167656E744C6162 # "AgentLab" + + +async def create_schema(target: AsyncEngine) -> None: + """建表、补列补索引(create_all 只建新表,老库的新列靠 _migrate)。两种库同一套;copy-to-postgres 建目标库也用它。""" + from app.db import models # noqa: F401 确保模型已注册到 metadata + + async with target.begin() as conn: + if conn.dialect.name == "postgresql": + # 事务结束自动放锁。CREATE TABLE IF NOT EXISTS 在 PostgreSQL 上并发执行也会撞(系统表的唯一约束) + from sqlalchemy import text + + await conn.execute(text("SELECT pg_advisory_xact_lock(:k)"), {"k": _INIT_LOCK}) + await conn.run_sync(Base.metadata.create_all) + await conn.run_sync(_migrate) async def init_db() -> None: - from app.db import models # noqa: F401 确保模型已注册到 metadata - - async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) - await conn.run_sync(_migrate) + logger.info("应用库:%s", describe_database()) + await create_schema(engine) diff --git a/backend/app/db/models.py b/backend/app/db/models.py index 860819a1..f10efdea 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -362,7 +362,8 @@ class CatalogDraftJob(Base, TimestampMixin): __table_args__ = ( Index("ix_catalog_draft_jobs_source_created", "source_id", "created_at"), Index("ux_catalog_draft_jobs_active", "source_id", unique=True, - sqlite_where=text("status IN ('queued', 'running')")), + sqlite_where=text("status IN ('queued', 'running')"), + postgresql_where=text("status IN ('queued', 'running')")), ) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) @@ -450,7 +451,8 @@ class CatalogProfileJob(Base, TimestampMixin): __table_args__ = ( Index("ix_catalog_profile_jobs_source_created", "source_id", "created_at"), Index("ux_catalog_profile_jobs_active", "source_id", unique=True, - sqlite_where=text("status IN ('queued', 'running')")), + sqlite_where=text("status IN ('queued', 'running')"), + postgresql_where=text("status IN ('queued', 'running')")), ) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) @@ -904,7 +906,8 @@ class Conversation(Base, TimestampMixin): # 一个人在一个应用里只有一段(未归档的)助手对话(D-e)。存量库由 db/base.py 的 _UNIQUE_INDEX_MIGRATIONS 补建。 # created_by 为空(关登录又没填署名)的不受约束:SQLite 的唯一索引里 NULL 互不相等,那种情况靠接口里的进程内锁 __table_args__ = (Index("ux_conversations_app_owner", "app_id", "created_by", unique=True, - sqlite_where=text("kind = 'app' AND archived = 0")),) + sqlite_where=text("kind = 'app' AND archived = 0"), + postgresql_where=text("kind = 'app' AND archived = false")),) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) title: Mapped[str] = mapped_column(String(200), default="") diff --git a/backend/app/db/models_sources.py b/backend/app/db/models_sources.py index 5495855e..2423fa50 100644 --- a/backend/app/db/models_sources.py +++ b/backend/app/db/models_sources.py @@ -19,10 +19,10 @@ from __future__ import annotations from datetime import datetime from typing import Any -from sqlalchemy import JSON, Boolean, ForeignKey, Index, Integer, String, Text, UniqueConstraint, text +from sqlalchemy import Boolean, ForeignKey, Index, Integer, String, Text, UniqueConstraint, text from sqlalchemy.orm import Mapped, mapped_column -from app.db.base import Base, UTCDateTime, new_id, utcnow +from app.db.base import Base, PortableJSON, UTCDateTime, new_id, utcnow class SourceObject(Base): @@ -92,7 +92,8 @@ class IntrospectJob(Base): __table_args__ = ( Index("ix_introspect_jobs_source_created", "source_id", "created_at"), Index("ux_introspect_jobs_active", "source_id", unique=True, - sqlite_where=text("status IN ('queued', 'running')")), + sqlite_where=text("status IN ('queued', 'running')"), + postgresql_where=text("status IN ('queued', 'running')")), ) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id) @@ -219,8 +220,8 @@ class SourceChange(Base): seen_count: Mapped[int] = mapped_column(Integer, default=1) detected_by: Mapped[dict[str, Any]] = mapped_column(default=dict) basis: Mapped[dict[str, Any] | None] = mapped_column(default=None) - before: Mapped[Any] = mapped_column(JSON, nullable=True, default=None) - after: Mapped[Any] = mapped_column(JSON, nullable=True, default=None) + before: Mapped[Any] = mapped_column(PortableJSON, nullable=True, default=None) + after: Mapped[Any] = mapped_column(PortableJSON, nullable=True, default=None) affects: Mapped[dict[str, Any] | None] = mapped_column(default=None) struct_version_before: Mapped[int | None] = mapped_column(Integer, default=None) struct_version_after: Mapped[int | None] = mapped_column(Integer, default=None) diff --git a/backend/app/db/portable.py b/backend/app/db/portable.py new file mode 100644 index 00000000..4d643d1d --- /dev/null +++ b/backend/app/db/portable.py @@ -0,0 +1,235 @@ +"""应用库上 SQLite、PostgreSQL 写法不一样的那几样,收成按方言编译的表达式。 + +只给应用库(agentlab.db / AGENTLAB_DATABASE_URL)用;发给业务数据源的 SQL 不走这里。 + +**按 JSON 列里的键筛选。** SQLite 用 json_extract:取出来是 SQL 标量(字符串是文本、数字是数字、true 是 1、 +没有这个键或值为 null 是 NULL),和 Python 值比是带类型的比较(数字 3 不等于字符串 "3")。PostgreSQL 的 JSON 列是 +存原文的 json(db/base.PortableJSON),用 -> / ->> 取,要比较的只把取出来的那一小段转成 JSONB: + +- json_eq:PostgreSQL 上按 JSONB 相等比(`(col -> 'k')::jsonb = '3'::jsonb`),同样带类型,和 SQLite 一致; +- json_text:只要文本(字符串原样、数字是它的写法),PostgreSQL 上是 ->>; +- json_flag:标记位「真」(true、非零的数):没有这个键、null、false、0 都算假; +- json_is_string:这个键的值是不是字符串(SQLite 的 json_type 叫 text,PostgreSQL 的 jsonb_typeof 叫 string)。 + +键名只收标识符(字母、数字、下划线),直接写进 SQL:键都是代码里的常量,写进去语句缓存也按它区分。 + +**插入,撞了唯一键就什么都不做。** insert_ignore 按会话连着的库挑 sqlite / postgresql 的 insert,两边的 +on_conflict_do_nothing 是同一个接口。 + +**撞的是哪个唯一约束。** unique_violation 认 IntegrityError:SQLite 的报错是「UNIQUE constraint failed: 表.列」, +PostgreSQL(asyncpg)的是 23505,表名、DETAIL 里的「Key (列)=」分别带着。不靠约束名(命名规则改了也认得出)。 +""" + +from __future__ import annotations + +import json +import re +from typing import Any + +from sqlalchemy import Boolean, String, bindparam, literal_column +from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.functions import FunctionElement +from sqlalchemy.types import TypeDecorator + +_KEY = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") + + +def _key(name: str) -> Any: + if not _KEY.match(name): + raise ValueError(f"JSON 键名只能是标识符:{name!r}") + return literal_column(name) + + +class _JsonScalar(TypeDecorator): + """json_eq 右边的值:SQLite 原样绑定(和 json_extract 取出的标量比),PostgreSQL 绑定成 JSON 文本再转 JSONB。""" + + impl = String + cache_ok = True + + def process_bind_param(self, value: Any, dialect: Any) -> Any: + if dialect.name == "postgresql": + return json.dumps(value, ensure_ascii=False) + return value + + +class json_text(FunctionElement[str]): + """JSON 列顶层键的值,当文本用。""" + + type = String() + inherit_cache = True + name = "json_text" + + def __init__(self, column: Any, key: str) -> None: + super().__init__(column, _key(key)) + + +class _JsonEq(FunctionElement[bool]): + type = Boolean() + inherit_cache = True + name = "json_eq" + + def __init__(self, column: Any, key: str, value: Any) -> None: + super().__init__(column, _key(key), bindparam(None, value, type_=_JsonScalar())) + + +class _JsonNull(FunctionElement[bool]): + type = Boolean() + inherit_cache = True + name = "json_null" + + def __init__(self, column: Any, key: str) -> None: + super().__init__(column, _key(key)) + + +def json_eq(column: Any, key: str, value: Any) -> Any: + """JSON 列顶层键的值等于 value(带类型的比较,见模块说明)。value 为 None 时是「没有这个键或值为 null」。 + + None 用另一个表达式:语句缓存不按绑定参数的值区分,同一个表达式不能因为值是 None 编译成两样。""" + return _JsonNull(column, key) if value is None else _JsonEq(column, key, value) + + +class json_flag(FunctionElement[bool]): + """JSON 列顶层键是「真」的标记位:true 或非零的数(字符串也算真,同 SQLite)。没有这个键、null、false、0 是假。""" + + type = Boolean() + inherit_cache = True + name = "json_flag" + + def __init__(self, column: Any, key: str) -> None: + super().__init__(column, _key(key)) + + +class json_is_string(FunctionElement[bool]): + """JSON 列顶层键的值是字符串。""" + + type = Boolean() + inherit_cache = True + name = "json_is_string" + + def __init__(self, column: Any, key: str) -> None: + super().__init__(column, _key(key)) + + +def _parts(element: Any, compiler: Any, **kw: Any) -> tuple[str, str, list[Any]]: + column, key, *rest = list(element.clauses) + return compiler.process(column, **kw), key.name, rest + + +# ---- SQLite(缺省)---- + + +@compiles(json_text) +def _json_text_sqlite(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"json_extract({col}, '$.{key}')" + + +@compiles(_JsonEq) +def _json_eq_sqlite(element: Any, compiler: Any, **kw: Any) -> str: + col, key, (value,) = _parts(element, compiler, **kw) + return f"(json_extract({col}, '$.{key}') = {compiler.process(value, **kw)})" + + +@compiles(_JsonNull) +def _json_null_sqlite(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"(json_extract({col}, '$.{key}') IS NULL)" + + +@compiles(json_flag) +def _json_flag_sqlite(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"(coalesce(json_extract({col}, '$.{key}'), 0) != 0)" + + +@compiles(json_is_string) +def _json_is_string_sqlite(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"(json_type({col}, '$.{key}') = 'text')" + + +# ---- PostgreSQL(应用库的 JSON 列是 json)---- + + +def _got(col: str, key: str) -> str: + """取出的那一段转成 JSONB:json 没有相等运算符。列本身是 JSONB 也照样成立。""" + return f"CAST(({col} -> '{key}') AS JSONB)" + + +@compiles(json_text, "postgresql") +def _json_text_pg(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"({col} ->> '{key}')" + + +@compiles(_JsonEq, "postgresql") +def _json_eq_pg(element: Any, compiler: Any, **kw: Any) -> str: + col, key, (value,) = _parts(element, compiler, **kw) + return f"({_got(col, key)} = CAST({compiler.process(value, **kw)} AS JSONB))" + + +@compiles(_JsonNull, "postgresql") +def _json_null_pg(element: Any, compiler: Any, **kw: Any) -> str: + # json_extract 对「没有这个键」和「值为 null」都给 NULL;-> 对后者给的是 JSON 的 null + col, key, _ = _parts(element, compiler, **kw) + return f"(coalesce({_got(col, key)}, 'null'::jsonb) = 'null'::jsonb)" + + +@compiles(json_flag, "postgresql") +def _json_flag_pg(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"(coalesce({_got(col, key)}, 'null'::jsonb) NOT IN ('null'::jsonb, 'false'::jsonb, '0'::jsonb))" + + +@compiles(json_is_string, "postgresql") +def _json_is_string_pg(element: Any, compiler: Any, **kw: Any) -> str: + col, key, _ = _parts(element, compiler, **kw) + return f"(jsonb_typeof({_got(col, key)}) = 'string')" + + +# ---- 插入,撞了唯一键就跳过 ---- + + +def dialect_of(session: Any) -> str: + """会话(AsyncSession / Session)连着的库的方言名。""" + return session.get_bind().dialect.name + + +def insert_ignore(session: Any, table: Any) -> Any: + """按会话连着的库挑方言的 insert,调用方接着 .values(...).on_conflict_do_nothing(index_elements=[...])。""" + if dialect_of(session) == "postgresql": + from sqlalchemy.dialects.postgresql import insert + else: + from sqlalchemy.dialects.sqlite import insert + return insert(table) + + +# ---- 撞的是哪个唯一约束 ---- + +_SQLITE_UNIQUE = "UNIQUE constraint failed:" +_PG_KEY = re.compile(r"Key \(([^)]*)\)=") + + +def _driver_error(error: BaseException) -> Any: + """SQLAlchemy 的异常 → asyncpg 自己的异常(在 orig 的 __cause__ 上;orig 只抄了 sqlstate,没有表名和 DETAIL)。""" + seen: Any = getattr(error, "orig", None) or error + for _ in range(4): + if getattr(seen, "sqlstate", None) and hasattr(seen, "table_name"): + return seen + seen = getattr(seen, "__cause__", None) + if seen is None: + break + return None + + +def unique_violation(error: BaseException, table: str, *columns: str) -> bool: + """这个 IntegrityError 是不是撞了 table 上正好这几列的唯一约束(主键、唯一索引都算)。""" + text = str(getattr(error, "orig", None) or error) + if _SQLITE_UNIQUE in text: + failed = text.split(_SQLITE_UNIQUE, 1)[1].splitlines()[0] + return {part.strip() for part in failed.split(",")} == {f"{table}.{c}" for c in columns} + cause = _driver_error(error) + if cause is None or getattr(cause, "sqlstate", None) != "23505" or getattr(cause, "table_name", None) != table: + return False + match = _PG_KEY.search(getattr(cause, "detail", None) or "") + return bool(match) and [c.strip().strip('"') for c in match.group(1).split(",")] == list(columns) diff --git a/backend/app/db/postgres.py b/backend/app/db/postgres.py new file mode 100644 index 00000000..74318aa3 --- /dev/null +++ b/backend/app/db/postgres.py @@ -0,0 +1,76 @@ +"""应用库是 PostgreSQL 时,让它在 SQLite 上写成的那些假设下照样成立。 + +只装在应用库那个引擎的方言实例上(db/base.make_engine),业务数据源的 PostgreSQL 引擎不受影响。 + +**建表的列类型。** 应用是在 SQLite 上写成的:SQLite 不查 VARCHAR 的长度,INTEGER 是 64 位。PostgreSQL 上 +VARCHAR(n) 超长直接报错、INTEGER 只有 32 位。模型里的 String(n) 是写给人看的「大致多长」,从来没有被库约束过: +逐列核对只核得到测试走到的值,漏一处就是线上的写入失败,从 SQLite 导过来的存量数据里也可能已经有超长的。所以 +应用库上 String(n) 建成不限长的 VARCHAR(PostgreSQL 里 varchar、varchar(n)、text 存储和性能一样,限长只多一次 +检查)。CAST(x AS VARCHAR(n)) 这种查询里的写法不动——那是截断的意思。Integer 列在 PostgreSQL 上是 BIGINT: +那是在模型这一层换的(db/base._portable_integer),绑定参数的类型转换跟着列走,大数不在半路溢出。 + +**ORDER BY 里 NULL 的位置。** SQLite 升序 NULL 在前、降序 NULL 在后,PostgreSQL 正好相反。应用里两百多处排序, +可空的列(开始时间、到期时间、发布时间……)排出来的先后要和 SQLite 一致:可空列上没写 nulls_first / nulls_last +的 ASC 补 NULLS FIRST、DESC 补 NULLS LAST,ORDER BY 里直接写列名(隐含 ASC)的同样补上。不可空的列不动—— +补了会让索引用不上,又没有 NULL 可排。窗口函数 OVER (ORDER BY …) 里同样补。原生 text() 里写的 ORDER BY 不经过 +这里,要自己写明(db/base 的去重语句里就写了)。 +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import Column +from sqlalchemy.dialects.postgresql.asyncpg import PGCompiler_asyncpg +from sqlalchemy.dialects.postgresql.base import PGTypeCompiler +from sqlalchemy.sql import operators +from sqlalchemy.sql.elements import ClauseList, UnaryExpression + +#: 外层已经写了 nulls_first / nulls_last:里面那一层 ASC / DESC 不再补 +_EXPLICIT = "_agentlab_explicit_nulls" + + +def _nullable(element: Any) -> bool: + return isinstance(element, Column) and bool(element.nullable) and not element.primary_key + + +class AppTypeCompiler(PGTypeCompiler): + def visit_string(self, type_: Any, **kw: Any) -> str: + if type_.length and isinstance(kw.get("type_expression"), Column): + return "VARCHAR" + return super().visit_string(type_, **kw) + + +def _ordered(clauses: Any) -> list[Any]: + """ORDER BY 里直接写的可空列补成 ASC(visit_unary 再给它补 NULLS FIRST)。""" + return [c.asc() if _nullable(c) else c for c in clauses] + + +class AppCompiler(PGCompiler_asyncpg): + def order_by_clause(self, select: Any, **kw: Any) -> str: + order_by = self._generate_delimited_list(_ordered(select._order_by_clauses), ", ", **kw) + return " ORDER BY " + order_by if order_by else "" + + def visit_over(self, over: Any, **kw: Any) -> str: + clauses = list(over.order_by) if over.order_by is not None else [] + if any(_nullable(c) for c in clauses): + # 在副本上换:语句对象可能被缓存、被别处共用,不能改原件 + over = over._clone() + over.order_by = ClauseList(*_ordered(clauses)) + return super().visit_over(over, **kw) + + def visit_unary(self, unary: UnaryExpression, **kw: Any) -> str: + modifier = unary.modifier + if modifier in (operators.nulls_first_op, operators.nulls_last_op): + return super().visit_unary(unary, **{**kw, _EXPLICIT: True}) + explicit = kw.pop(_EXPLICIT, False) + text = super().visit_unary(unary, **kw) + if not explicit and modifier in (operators.asc_op, operators.desc_op) and _nullable(unary.element): + text += " NULLS FIRST" if modifier is operators.asc_op else " NULLS LAST" + return text + + +def install(dialect: Any) -> None: + """装到应用库引擎的方言实例上。""" + dialect.statement_compiler = AppCompiler + dialect.type_compiler_instance = dialect.type_compiler = AppTypeCompiler(dialect) diff --git a/backend/app/db/urls.py b/backend/app/db/urls.py new file mode 100644 index 00000000..c1b8855f --- /dev/null +++ b/backend/app/db/urls.py @@ -0,0 +1,57 @@ +"""应用库、断点库的地址:规范写法、换算成 psycopg 用的连接串、写进日志的样子(不带密码)。 + +不依赖 app.core.config:配置的校验器要调这里,这里再导 config 就成了循环导入。 +""" + +from __future__ import annotations + +from sqlalchemy.engine import URL, make_url +from sqlalchemy.exc import ArgumentError + +#: 应用库认的写法。asyncpg 是 SQLAlchemy 这边的驱动;postgresql:// / postgres:// 是常见的通用写法,一律换成 asyncpg +_APP_SCHEMES = {"postgresql+asyncpg", "postgresql", "postgres"} +#: 断点库认的写法:psycopg 只认 libpq 的 postgresql:// / postgres://;带驱动名的 SQLAlchemy 写法也收,去掉驱动名 +_CHECKPOINT_SCHEMES = {"postgresql", "postgres", "postgresql+psycopg", "postgresql+asyncpg"} + +#: asyncpg 和 libpq 对 SSL 的参数名不一样:asyncpg 是 ssl=require,libpq 是 sslmode=require(取值同一套) +_ASYNCPG_TO_LIBPQ = {"ssl": "sslmode"} + + +def _parse(raw: str, schemes: set[str], example: str) -> URL: + try: + url = make_url(raw) + except ArgumentError: + raise ValueError(f"认不出这个地址,应当写成 {example}") from None + if url.drivername not in schemes: + raise ValueError(f"只认 PostgreSQL({example}),不认 {url.drivername}://;用 SQLite 就不要设这一项") + if not url.database: + raise ValueError(f"地址里要写库名,例如 {example}") + return url + + +def normalize_database_url(raw: str) -> str: + """AGENTLAB_DATABASE_URL → postgresql+asyncpg://…(密码原样保留)。认不出的写法抛 ValueError。""" + url = _parse(raw, _APP_SCHEMES, "postgresql+asyncpg://用户:密码@主机:5432/库名") + return url.set(drivername="postgresql+asyncpg").render_as_string(hide_password=False) + + +def normalize_checkpoint_url(raw: str) -> str: + """AGENTLAB_CHECKPOINT_URL → libpq 认的 postgresql://…。认不出的写法抛 ValueError。""" + url = _parse(raw, _CHECKPOINT_SCHEMES, "postgresql://用户:密码@主机:5432/库名") + return libpq_conninfo(url.render_as_string(hide_password=False)) + + +def libpq_conninfo(database_url: str) -> str: + """SQLAlchemy 的地址 → psycopg 的连接串:去掉驱动名,asyncpg 写法的参数换成 libpq 的。""" + url = make_url(database_url) + query = {_ASYNCPG_TO_LIBPQ.get(k, k): v for k, v in url.query.items()} + return url.set(drivername="postgresql", query=query).render_as_string(hide_password=False) + + +def describe(database_url: str) -> str: + """写进日志的样子:主机、端口、库名、用户,不写密码。""" + url = make_url(database_url) + host = url.host or "本机套接字" + port = f":{url.port}" if url.port else "" + user = f",用户 {url.username}" if url.username else "" + return f"{host}{port}/{url.database}{user}" diff --git a/backend/app/engine/checkpoints.py b/backend/app/engine/checkpoints.py new file mode 100644 index 00000000..bd267ccf --- /dev/null +++ b/backend/app/engine/checkpoints.py @@ -0,0 +1,87 @@ +"""LangGraph 的断点库:应用库是 SQLite 时用 AsyncSqliteSaver(数据目录里的 checkpoints.db,和业务表分开免得抢写锁), +是 PostgreSQL 时用官方的 AsyncPostgresSaver(psycopg 3 + 连接池),放在同一个库里,AGENTLAB_CHECKPOINT_URL 另给就放那里。 + +它建的四张表(checkpoints、checkpoint_blobs、checkpoint_writes、checkpoint_migrations)和应用的表不重名 +(test_db_portable 守着),所以同一个库、同一个 schema 就行。 + +应用里读断点一律走 saver 的公开接口(aget_state / aget_state_history / aget_tuple / alist),不直接查这几张表: +两种库的表结构完全不一样(SQLite 整份断点是一个 msgpack,PostgreSQL 拆成 JSONB 加 blob)。 +""" + +from __future__ import annotations + +import contextlib +from collections.abc import AsyncIterator +from pathlib import Path +from typing import Any + +from langgraph.checkpoint.base import BaseCheckpointSaver + +from app.core.config import settings + +#: 建断点表时拿的会话级咨询锁的号:几个进程同时启动,各自的 setup(建表、记迁移版本)一个一个来 +_SETUP_LOCK = 0x4167656E74436B70 # "AgentCkp" + + +def _pool_size() -> int: + # 运行里每个超步写一次断点、每个任务写一次 pending writes,都是一来一回的短操作。并发运行多时多给几条 + return max(4, min(settings.max_concurrent_runs, 20)) + + +@contextlib.asynccontextmanager +async def open_postgres_saver(conninfo: str, *, max_size: int | None = None, + application_name: str = "agentlab") -> AsyncIterator[BaseCheckpointSaver]: + """PostgreSQL 的断点库:开连接池、建表(幂等),退出时关池子。application_name 同应用库(db/base.APPLICATION_NAME)。""" + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + from psycopg.rows import dict_row + from psycopg_pool import AsyncConnectionPool + + # 连接参数照 AsyncPostgresSaver.from_conn_string:自动提交(它不自己管事务)、语句第一次执行就预备、行按 dict 取 + pool = AsyncConnectionPool( + conninfo, min_size=1, max_size=max_size or _pool_size(), open=False, + kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row, + "application_name": application_name}, + check=AsyncConnectionPool.check_connection, name="agentlab-checkpoints") + await pool.open(wait=True) + try: + saver = AsyncPostgresSaver(pool) # type: ignore[arg-type] + async with pool.connection() as conn: + await conn.execute("SELECT pg_advisory_lock(%s)", (_SETUP_LOCK,)) + try: + await saver.setup() + finally: + await conn.execute("SELECT pg_advisory_unlock(%s)", (_SETUP_LOCK,)) + yield saver + finally: + await pool.close() + + +@contextlib.asynccontextmanager +async def open_sqlite_saver(path: Path | str) -> AsyncIterator[BaseCheckpointSaver]: + from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + + async with AsyncSqliteSaver.from_conn_string(str(path)) as saver: + yield saver + + +def open_saver() -> contextlib.AbstractAsyncContextManager[BaseCheckpointSaver]: + """按配置开这个进程的断点库。""" + conninfo = settings.checkpoint_conninfo + if conninfo: + return open_postgres_saver(conninfo) + return open_sqlite_saver(settings.checkpoint_path) + + +def kind() -> str: + """"sqlite" 或 "postgresql":断点库放在哪种库里。""" + return "postgresql" if settings.checkpoint_conninfo else "sqlite" + + +async def count(saver: Any, thread_id: str | None = None) -> tuple[int, int]: + """(断点条数, pending writes 条数)。thread_id 为空数全部。copy-to-postgres 核对用,走公开接口。""" + config = {"configurable": {"thread_id": thread_id}} if thread_id else None + checkpoints = writes = 0 + async for tup in saver.alist(config): + checkpoints += 1 + writes += len(tup.pending_writes or ()) + return checkpoints, writes diff --git a/backend/app/engine/runner.py b/backend/app/engine/runner.py index 1b3d020c..b33e472a 100644 --- a/backend/app/engine/runner.py +++ b/backend/app/engine/runner.py @@ -10,7 +10,7 @@ from collections.abc import Callable, Sequence from datetime import datetime, timezone from typing import Any -from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.errors import GraphRecursionError from langgraph.types import Command from sqlalchemy import case, func, select, update @@ -71,7 +71,7 @@ class RunManager: def __init__(self) -> None: self._tasks: dict[str, asyncio.Task[None]] = {} - self._checkpointer: AsyncSqliteSaver | None = None + self._checkpointer: BaseCheckpointSaver | None = None self._cm: Any = None self._semaphore = asyncio.Semaphore(settings.max_concurrent_runs) #: 正在关停。此时收到的取消来自关停而不是用户,运行要记成可恢复的中断 @@ -86,7 +86,10 @@ class RunManager: async def setup(self) -> None: self._closing = False settings.ensure_dirs() - self._cm = AsyncSqliteSaver.from_conn_string(str(settings.checkpoint_path)) + # SQLite:数据目录里的 checkpoints.db;PostgreSQL:同一个库(或 AGENTLAB_CHECKPOINT_URL),见 engine/checkpoints.py + from app.engine.checkpoints import open_saver + + self._cm = open_saver() self._checkpointer = await self._cm.__aenter__() # 上次进程被强杀时留下的记录。running 和 queued 要分开处理: # running 至少跑过一步、有 checkpoint,可以续;queued 还在排队, @@ -163,7 +166,7 @@ class RunManager: await self._cm.__aexit__(None, None, None) @property - def checkpointer(self) -> AsyncSqliteSaver: + def checkpointer(self) -> BaseCheckpointSaver: if self._checkpointer is None: raise RuntimeError("RunManager 还没初始化") return self._checkpointer @@ -889,6 +892,15 @@ class RunManager: try: await self._finalize(run_id, status, error, final_state, elapsed, terminal=terminal, error_node=error_node, end_reason=end_reason) + # 记 action.unknown、安排核实也算收尾:做完才把自己摘掉。以前摘掉之后才做,等运行「不在跑了」 + # 再去读事件的一方(接口、测试)有一小段读不到这一条;库慢一点(PostgreSQL 满载时)就读不到 + if pending_action is not None: + from app.process.actions import on_unknown + + try: + await on_unknown(pending_action) + except Exception: # noqa: BLE001 - 运行已经收成挂起;没安排上的核实,下次启动时 actions.startup 补上 + logger.exception("run %s 的动作 %s 结果未知,安排核实时出错", run_id, pending_action) finally: # 收尾被打断(关停时的取消、库连接出错)也得把自己摘掉。以前写在 # _finalize 后面,一被打断就漏掉,留下的任务让下一次 shutdown 的 @@ -896,13 +908,6 @@ class RunManager: self._finalizing.discard(run_id) self._tasks.pop(run_id, None) self._stopped_by.pop(run_id, None) - if pending_action is not None: - from app.process.actions import on_unknown - - try: - await on_unknown(pending_action) - except Exception: # noqa: BLE001 - 运行已经收成挂起;没安排上的核实,下次启动时 actions.startup 补上 - logger.exception("run %s 的动作 %s 结果未知,安排核实时出错", run_id, pending_action) # 运行到终态就把沙箱会话收掉。thread_id 每个 run 都是新的, # 不收的话每跑一次带代码节点的图就多一台常驻 microVM(几百 MB), # 而且没有任何自动回收路径会碰它——闲置回收只在"下次有人执行 diff --git a/backend/app/files/store.py b/backend/app/files/store.py index f3f74aa5..4b466005 100644 --- a/backend/app/files/store.py +++ b/backend/app/files/store.py @@ -294,7 +294,7 @@ def accept_problem(setting: dict[str, Any] | None, *, name: str, mime: str, size # -------------------------------------------------------------------------- #: 还没绑到记录上:JSON 列显式写过 None 的存的是 JSON 的 null,不是 SQL 的 NULL,两种都算 -_NO_RECORD = text("(files.bound_record IS NULL OR files.bound_record = 'null')") +_NO_RECORD = text("(files.bound_record IS NULL OR CAST(files.bound_record AS TEXT) = 'null')") def file_fields(form: Any) -> list[str]: diff --git a/backend/app/ledger/issue.py b/backend/app/ledger/issue.py index 66f701bd..0c451172 100644 --- a/backend/app/ledger/issue.py +++ b/backend/app/ledger/issue.py @@ -44,6 +44,7 @@ from sqlalchemy import func, select from app.db.base import SessionLocal from app.db.models import Run, Workflow +from app.db.portable import json_eq logger = logging.getLogger(__name__) @@ -150,9 +151,9 @@ async def issued_today(session: Any) -> int: async def issued_runs(session: Any, run_id: str) -> list[Run]: """这次运行按口径出具过的正式运行(origin.source_run 是它),新的在前。origin_kind 有索引,先按它筛再看 origin - (json_extract,同 auth/keys 认幂等键的写法)。""" + (db/portable.json_eq,同 auth/keys 认幂等键的写法)。""" rows = (await session.execute(select(Run).where( - Run.origin_kind == ORIGIN, func.json_extract(Run.origin, "$.source_run") == run_id) + Run.origin_kind == ORIGIN, json_eq(Run.origin, "source_run", run_id)) .order_by(Run.created_at.desc()))).scalars() return list(rows) diff --git a/backend/app/ledger/recall.py b/backend/app/ledger/recall.py index 62bf2dbe..3227c66f 100644 --- a/backend/app/ledger/recall.py +++ b/backend/app/ledger/recall.py @@ -66,11 +66,12 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from typing import Any -from sqlalchemy import func, select +from sqlalchemy import func, not_, select from app.core.events import EventType from app.db.base import SessionLocal, utcnow from app.db.models import Run, RunEvent, Workflow +from app.db.portable import json_eq, json_flag, json_text from app.ledger import uses as ledger_uses logger = logging.getLogger(__name__) @@ -225,10 +226,9 @@ def source_cards(graph: Any) -> list[str]: async def recall_runs(session: Any, workflow_id: str, new_version: int | None = None) -> list[Run]: """以 W(新版本 N)发起过的召回运行(试算的不算),新的在前。origin_kind 有索引,先按它筛再看 origin。""" stmt = select(Run).where(Run.origin_kind == ORIGIN, - func.json_extract(Run.origin, "$.workflow_id") == workflow_id, - func.coalesce(func.json_extract(Run.origin, "$.probe"), 0) == 0) + json_eq(Run.origin, "workflow_id", workflow_id), not_(json_flag(Run.origin, "probe"))) if new_version is not None: - stmt = stmt.where(func.json_extract(Run.origin, "$.to_version") == int(new_version)) + stmt = stmt.where(json_eq(Run.origin, "to_version", int(new_version))) return list((await session.execute(stmt.order_by(Run.created_at.desc()))).scalars()) @@ -567,9 +567,8 @@ async def version_recall_counts(session: Any, workflow_id: str, version: int) -> state, _ = await outcome_of(session, run) changed += state == "changed" as_old = int(await session.scalar(select(func.count()).select_from(Run).where( - Run.origin_kind == ORIGIN, func.json_extract(Run.origin, "$.workflow_id") == workflow_id, - func.json_extract(Run.origin, "$.from_version") == int(version), - func.coalesce(func.json_extract(Run.origin, "$.probe"), 0) == 0)) or 0) + Run.origin_kind == ORIGIN, json_eq(Run.origin, "workflow_id", workflow_id), + json_eq(Run.origin, "from_version", int(version)), not_(json_flag(Run.origin, "probe")))) or 0) return {"groups": len(as_new), "changed": changed, "recalled_from": as_old} @@ -986,7 +985,7 @@ async def answer_marks(session: Any, run_id: str, principal: Any) -> dict[str, A select(Task, Run.origin).join(Run, Run.id == Task.run_id) .where(Task.kind == TASK_KIND, Task.status == "pending", (Task.run_id == run_id) | ((Run.origin_kind == ORIGIN) - & (func.instr(func.json_extract(Run.origin, "$.uses"), run_id) > 0))) + & json_text(Run.origin, "uses").contains(run_id, autoescape=True))) .order_by(Task.created_at, Task.id))).all() mine = await _assignments(session, [t.id for t, _ in candidates]) tasks: list[dict[str, Any]] = [] @@ -1134,8 +1133,8 @@ async def _explore_snapshots(session: Any, source: Any, tables: Iterable[str], * stmt = (select(Artifact.id, Artifact.meta, Run.id, Run.finished_at).join(Run, Run.id == Artifact.run_id) .where(Artifact.kind == QUERY_SNAPSHOT_KIND, Run.run_class == "exploratory", Run.status == "succeeded", Run.manifest_hash.is_not(None), counted_runs(), - or_(func.json_extract(Artifact.meta, "$.source_id") == source.id, - func.json_extract(Artifact.meta, "$.source") == source.name)) + or_(json_eq(Artifact.meta, "source_id", source.id), + json_eq(Artifact.meta, "source", source.name))) .order_by(Run.finished_at.desc(), Run.id, Artifact.id)) cut = _aware(before) offset = 0 @@ -1473,10 +1472,10 @@ async def _probe_calibers(change: Mapping[str, Any]) -> list[str]: continue async with SessionLocal() as session: done = await session.scalar(select(func.count()).select_from(Run).where( - Run.origin_kind == ORIGIN, func.json_extract(Run.origin, "$.probe") == 1, - func.json_extract(Run.origin, "$.workflow_id") == caliber.workflow_id, - func.json_extract(Run.origin, "$.version") == caliber.version, - func.json_extract(Run.origin, "$.change") == change.get("id"))) + Run.origin_kind == ORIGIN, json_flag(Run.origin, "probe"), + json_eq(Run.origin, "workflow_id", caliber.workflow_id), + json_eq(Run.origin, "version", caliber.version), + json_eq(Run.origin, "change", change.get("id")))) if done: continue latest = (await session.execute(select(Run.input).where( diff --git a/backend/app/ledger/reconcile.py b/backend/app/ledger/reconcile.py index db2edfff..ea114515 100644 --- a/backend/app/ledger/reconcile.py +++ b/backend/app/ledger/reconcile.py @@ -143,10 +143,10 @@ def actions_of(content: dict[str, Any], names: dict[str, str], hide: Any = None) async def reconciled_runs(session: Any, run_id: str) -> list[Run]: """对这次运行对过的账(origin.source_run 是它),新的在前。""" - from sqlalchemy import func + from app.db.portable import json_eq rows = (await session.execute(select(Run).where( - Run.origin_kind == ORIGIN, func.json_extract(Run.origin, "$.source_run") == run_id) + Run.origin_kind == ORIGIN, json_eq(Run.origin, "source_run", run_id)) .order_by(Run.created_at.desc()))).scalars() return list(rows) diff --git a/backend/app/process/tasks.py b/backend/app/process/tasks.py index 0080edb5..4d62c7eb 100644 --- a/backend/app/process/tasks.py +++ b/backend/app/process/tasks.py @@ -35,6 +35,7 @@ from app.db.base import SessionLocal, utcnow from app.db.models import ( App, Approval, AppVersion, AuditEvent, Group, GroupMember, Run, Task, TaskAssignment, User, ) +from app.db.portable import json_text from app.engine.lint_process import normalize_policy, parse_duration from app.process.events import record @@ -913,7 +914,7 @@ async def tick(now: datetime | None = None) -> list[str]: """调度循环每一拍:提醒、升级、超时驳回(永不自动批准)。交回这一拍处理过的待办 id(生效可查:办了哪几条, 不只是几条;测试按自己造的那条断言,不受同一个库里别的待办影响)。试运行的待办不处理(只显示时限)。""" now = now or utcnow() - action = func.json_extract(Task.on_timeout, "$.action") + action = json_text(Task.on_timeout, "action") trial = select(Run.id).where(Run.origin_kind == TRIAL_ORIGIN) async with SessionLocal() as session: due = list((await session.execute(select(Task.id).where( diff --git a/backend/pyproject.toml b/backend/pyproject.toml index d7af0048..dc027c29 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -20,6 +20,9 @@ dependencies = [ # 所以上限卡到下一个小版本——升级要主动改这里、重新锁、跑一遍测试。 # 精确版本在 requirements.lock 里 "langgraph>=1.2.11,<1.3", + # 断点的序列化格式就在这个包里(langgraph 只要求 <5,上面那条卡不住它);SQLite、PostgreSQL 两种断点库 + # 共用这一份格式,copy-to-postgres 搬断点也靠它前后一致 + "langgraph-checkpoint>=4.2.0,<4.3", "langgraph-checkpoint-sqlite>=3.1.1,<3.2", "langchain-core>=1.6.3,<1.7", "langchain-anthropic>=1.7.2,<1.8", @@ -49,6 +52,15 @@ docs = ["pypdf>=5.0", "python-docx>=1.1", "python-pptx>=1.0", "openpyxl>=3.1"] # 数据库接入。三个驱动都是 asyncio 原生且免装数据库客户端—— # oracledb 走 thin 模式,省掉了 Oracle Instant Client 这个历来最劝退的一步。 db = ["aiomysql>=0.2", "oracledb>=2.0", "asyncpg>=0.30"] +# 应用自己的库用 PostgreSQL(AGENTLAB_DATABASE_URL=postgresql+asyncpg://…)。应用库走 asyncpg;LangGraph 的断点 +# 走官方的 langgraph-checkpoint-postgres,它只认 psycopg 3 和它的连接池。psycopg 带 binary:wheel 里自带 libpq, +# 镜像和本机都不用另装 PostgreSQL 客户端。断点格式归 LangGraph 一族管,上限同上卡到下一个小版本 +postgres = [ + "asyncpg>=0.30", + "psycopg[binary]>=3.2,<4", + "psycopg-pool>=3.2,<4", + "langgraph-checkpoint-postgres>=3.1.2,<3.2", +] [build-system] requires = ["hatchling"] diff --git a/backend/pytest.ini b/backend/pytest.ini index b30d93f3..facc6e1c 100644 --- a/backend/pytest.ini +++ b/backend/pytest.ini @@ -3,3 +3,5 @@ asyncio_mode = auto testpaths = tests addopts = -q --tb=short +markers = + sqlite_only: 只对 SQLite 应用库成立的测试(数据目录里的库文件、PRAGMA);AGENTLAB_TEST_DATABASE_URL 跑 PostgreSQL 时跳过 diff --git a/backend/requirements.lock b/backend/requirements.lock index ee31c969..20b8fb3d 100644 --- a/backend/requirements.lock +++ b/backend/requirements.lock @@ -136,9 +136,13 @@ langgraph==1.2.11 # via agentlab-backend (pyproject.toml) langgraph-checkpoint==4.2.0 # via + # agentlab-backend (pyproject.toml) # langgraph + # langgraph-checkpoint-postgres # langgraph-checkpoint-sqlite # langgraph-prebuilt +langgraph-checkpoint-postgres==3.1.2 + # via agentlab-backend (pyproject.toml) langgraph-checkpoint-sqlite==3.1.1 # via agentlab-backend (pyproject.toml) langgraph-prebuilt==1.1.0 @@ -167,6 +171,7 @@ oracledb==26.0.0 # via agentlab-backend (pyproject.toml) orjson==3.12.0 # via + # langgraph-checkpoint-postgres # langgraph-sdk # langsmith ormsgpack==1.12.2 @@ -180,6 +185,16 @@ pillow==12.3.0 # via python-pptx pluggy==1.6.0 # via pytest +psycopg==3.3.6 + # via + # agentlab-backend (pyproject.toml) + # langgraph-checkpoint-postgres +psycopg-binary==3.3.6 ; implementation_name != 'pypy' + # via psycopg +psycopg-pool==3.3.3 + # via + # agentlab-backend (pyproject.toml) + # langgraph-checkpoint-postgres pycparser==3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy' # via cffi pydantic==2.13.5 @@ -294,6 +309,8 @@ typing-extensions==4.16.0 # mcp # openai # oracledb + # psycopg + # psycopg-pool # pydantic # pydantic-core # pytest-asyncio @@ -309,6 +326,8 @@ typing-inspection==0.4.4 # mcp # pydantic # pydantic-settings +tzdata==2026.5 ; sys_platform == 'win32' + # via psycopg urllib3==2.7.0 # via requests uuid-utils==0.17.1 diff --git a/backend/tests/appdb.py b/backend/tests/appdb.py new file mode 100644 index 00000000..bf217064 --- /dev/null +++ b/backend/tests/appdb.py @@ -0,0 +1,77 @@ +"""测试里另开一条连接直接改应用库:模拟别处(另一个进程、恢复、继续运行)抢在中间提交。 + +应用库是 SQLite 时用 sqlite3 连数据目录里的 agentlab.db(同以前各处的写法);是 PostgreSQL 时(AGENTLAB_TEST_DATABASE_URL, +见 conftest)用 psycopg 连同一个库。SQL 一律用 ? 占位,两种库写成同一句;锁等 5 秒(同以前 sqlite3 的 timeout=5): +另一条连接要等的锁正好被调用方拿着的话,等满报错,不会把测试卡死。 +""" + +from __future__ import annotations + +import re +from typing import Any + +from app.core.config import settings + + +def is_postgres() -> bool: + return settings.db_dialect == "postgresql" + + +def _pg_conninfo() -> str: + from app.db.urls import libpq_conninfo + + return libpq_conninfo(settings.db_url) + + +def execute(*statements: str | tuple[str, tuple[Any, ...]]) -> None: + """在另一条连接上依次执行、一起提交。每一项是 SQL,或者 (SQL, 参数)。""" + items = [(s, ()) if isinstance(s, str) else (s[0], tuple(s[1])) for s in statements] + if is_postgres(): + import psycopg + + with psycopg.connect(_pg_conninfo(), options="-c lock_timeout=5000") as conn: + for sql, params in items: + conn.execute(re.sub(r"\?", "%s", sql), params) + conn.commit() + return + import sqlite3 + + db = sqlite3.connect(settings.db_path, timeout=5) + try: + for sql, params in items: + db.execute(sql, params) + db.commit() + finally: + db.close() + + +def query(sql: str, *params: Any) -> list[tuple[Any, ...]]: + """在另一条连接上读:看别人这时提交了什么(不是调用方事务里的)。""" + if is_postgres(): + import psycopg + + with psycopg.connect(_pg_conninfo(), options="-c lock_timeout=5000") as conn: + return [tuple(r) for r in conn.execute(re.sub(r"\?", "%s", sql), params).fetchall()] + import sqlite3 + + db = sqlite3.connect(settings.db_path, timeout=5) + try: + return [tuple(r) for r in db.execute(sql, params).fetchall()] + finally: + db.close() + + +def refuse_inserts(table: str, condition: str, *, name: str) -> tuple[list[str], list[str]]: + """建一个触发器:往 table 插入满足 condition(用 NEW.列 写)的行时报错,报的是完整性错误(IntegrityError),同 SQLite + 的 RAISE(ABORT)。返回 (建的语句, 删的语句),各是一串语句,逐条在应用库的连接上执行。 + + PostgreSQL 要先建一个触发器函数(plpgsql)再挂上去,错误码给 23000(integrity_constraint_violation)。""" + if is_postgres(): + create = [f"CREATE FUNCTION {name}_fn() RETURNS trigger AS $$ BEGIN IF {condition} THEN " + f"RAISE EXCEPTION 'refused by test' USING ERRCODE = 'integrity_constraint_violation'; END IF; " + f"RETURN NEW; END $$ LANGUAGE plpgsql", + f"CREATE TRIGGER {name} BEFORE INSERT ON {table} FOR EACH ROW EXECUTE FUNCTION {name}_fn()"] + drop = [f"DROP TRIGGER IF EXISTS {name} ON {table}", f"DROP FUNCTION IF EXISTS {name}_fn()"] + return create, drop + return ([f"CREATE TRIGGER {name} BEFORE INSERT ON {table} WHEN {condition} " + f"BEGIN SELECT RAISE(ABORT, 'refused by test'); END"], [f"DROP TRIGGER IF EXISTS {name}"]) diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index c8f7d71d..cd074b0a 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -12,12 +12,58 @@ app.db.base 导入时就按 db_url 建好了。conftest 早于测试模块加载 from __future__ import annotations +import atexit import os import tempfile from pathlib import Path _TMP = Path(tempfile.mkdtemp(prefix="agentlab-tests-")) os.environ["AGENTLAB_DATA_DIR"] = str(_TMP) + + +def _postgres_database(admin_url: str) -> str: + """AGENTLAB_TEST_DATABASE_URL 给的是管理用的库:每个 xdist 进程在同一台服务器上建一个自己的空库(库名加进程号后缀, + 已有的同名库先删掉),进程退出时删掉。返回这个库的地址。scripts/test-postgres.sh 起临时容器、设好这个变量。""" + import psycopg + from sqlalchemy.engine import make_url + + from app.db.urls import libpq_conninfo, normalize_database_url + + url = make_url(normalize_database_url(admin_url)) + name = f"{url.database}_{os.environ.get('PYTEST_XDIST_WORKER', 'main')}" + admin = libpq_conninfo(url.render_as_string(hide_password=False)) + + def run(sql: str) -> None: + with psycopg.connect(admin, autocommit=True) as conn: + conn.execute(sql) + + run(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)') + run(f'CREATE DATABASE "{name}"') + + def drop() -> None: + try: + run(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)') + except Exception: # noqa: BLE001 - 服务器先走了(容器已删)就算了 + pass + + atexit.register(drop) + return url.set(database=name).render_as_string(hide_password=False) + + +# 应用库:缺省数据目录里的 SQLite。设了 AGENTLAB_TEST_DATABASE_URL 就跑在 PostgreSQL 上。不管哪种,都盖掉环境和 .env +# 里可能有的 AGENTLAB_DATABASE_URL / AGENTLAB_CHECKPOINT_URL:测试绝不能连到真在用的库上 +_TEST_PG = os.environ.get("AGENTLAB_TEST_DATABASE_URL", "").strip() +os.environ["AGENTLAB_CHECKPOINT_URL"] = "" +os.environ["AGENTLAB_DATABASE_URL"] = _postgres_database(_TEST_PG) if _TEST_PG else "" +if _TEST_PG: + # 每个测试一个事件循环(pytest-asyncio 的缺省),asyncpg 的连接不能跨循环用:引擎先按不用连接池建, + # 再换成按事件循环分的池子(tests/pgpool.py),一个测试里的查库照样复用连接 + os.environ["AGENTLAB_DB_POOL_SIZE"] = "0" + + from app.db.base import engine as _app_engine # noqa: E402 + from tests.pgpool import install as _loop_local_pool # noqa: E402 + + _loop_local_pool(_app_engine) # 别把真实的 key 带进测试环境 os.environ.setdefault("AGENTLAB_SECRET_KEY", "test-only-not-a-real-key") # 定时任务的调度循环不在测试里起:走 lifespan 的测试不该到点替人发起运行。要测调度的直接调它 @@ -42,7 +88,18 @@ def _prepare_schema(): def pytest_report_header(config) -> str: - return f"测试数据目录:{_TMP}" + where = "PostgreSQL(AGENTLAB_TEST_DATABASE_URL)" if _TEST_PG else "SQLite" + return f"测试数据目录:{_TMP};应用库:{where}" + + +def pytest_collection_modifyitems(config, items) -> None: + """标了 sqlite_only 的测试只在应用库是 SQLite 时跑:它们断言的是 SQLite 自己的东西(数据目录里的库文件、PRAGMA)。""" + if not _TEST_PG: + return + skip = pytest.mark.skip(reason="只对 SQLite 应用库成立") + for item in items: + if "sqlite_only" in item.keywords: + item.add_marker(skip) @pytest.fixture(scope="session") diff --git a/backend/tests/pgpool.py b/backend/tests/pgpool.py new file mode 100644 index 00000000..ae1b4a65 --- /dev/null +++ b/backend/tests/pgpool.py @@ -0,0 +1,123 @@ +"""应用库是 PostgreSQL 时测试用的连接池:每个事件循环一个池子,连接只给建它的那个循环用。 + +pytest-asyncio 每个测试一个事件循环,TestClient 在自己的线程里另起一个;asyncpg 的连接绑在建它的循环上,换个循环用 +就报「attached to a different loop」。不用连接池(每次查库新建连接)能跑,但一次运行几十次查库,每次多几毫秒, +合起来让计时类的测试(审批收尾不计时长、反查的时限)在满载时越界——生产上有连接池,没有这笔开销。 + +这里:连接按建它的循环归池子,取的时候只从当前循环的池子里拿,还回来放回它自己的池子;循环关掉之前(测试结束、 +TestClient 退出)把它池子里的连接在它自己的循环里关掉。只在 conftest 里、AGENTLAB_TEST_DATABASE_URL 时装上, +生产代码不用它。 +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import weakref +from typing import Any + +from sqlalchemy.pool import Pool + +_LOOP = "agentlab_loop" + + +class LoopLocalPool(Pool): + _is_asyncio = True + + def __init__(self, creator: Any, **kw: Any) -> None: + super().__init__(creator, **kw) + self._idle: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, list[Any]] = weakref.WeakKeyDictionary() + self._lock = threading.Lock() + + @classmethod + def replacing(cls, pool: Pool) -> LoopLocalPool: + """换掉引擎原来的池子:同一个建连接的函数、同一套连接事件(方言的 on_connect 挂在上面)。""" + return cls(pool._creator, recycle=pool._recycle, echo=pool.echo, logging_name=pool._orig_logging_name, + reset_on_return=pool._reset_on_return, pre_ping=False, _dispatch=pool.dispatch, + dialect=pool._dialect) + + def recreate(self) -> LoopLocalPool: + return LoopLocalPool.replacing(self) + + def _bucket(self, loop: asyncio.AbstractEventLoop) -> list[Any]: + with self._lock: + return self._idle.setdefault(loop, []) + + def _do_get(self) -> Any: + bucket = self._bucket(asyncio.get_running_loop()) + with self._lock: + if bucket: + return bucket.pop() + record = self._create_connection() + record.info[_LOOP] = weakref.ref(asyncio.get_running_loop()) + return record + + def _do_return_conn(self, record: Any) -> None: + ref = record.info.get(_LOOP) + loop = ref() if ref is not None else None + if loop is None or loop.is_closed(): + with contextlib.suppress(Exception): + record.close() + return + bucket = self._bucket(loop) + with self._lock: + bucket.append(record) + + def take(self, loop: asyncio.AbstractEventLoop) -> list[Any]: + """这个循环池子里空闲的连接,取出来(交给调用方在这个循环里关掉)。""" + with self._lock: + return self._idle.pop(loop, []) + + def dispose(self) -> None: + """只关当前循环的:别的循环的连接不能在这里关。没有正在跑的循环时什么都不做(各循环关掉时自己收拾)。""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + for record in self.take(loop): + with contextlib.suppress(Exception): + record.close() + + def status(self) -> str: + with self._lock: + idle = sum(len(v) for v in self._idle.values()) + return f"LoopLocalPool id:{id(self)} idle:{idle} loops:{len(self._idle)}" + + +def install(engine: Any) -> LoopLocalPool: + """把应用库引擎的池子换成 LoopLocalPool,并让每个事件循环关掉之前先关它自己的连接。""" + from sqlalchemy.util import greenlet_spawn + + pool = LoopLocalPool.replacing(engine.sync_engine.pool) + engine.sync_engine.pool = pool + + async def close_mine(records: list[Any]) -> None: + def run() -> None: + for record in records: + with contextlib.suppress(Exception): + record.close() + + await greenlet_spawn(run) + + base = asyncio.get_event_loop_policy() + + class Policy(type(base)): # type: ignore[misc] + def new_event_loop(self) -> asyncio.AbstractEventLoop: + loop = super().new_event_loop() + close = loop.close + + def closing() -> None: + current = engine.sync_engine.pool + if isinstance(current, LoopLocalPool) and not loop.is_closed() and not loop.is_running(): + records = current.take(loop) + if records: + with contextlib.suppress(Exception): + loop.run_until_complete(close_mine(records)) + close() + + loop.close = closing # type: ignore[method-assign] + return loop + + asyncio.set_event_loop_policy(Policy()) + return pool diff --git a/backend/tests/test_datasources_upload.py b/backend/tests/test_datasources_upload.py index 51ca5fb7..30060844 100644 --- a/backend/tests/test_datasources_upload.py +++ b/backend/tests/test_datasources_upload.py @@ -556,9 +556,16 @@ async def test_startup_locks_legacy_sources_registered_under_other_spellings(cli def _app_db() -> str: - """测试进程实际连着的元数据库(conftest 建的那个)。data_dir 被夹具换过,和 settings.db_path 不是同一个。""" + """测试进程实际连着的元数据库(conftest 建的那个)。data_dir 被夹具换过,和 settings.db_path 不是同一个。 + + 应用库是 PostgreSQL 时没有连着的库文件:拿数据目录里的 agentlab.db 顶上(导入 PostgreSQL 之前留下的那一份, + 同一批元数据,照样不许登记),先建出这个文件,硬链接、大小写这些变体才有东西可指。""" from app.db import base as db_base + if db_base.engine.dialect.name != "sqlite": + settings.db_path.parent.mkdir(parents=True, exist_ok=True) + settings.db_path.touch() + return str(settings.db_path) return str(db_base.engine.url.database) diff --git a/backend/tests/test_db_portable.py b/backend/tests/test_db_portable.py new file mode 100644 index 00000000..7dcbfc11 --- /dev/null +++ b/backend/tests/test_db_portable.py @@ -0,0 +1,357 @@ +"""应用库在 SQLite、PostgreSQL 上行为一致(db/portable.py、db/postgres.py、db/base.py 的迁移)。 + +连着库的那几项在两种库上各跑一遍(scripts/test-postgres.sh 跑 PostgreSQL),断言同一套结果;只看编译结果的不连库; +PostgreSQL 老库迁移另建一个库(只在 AGENTLAB_TEST_DATABASE_URL 时跑)。全是虚构数据。 +""" + +from __future__ import annotations + +import os +import re +import uuid + +import pytest +from sqlalchemy import func, select, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.schema import CreateTable + +from app.db import base as db_base +from app.db.base import Base, SessionLocal +from app.db.models import Run, TableRecipe +from app.db.portable import json_eq, json_flag, json_is_string, json_text, unique_violation + +PG = bool(os.environ.get("AGENTLAB_TEST_DATABASE_URL", "").strip()) + + +def _pg_dialect(): + from sqlalchemy.dialects.postgresql.asyncpg import dialect + + from app.db import postgres + + made = dialect() + postgres.install(made) + return made + + +# -------------------------------------------------------------------------- +# 守着的约定 +# -------------------------------------------------------------------------- + + +def test_every_migrated_column_is_in_the_models(): + """PostgreSQL 上补列的类型取模型里这一列的(db/base._column_ddl):迁移表里的每一列模型里都得有。""" + missing = [(t, c) for t, c, _ in db_base._COLUMN_MIGRATIONS + if t not in Base.metadata.tables or c not in Base.metadata.tables[t].c] + assert missing == [] + + +def test_checkpoint_tables_do_not_clash_with_app_tables(): + """断点和应用表放在同一个库、同一个 schema:两边的表名不能撞。名字从两个 saver 自己的建表语句里取。""" + import inspect as pyinspect + + from langgraph.checkpoint.postgres.base import BasePostgresSaver + from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + + created = set(re.findall(r"CREATE TABLE IF NOT EXISTS (\w+)", " ".join(BasePostgresSaver.MIGRATIONS))) + created |= set(re.findall(r"CREATE TABLE IF NOT EXISTS (\w+)", pyinspect.getsource(AsyncSqliteSaver.setup))) + assert {"checkpoints", "checkpoint_blobs", "checkpoint_writes", "checkpoint_migrations"} <= created + assert created & set(Base.metadata.tables) == set() + + +def test_postgres_columns_are_unbounded_varchar_bigint_and_json(): + """应用库上 String(n) 建成不限长的 VARCHAR、Integer 是 BIGINT(自增主键 BIGSERIAL)、JSON 是 json(不是 JSONB: + JSONB 会重排对象里的键)、时间是不带时区的 TIMESTAMP。SQLite 上的建表语句不变。""" + from sqlalchemy.dialects.sqlite.aiosqlite import dialect as sqlite_dialect + + runs = Base.metadata.tables["runs"] + pg = CreateTable(runs).compile(dialect=_pg_dialect()).string + assert "id VARCHAR NOT NULL" in pg and "VARCHAR(" not in pg + assert "last_seq BIGINT" in pg and "INTEGER" not in pg + assert "input JSON NOT NULL" in pg and "JSONB" not in pg + assert "created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL" in pg + changes = CreateTable(Base.metadata.tables["source_changes"]).compile(dialect=_pg_dialect()).string + assert "seq BIGSERIAL NOT NULL" in changes + lite = CreateTable(runs).compile(dialect=sqlite_dialect()).string + assert "id VARCHAR(32) NOT NULL" in lite and "last_seq INTEGER" in lite and "input JSON NOT NULL" in lite + + +def test_postgres_casts_keep_their_length(): + """CAST(x AS VARCHAR(n)) 是截断的意思,照旧写长度;只有建表时的列不限长。""" + from sqlalchemy import String, cast + + sql = str(select(cast(Run.workflow_name, String(4))).compile(dialect=_pg_dialect())) + assert "CAST(runs.workflow_name AS VARCHAR(4))" in sql + + +def test_postgres_orders_nulls_like_sqlite(): + """SQLite 升序 NULL 在前、降序 NULL 在后:PostgreSQL 上可空列照这个补写,不可空的列、写明了的不动。""" + from app.db.models import Task + + sql = str(select(Run.id, func.row_number().over(order_by=Run.started_at).label("n")) + .order_by(Run.started_at, Run.finished_at.desc(), Run.created_at.desc(), + Task.due_at.desc().nulls_first(), Task.remind_at.asc().nulls_last()) + .compile(dialect=_pg_dialect())) + assert "OVER (ORDER BY runs.started_at ASC NULLS FIRST)" in sql + order = sql.split("ORDER BY", 2)[2] + assert "runs.started_at ASC NULLS FIRST" in order + assert "runs.finished_at DESC NULLS LAST" in order + assert "runs.created_at DESC," in order # 不可空:不补 + assert "tasks.due_at DESC NULLS FIRST" in order and "NULLS FIRST NULLS" not in order + assert "tasks.remind_at ASC NULLS LAST" in order + + +def test_postgres_integer_binds_are_bigint(): + """和 Integer 列比较、写进 Integer 列的参数按 BIGINT 转换,超过 21 亿的值不在半路溢出。""" + sql = str(select(Run.id).where(Run.last_seq > 5).compile(dialect=_pg_dialect())) + assert "runs.last_seq > $1::BIGINT" in sql + + +# -------------------------------------------------------------------------- +# 连着库:两种库同一套结果 +# -------------------------------------------------------------------------- + + +async def _runs(origins: list[dict | None]) -> tuple[str, list[str]]: + tag = uuid.uuid4().hex[:12] + ids = [] + async with SessionLocal() as session: + for i, origin in enumerate(origins): + row = Run(workflow_name=f"便携-{tag}-{i}", origin=origin, input={"question": "虚构的问题"} if i % 2 else + {"question": 7}) + session.add(row) + await session.flush() + ids.append(row.id) + await session.commit() + return tag, ids + + +async def _match(tag: str, *where) -> list[int]: + async with SessionLocal() as session: + names = (await session.execute(select(Run.workflow_name).where( + Run.workflow_name.like(f"便携-{tag}-%"), *where))).scalars().all() + return sorted(int(n.rsplit("-", 1)[1]) for n in names) + + +async def test_json_helpers_match_the_same_rows_on_both_databases(): + tag, _ = await _runs([ + {"v": 3, "s": "3", "flag": True, "k": "甲"}, # 0 + {"v": "3", "s": 3, "flag": 1, "k": None}, # 1 + {"v": 4, "flag": 0, "k": "乙", "uses": ["r-1", "r-22"]}, # 2 + {"flag": False}, # 3 + {}, # 4 + None, # 5(JSON 的 null) + ]) + # 带类型的相等:数字 3 不等于字符串 "3" + assert await _match(tag, json_eq(Run.origin, "v", 3)) == [0] + assert await _match(tag, json_eq(Run.origin, "v", "3")) == [1] + assert await _match(tag, json_eq(Run.origin, "s", "3")) == [0] + # None:没有这个键、或者值是 null + assert await _match(tag, json_eq(Run.origin, "k", None)) == [1, 3, 4, 5] + # 标记位:true 和非零的数算真 + assert await _match(tag, json_flag(Run.origin, "flag")) == [0, 1] + from sqlalchemy import not_ + + assert await _match(tag, not_(json_flag(Run.origin, "flag"))) == [2, 3, 4, 5] + # 文本 + assert await _match(tag, json_text(Run.origin, "k") == "乙") == [2] + assert await _match(tag, func.coalesce(json_text(Run.origin, "k"), "").in_(["甲", ""])) == [0, 1, 3, 4, 5] + assert await _match(tag, json_text(Run.origin, "uses").contains("r-22", autoescape=True)) == [2] + # 是不是字符串:input.question 偶数行是数字、奇数行是字符串 + assert await _match(tag, json_is_string(Run.input, "question")) == [1, 3, 5] + + +async def test_json_columns_keep_the_order_of_keys(): + """JSON 列存原文:对象里键的先后读回来不变(JSONB 会重排,应用里按键的先后出列名、出段落的地方就乱了)。""" + keys = ["单号", "金额", "电话", "b", "a", "aa"] + async with SessionLocal() as session: + row = Run(workflow_name="键序", origin={k: i for i, k in enumerate(keys)}) + session.add(row) + await session.commit() + rid = row.id + async with SessionLocal() as session: + got = (await session.execute(select(Run.origin).where(Run.id == rid))).scalar_one() + assert list(got) == keys + + +async def test_json_null_and_sql_null_stay_apart(): + """JSON 列写 None 存的是 JSON 的 null(不是 SQL 的 NULL),两种库一样:按 IS NULL 筛不到它。""" + async with SessionLocal() as session: + row = Run(workflow_name="空值", origin=None) + session.add(row) + await session.commit() + rid = row.id + async with SessionLocal() as session: + is_sql_null = (await session.execute(select(Run.origin.is_(None)).where(Run.id == rid))).scalar_one() + raw = (await session.execute(text("SELECT CAST(origin AS TEXT) FROM runs WHERE id = :i"), {"i": rid})).scalar_one() + assert is_sql_null is False and raw == "null" + + +async def test_long_strings_and_big_integers_are_kept(): + """String(n) 超长、Integer 过 21 亿:SQLite 照存,PostgreSQL 上也照存(列不限长、是 BIGINT)。""" + long_name = "长" * 500 # workflow_name 是 String(200) + big = 5_000_000_000 + async with SessionLocal() as session: + row = Run(workflow_name=long_name, last_seq=big) + session.add(row) + await session.commit() + rid = row.id + async with SessionLocal() as session: + got = (await session.execute(select(Run.workflow_name, Run.last_seq).where( + Run.id == rid, Run.last_seq > big - 1))).one() + assert tuple(got) == (long_name, big) + + +async def test_non_finite_floats_in_json_are_written(): + """JSON 里的 NaN:SQLite 照存;PostgreSQL 的 json 不收,写成 null(同浏览器的 JSON.stringify),不让整条写入失败。""" + async with SessionLocal() as session: + row = Run(workflow_name="非有限", origin={"x": float("nan"), "y": 1.5, "s": "a\x00b"}) + session.add(row) + await session.commit() + rid = row.id + async with SessionLocal() as session: + got = (await session.execute(select(Run.origin).where(Run.id == rid))).scalar_one() + assert got["y"] == 1.5 + if db_base.engine.dialect.name == "postgresql": + assert got["x"] is None and got["s"] == "ab" + else: + assert got["x"] != got["x"] and got["s"] == "a\x00b" # NaN、NUL 原样 + + +async def test_unique_violation_names_the_table_and_columns(): + sid = uuid.uuid4().hex + async with SessionLocal() as session: + session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256="c" * 64)) + await session.commit() + session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256="c" * 64)) + with pytest.raises(IntegrityError) as caught: + await session.commit() + assert unique_violation(caught.value, "table_recipes", "source_id", "seq") + assert not unique_violation(caught.value, "table_recipes", "source_id") + assert not unique_violation(caught.value, "data_sources", "source_id", "seq") + + +async def test_order_by_puts_nulls_where_sqlite_does(): + tag = uuid.uuid4().hex[:12] + from datetime import datetime, timezone + + async with SessionLocal() as session: + for i, started in enumerate([datetime(2026, 1, 2, tzinfo=timezone.utc), None, + datetime(2026, 1, 1, tzinfo=timezone.utc)]): + session.add(Run(workflow_name=f"排序-{tag}-{i}", started_at=started)) + await session.commit() + async with SessionLocal() as session: + where = Run.workflow_name.like(f"排序-{tag}-%") + up = (await session.execute(select(Run.workflow_name).where(where).order_by(Run.started_at))).scalars().all() + down = (await session.execute(select(Run.workflow_name).where(where) + .order_by(Run.started_at.desc()))).scalars().all() + assert [n[-1] for n in up] == ["1", "2", "0"] # 升序:NULL 在前 + assert [n[-1] for n in down] == ["0", "2", "1"] # 降序:NULL 在后 + + +async def test_timestamps_read_back_as_aware_utc(): + from datetime import datetime, timedelta, timezone + + local = datetime(2026, 3, 4, 9, 30, tzinfo=timezone(timedelta(hours=8))) + async with SessionLocal() as session: + row = Run(workflow_name="时区", started_at=local) + session.add(row) + await session.commit() + rid = row.id + async with SessionLocal() as session: + got = (await session.execute(select(Run.started_at).where(Run.id == rid, Run.started_at == local))).scalar_one() + assert got == local and got.tzinfo == timezone.utc and got.hour == 1 + + +# -------------------------------------------------------------------------- +# PostgreSQL 上的老库迁移 +# -------------------------------------------------------------------------- + + +@pytest.fixture +def scratch_pg(): + """同一台服务器上另建一个空库,给迁移测试造「老版本建的库」。只在 PostgreSQL 时有。""" + if not PG: + pytest.skip("要 AGENTLAB_TEST_DATABASE_URL(scripts/test-postgres.sh)") + import psycopg + from sqlalchemy import create_engine + from sqlalchemy.engine import make_url + + from app.db.urls import libpq_conninfo + + admin = make_url(db_base.settings.db_url) + name = f"{admin.database}_mig_{uuid.uuid4().hex[:8]}" + conninfo = libpq_conninfo(admin.render_as_string(hide_password=False)) + with psycopg.connect(conninfo, autocommit=True) as conn: + conn.execute(f'CREATE DATABASE "{name}"') + url = admin.set(drivername="postgresql+psycopg", database=name) + engine = create_engine(url) + try: + yield engine + finally: + engine.dispose() + with psycopg.connect(conninfo, autocommit=True) as conn: + conn.execute(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)') + + +def _columns(conn, table: str) -> dict[str, tuple[str, str | None]]: + rows = conn.execute(text("SELECT column_name, data_type, column_default FROM information_schema.columns " + "WHERE table_name = :t"), {"t": table}).all() + return {r[0]: (r[1], r[2]) for r in rows} + + +def test_old_postgres_databases_get_the_missing_columns_and_indexes(scratch_pg): + """老版本建的 PostgreSQL 库(少几列、少索引、有并发留下的重复行):补上的列类型按模型的、缺省值回填,布尔的缺省是 + false 不是 0,索引补上,唯一索引建之前先去重;再跑一遍什么都不做。""" + with scratch_pg.begin() as conn: + conn.execute(text("CREATE TABLE runs (id VARCHAR PRIMARY KEY, workflow_id VARCHAR, version BIGINT, " + "status VARCHAR, last_seq BIGINT)")) + conn.execute(text("INSERT INTO runs (id, status, last_seq) VALUES ('r1', 'succeeded', 9)")) + conn.execute(text("CREATE TABLE schedules (id VARCHAR PRIMARY KEY, name VARCHAR)")) + conn.execute(text("INSERT INTO schedules (id, name) VALUES ('s1', '虚构的定时')")) + conn.execute(text("CREATE TABLE conversations (id VARCHAR PRIMARY KEY, kind VARCHAR, archived BOOLEAN, " + "workflow_id VARCHAR, last_active_at TIMESTAMP)")) + for cid, at in (("c1", "2026-01-01"), ("c2", "2026-01-02")): + conn.execute(text("INSERT INTO conversations (id, kind, archived, last_active_at) " + "VALUES (:i, 'app', false, :at)"), {"i": cid, "at": at}) + db_base._migrate(conn) + db_base._migrate(conn) + runs = _columns(conn, "runs") + assert runs["options"][0] == "json" and runs["options"][1].startswith("'{}'") + assert runs["schedule_id"][0] == "character varying" and runs["manifest_seq"][0] == "bigint" + assert runs["started_by"][0] == "character varying" + assert conn.execute(text("SELECT CAST(options AS TEXT), run_class FROM runs WHERE id = 'r1'")).one() \ + == ("{}", "exploratory") + schedules = _columns(conn, "schedules") + assert schedules["trigger_kind"][1].startswith("'cron'") + assert conn.execute(text("SELECT trigger_kind FROM schedules")).scalar_one() == "cron" + indexes = {r[0] for r in conn.execute(text("SELECT indexname FROM pg_indexes WHERE tablename = 'runs'"))} + assert {"ix_runs_schedule_id", "ix_runs_wf_ver_status"} <= indexes + # 同一个人在同一个应用里两段没归档的助手对话:created_by 是迁移补上的列(NULL 不受约束),唯一索引照样建上 + convs = {r[0] for r in conn.execute(text("SELECT indexname FROM pg_indexes WHERE tablename = 'conversations'"))} + assert "ux_conversations_app_owner" in convs + + +def test_old_postgres_duplicates_are_settled_before_the_unique_index(scratch_pg): + with scratch_pg.begin() as conn: + conn.execute(text("CREATE TABLE conversations (id VARCHAR PRIMARY KEY, kind VARCHAR, archived BOOLEAN, " + "app_id VARCHAR, created_by VARCHAR, last_active_at TIMESTAMP)")) + for cid, at in (("c1", "2026-01-01"), ("c2", "2026-01-03"), ("c3", "2026-01-02")): + conn.execute(text("INSERT INTO conversations VALUES (:i, 'app', false, 'a1', '甲', :at)"), + {"i": cid, "at": at}) + db_base._migrate(conn) + kept = conn.execute(text("SELECT id FROM conversations WHERE archived = false")).scalars().all() + assert kept == ["c2"], "留最近活跃的那一段,别的归档" + with pytest.raises(Exception, match="ux_conversations_app_owner"): + with conn.begin_nested(): + conn.execute(text("INSERT INTO conversations VALUES ('c4', 'app', false, 'a1', '甲', '2026-02-01')")) + + +async def test_new_postgres_databases_have_the_partial_unique_indexes(): + """新库靠 create_all:部分唯一索引在 PostgreSQL 上也是部分索引(postgresql_where),条件里的布尔写成 false。""" + if not PG: + pytest.skip("要 AGENTLAB_TEST_DATABASE_URL(scripts/test-postgres.sh)") + async with db_base.engine.connect() as conn: + rows = await conn.execute(text("SELECT indexname, indexdef FROM pg_indexes WHERE indexname LIKE 'ux_%'")) + got = {r[0]: r[1] for r in rows} + assert "WHERE (((kind)::text = 'app'::text) AND (archived = false))" in got["ux_conversations_app_owner"] + assert "WHERE ((status)::text = ANY" in got["ux_catalog_draft_jobs_active"] diff --git a/backend/tests/test_evidence_judge_api.py b/backend/tests/test_evidence_judge_api.py index 923aec55..25cea0dd 100644 --- a/backend/tests/test_evidence_judge_api.py +++ b/backend/tests/test_evidence_judge_api.py @@ -446,7 +446,7 @@ async def test_a_continue_that_wins_the_row_between_check_and_write_gets_no_even """核对过状态、还没写入的那一下,接着跑的条件更新(runner._claim)先提交了:追加判定的条件更新落空, 照样不记——核对和写入是同一行上的条件更新,不是先读后写。""" from app.core.bus import bus - from app.core.config import settings + from tests import appdb Models(monkeypatch, verdicts={"新客": "contradicted"}) row = await finish(weekly(claims="judge")) @@ -456,12 +456,7 @@ async def test_a_continue_that_wins_the_row_between_check_and_write_gets_no_even def claimed_in_between(run_id, value): # _append_judged 读完运行、分配序号之前正好调用它:在这里用另一条连接把运行占成 running if run_id == row.id: - db = sqlite3.connect(settings.db_path, timeout=5) - try: - db.execute("UPDATE runs SET status = 'running' WHERE id = ?", (row.id,)) - db.commit() - finally: - db.close() + appdb.execute(("UPDATE runs SET status = 'running' WHERE id = ?", (row.id,))) set_seq(run_id, value) monkeypatch.setattr(bus, "set_seq", claimed_in_between) diff --git a/backend/tests/test_run_append.py b/backend/tests/test_run_append.py index e73d5f8d..2032c090 100644 --- a/backend/tests/test_run_append.py +++ b/backend/tests/test_run_append.py @@ -16,7 +16,6 @@ from __future__ import annotations import asyncio import json -import sqlite3 from collections import Counter import pytest @@ -25,7 +24,6 @@ from pydantic import BaseModel from sqlalchemy import select from app.core.bus import bus -from app.core.config import settings from app.core.events import EventType from app.db.base import SessionLocal from app.db.models import Approval, Run, RunEvent @@ -34,6 +32,7 @@ from app.engine.context import NodeError from app.engine.runner import RunManager, run_manager, verify_manifest from app.engine.schema import NodeType from app.tools.registry import register +from tests import appdb #: Agent 场景里工具真正执行的次数。注册表是进程级的,工具只注册一次,计数按测试清零 CALLS: list[str] = [] @@ -162,12 +161,7 @@ async def _pending(run_id: str) -> list[Approval]: def _sql(statement: str, *params) -> None: """另开一条连接改库:模拟别处(恢复、继续运行、另一个写入者)抢在中间提交。""" - db = sqlite3.connect(settings.db_path, timeout=5) - try: - db.execute(statement, params) - db.commit() - finally: - db.close() + appdb.execute((statement, params)) # -------------------------------------------------------------------------- @@ -226,12 +220,7 @@ async def test_a_resume_that_wins_the_row_leaves_nothing_behind(monkeypatch): def resumed_in_between(run_id, value): if run_id == paused.id: - db = sqlite3.connect(settings.db_path, timeout=5) - try: - db.execute("UPDATE runs SET status = 'running' WHERE id = ?", (paused.id,)) - db.commit() - finally: - db.close() + appdb.execute(("UPDATE runs SET status = 'running' WHERE id = ?", (paused.id,))) set_seq(run_id, value) monkeypatch.setattr(bus, "set_seq", resumed_in_between) @@ -252,14 +241,9 @@ async def test_events_written_elsewhere_in_between_win(monkeypatch): def written_in_between(run_id, value): if run_id == paused.id: - db = sqlite3.connect(settings.db_path, timeout=5) - try: - db.execute("INSERT INTO run_events (id, run_id, seq, type, node_id, ts, data) " - "VALUES ('elsewhere-1', ?, ?, 'log', NULL, 0, '{}')", (paused.id, taken)) - db.execute("UPDATE runs SET last_seq = ? WHERE id = ?", (taken, paused.id)) - db.commit() - finally: - db.close() + appdb.execute(("INSERT INTO run_events (id, run_id, seq, type, node_id, ts, data) " + "VALUES ('elsewhere-1', ?, ?, 'log', NULL, 0, '{}')", (paused.id, taken)), + ("UPDATE runs SET last_seq = ? WHERE id = ?", (taken, paused.id))) set_seq(run_id, value) monkeypatch.setattr(bus, "set_seq", written_in_between) @@ -352,12 +336,8 @@ async def test_published_only_once_stored_and_never_when_not_written(monkeypatch async def stored_when_published(event): if event.type == EventType.APPROVAL_INQUIRY: - db = sqlite3.connect(settings.db_path, timeout=5) - try: - (count,) = db.execute("SELECT count(*) FROM run_events WHERE run_id = ? AND seq = ?", - (event.run_id, event.seq)).fetchone() - finally: - db.close() + [(count,)] = appdb.query("SELECT count(*) FROM run_events WHERE run_id = ? AND seq = ?", + event.run_id, event.seq) published.append((event.seq, count)) await publish(event) @@ -430,12 +410,7 @@ async def test_a_sealed_run_that_changes_in_between_gets_nothing(monkeypatch, ch def changed_in_between(run_id, value): if run_id == final.id: - db = sqlite3.connect(settings.db_path, timeout=5) - try: - db.execute(change, (final.id,)) - db.commit() - finally: - db.close() + appdb.execute((change, (final.id,))) set_seq(run_id, value) monkeypatch.setattr(bus, "set_seq", changed_in_between) diff --git a/backend/tests/test_run_launch.py b/backend/tests/test_run_launch.py index baf0935c..5331592b 100644 --- a/backend/tests/test_run_launch.py +++ b/backend/tests/test_run_launch.py @@ -322,15 +322,16 @@ async def test_launch_events_are_written_with_the_run_or_not_at_all(): from app.db.base import engine from app.engine.launch import launch_run + from tests.appdb import refuse_inserts lib = await upstream_with_two_versions() wf, _ = await workflow(weekly(lib, upgrade_policy="dual"), published=1, versions=(1,)) before = await _run_count() # 口径升版处置那条写不进去。走异步引擎建删触发器:同步连接等写锁时会卡住事件循环,而持锁的正是这个循环里的写入 + create, drop = refuse_inserts("run_events", "NEW.type = 'caliber.upgrade'", name="refuse_upgrade_note") async with engine.begin() as conn: - await conn.execute(text( - "CREATE TRIGGER refuse_upgrade_note BEFORE INSERT ON run_events WHEN NEW.type = 'caliber.upgrade' " - "BEGIN SELECT RAISE(ABORT, 'refused by test'); END")) + for sql in create: + await conn.execute(text(sql)) try: async with SessionLocal() as session: with pytest.raises(IntegrityError): @@ -338,7 +339,8 @@ async def test_launch_events_are_written_with_the_run_or_not_at_all(): assert await _run_count() == before, "没写进治理事件,却留下了一条运行" finally: async with engine.begin() as conn: - await conn.execute(text("DROP TRIGGER IF EXISTS refuse_upgrade_note")) + for sql in drop: + await conn.execute(text(sql)) async def test_an_exploratory_launch_takes_the_saved_graph_and_the_run_defaults(): diff --git a/backend/tests/test_run_options.py b/backend/tests/test_run_options.py index ef456a27..f4e7fa53 100644 --- a/backend/tests/test_run_options.py +++ b/backend/tests/test_run_options.py @@ -112,9 +112,11 @@ def test_a_missing_table_is_left_alone(): async def test_new_databases_index_schedule_id(): + from sqlalchemy import inspect + async with app_engine.connect() as conn: - names = {row[1] for row in (await conn.execute(text("PRAGMA index_list(runs)"))).all()} - assert "ix_runs_schedule_id" in names + indexes = await conn.run_sync(lambda sync: inspect(sync).get_indexes("runs")) + assert {ix["name"]: ix["column_names"] for ix in indexes}.get("ix_runs_schedule_id") == ["schedule_id"] async def test_a_migrated_row_reads_back_empty_options(): diff --git a/backend/tests/test_run_races.py b/backend/tests/test_run_races.py index 6e89d179..a7d10951 100644 --- a/backend/tests/test_run_races.py +++ b/backend/tests/test_run_races.py @@ -254,11 +254,10 @@ async def test_abandoning_while_the_pause_is_being_recorded(slow, monkeypatch): async def test_an_approval_closed_just_before_it_is_marked_answered(slow, monkeypatch): """放弃恰好提交在「读到这条审批还是 pending」和「把它记成已回复」之间。""" - import sqlite3 - from sqlalchemy import event from app.db.base import engine + from tests import appdb run = await run_manager.start(graph=_graph(gate=True), input_payload={"q": "x"}) await _until(lambda: _is(run.id, "interrupted")) @@ -268,11 +267,8 @@ async def test_an_approval_closed_just_before_it_is_marked_answered(slow, monkey def abandon_first(conn, cursor, statement, params, context, executemany): if armed["on"] and statement.lstrip().upper().startswith("UPDATE APPROVALS"): armed["on"] = False - other = sqlite3.connect(engine.url.database, timeout=5) - other.execute("UPDATE approvals SET status='cancelled' WHERE run_id=?", (run.id,)) - other.execute("UPDATE runs SET status='cancelled' WHERE id=?", (run.id,)) - other.commit() - other.close() + appdb.execute(("UPDATE approvals SET status='cancelled' WHERE run_id=?", (run.id,)), + ("UPDATE runs SET status='cancelled' WHERE id=?", (run.id,))) event.listen(engine.sync_engine, "before_cursor_execute", abandon_first) try: diff --git a/backend/tests/test_table_versions_p2.py b/backend/tests/test_table_versions_p2.py index bacd95e6..f90bcae8 100644 --- a/backend/tests/test_table_versions_p2.py +++ b/backend/tests/test_table_versions_p2.py @@ -23,6 +23,7 @@ from typing import Any import pytest from sqlalchemy import create_engine, select, text +from sqlalchemy.exc import IntegrityError from app.core.config import settings from app.data import raw_store, recipe_parsers, recipe_types, table_versions @@ -32,6 +33,7 @@ from app.db.base import SessionLocal, utcnow from app.db.models import ( DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe, ) +from app.db.portable import unique_violation @pytest.fixture(autouse=True) @@ -268,8 +270,10 @@ async def test_recipe_seq_is_unique_per_source(): session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256=RECIPE_SHA)) await session.commit() session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256=RECIPE_SHA)) - with pytest.raises(Exception, match="UNIQUE"): + with pytest.raises(IntegrityError) as caught: await session.commit() + # 撞的正是 (source_id, seq) 这一条唯一约束(两种库的报错写法不同,按表和列认) + assert unique_violation(caught.value, "table_recipes", "source_id", "seq") # -------------------------------------------------------------------------- diff --git a/scripts/test-postgres.sh b/scripts/test-postgres.sh new file mode 100755 index 00000000..e337afdb --- /dev/null +++ b/scripts/test-postgres.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# 后端测试跑在 PostgreSQL 上(应用库和断点库都是 PostgreSQL):本机 docker 起一个临时的 postgres:16 容器,只绑 +# 127.0.0.1 的随机端口、数据放内存里,跑完(出错、Ctrl+C 也一样)删掉容器,留镜像。 +# +# scripts/test-postgres.sh 全量,-n auto +# scripts/test-postgres.sh tests/test_db_portable.py -x 参数原样交给 pytest(给了参数就不再加 -n auto) +# PYTHON=/path/to/python scripts/test-postgres.sh 指定装着后端依赖(含 [postgres]、[dev])的 Python,缺省 python +# AGENTLAB_PG_IMAGE=postgres:16 … 换镜像(要本机已经有,这里不拉) +# +# 测试怎么用这个库见 backend/tests/conftest.py:AGENTLAB_TEST_DATABASE_URL 给的是管理用的库,每个 xdist 进程在上面 +# 建自己的库(库名加 _gw0 之类的后缀),测试结束删掉。只对 SQLite 成立的测试标了 sqlite_only,这里跳过。 +# 不碰 scripts/dialect-it.sh 的容器(那边叫 agentlab-dialect-it-*)。 +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +PY="${PYTHON:-python}" +IMAGE="${AGENTLAB_PG_IMAGE:-postgres:16}" +NAME="agentlab-pgtest-$$" +USER_NAME="agentlab" +PASS="pgtest-only" +DB="agentlab_test" + +command -v docker >/dev/null || { echo "缺少 docker"; exit 1; } +docker image inspect "$IMAGE" >/dev/null 2>&1 || { echo "本机没有镜像 ${IMAGE}(这个脚本不拉镜像):先 docker pull ${IMAGE}"; exit 1; } + +cleanup() { docker rm -f "$NAME" >/dev/null 2>&1 || true; } +trap cleanup EXIT INT TERM + +# fsync、同步提交都关掉:测试库用完就扔,不需要落盘的保证。连接数放宽:每个 xdist 进程不用连接池(每个测试一个事件循环), +# 但同一时刻会开好几条,断点库还各有一个小池子。认证用 trust(只绑 127.0.0.1、用完即删的库):不用连接池时每次查库都 +# 新建连接,SCRAM 握手一次十几毫秒,trust 三毫秒,一次运行几十次查库,差出来的够让计时类的测试越界 +docker run -d --name "$NAME" -e POSTGRES_USER="$USER_NAME" -e POSTGRES_PASSWORD="$PASS" -e POSTGRES_DB="$DB" \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + -p 127.0.0.1::5432 --tmpfs /var/lib/postgresql/data:rw "$IMAGE" \ + -c fsync=off -c synchronous_commit=off -c full_page_writes=off -c max_connections=1000 -c shared_buffers=256MB \ + >/dev/null + +for _ in $(seq 1 60); do + if docker exec "$NAME" pg_isready -U "$USER_NAME" -d "$DB" -h 127.0.0.1 >/dev/null 2>&1; then break; fi + sleep 0.5 +done +docker exec "$NAME" pg_isready -U "$USER_NAME" -d "$DB" -h 127.0.0.1 >/dev/null || { echo "PostgreSQL 没起来"; docker logs "$NAME" | tail -20; exit 1; } +PORT="$(docker port "$NAME" 5432/tcp | head -1 | sed 's/.*://')" + +export AGENTLAB_TEST_DATABASE_URL="postgresql+asyncpg://$USER_NAME:$PASS@127.0.0.1:$PORT/$DB" +echo "测试库:127.0.0.1:${PORT}/${DB}(容器 ${NAME},结束后删除)" + +cd "$ROOT/backend" +args=("$@") +[ ${#args[@]} -eq 0 ] && args=(-n auto) +# 同 README「开发与验证」:环境里的密钥不带进测试 +env -u ANTHROPIC_API_KEY -u ANTHROPIC_BASE_URL -u ANTHROPIC_AUTH_TOKEN -u OPENAI_API_KEY \ + "$PY" -m pytest "${args[@]}"