应用库可移植到 PostgreSQL:AGENTLAB_DATABASE_URL 设成 postgresql+asyncpg://… 时整个应用库用 PostgreSQL,断点用官方的 langgraph-checkpoint-postgres(同一个库,AGENTLAB_CHECKPOINT_URL 可另给);不设时照旧是数据目录里的 SQLite,行为不变

- 配置:database_url、checkpoint_url(启动时校验写法,只认 PostgreSQL)、连接池四项;启动日志写用的是哪种库(主机、库名、用户,不写密码)
- 引擎(db/base.make_engine):SQLite 照旧 WAL、外键、忙等;PostgreSQL 用连接池、探活、会话时区 UTC,装上方言调整(db/postgres.py):建表时 String(n) 是不限长的 VARCHAR(SQLite 从不查长度,导过来的存量也可能超长)、ORDER BY 里可空列的 NULL 位置照 SQLite 补写(升序在前、降序在后,窗口函数里同样);Integer 列在 PostgreSQL 上是 BIGINT(列挂到表上时换成 with_variant,绑定参数跟着列按 BIGINT 转换);JSON 列两边都存原文(PostgreSQL 的 json,不用 JSONB:JSONB 重排对象里的键,快照的列序、表的先后读回来就变了);JSON 里的 NaN 写成 null、NUL 去掉(PostgreSQL 不收)
- 迁移按方言:已有的列和索引用 inspector 认;PostgreSQL 上补列的类型取模型里这一列的、缺省值取迁移表里的(布尔 0/1 换成 false/true),ADD COLUMN IF NOT EXISTS;去重语句里的布尔按方言写;部分唯一索引补 postgresql_where;init_db 在 PostgreSQL 上拿咨询锁(几个进程同时启动不撞)
- 应用库上的 SQLite 专用写法改成按方言编译的表达式(db/portable.py):json_extract 的十几处(口径召回、出具、对账、个人应用、幂等键、访问控制、数据源变化通知、剖析任务、待办超时)换成 json_eq / json_flag / json_text,json_type 换成 json_is_string(两种库同一套带类型的比较);能力登记的 sqlite insert 换成按会话方言的 insert_ignore;能力使用次数的 json_each 按方言分支;附件「没绑记录」的条件不靠 json 的等号;按报错文字认唯一约束的地方换成 unique_violation(按表和列认,两种库的报错都认得)
- 断点(engine/checkpoints.py):PostgreSQL 时开 psycopg 连接池、在咨询锁里 setup;runner 只认 BaseCheckpointSaver,读断点仍只走 saver 的公开接口
- 依赖:新 extra postgres(asyncpg、psycopg[binary] 3、psycopg-pool、langgraph-checkpoint-postgres 3.1.2);langgraph-checkpoint 卡在 <4.3(断点格式就在这个包里,postgres 包新版要 4.3,不动已装的 4.2.0);锁文件补上新包
- 运行收尾:记 action.unknown、安排核实做完才把运行摘掉(以前先摘掉,等「不在跑了」再读事件的一方在库慢时读不到这一条)
- 测试:AGENTLAB_TEST_DATABASE_URL 时跑在 PostgreSQL 上(每个 xdist 进程建自己的库、按事件循环分的连接池 tests/pgpool.py),不设时盖掉环境里的数据库地址照旧用 SQLite;scripts/test-postgres.sh 起临时的 postgres:16 容器跑全量;sqlite_only 标记;另开连接改库的几个竞态测试、触发器测试改成两种库都能跑(tests/appdb.py);新增 test_db_portable 守着方言调整、JSON 键序、超长字符串、大整数、NULL 排序、时区、PostgreSQL 老库迁移
This commit is contained in:
ThearchyHelios
2026-10-19 09:52:09 +08:00
parent c79f4e8bf7
commit 2e7d6103ea
39 changed files with 1522 additions and 155 deletions
+2 -1
View File
@@ -62,6 +62,7 @@ def hidden_runs(principal: Any) -> Any:
from app.db.models import Run from app.db.models import Run
from app.db.models_apps import App from app.db.models_apps import App
from app.db.models_blueprints import Blueprint from app.db.models_blueprints import Blueprint
from app.db.portable import json_text
if principal.mode == "off": if principal.mode == "off":
return false() return false()
@@ -72,7 +73,7 @@ def hidden_runs(principal: Any) -> Any:
kind = func.coalesce(Run.origin_kind, "") kind = func.coalesce(Run.origin_kind, "")
return or_( return or_(
and_(kind == "personal_app", or_(Run.app_id.is_(None), Run.app_id.not_in(mine))), 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)),
) )
+4 -3
View File
@@ -39,7 +39,7 @@ from dataclasses import dataclass
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from typing import Any, Awaitable, Callable from typing import Any, Awaitable, Callable
from sqlalchemy import delete, func, select from sqlalchemy import delete, select
from starlette.requests import HTTPConnection from starlette.requests import HTTPConnection
from app.api.coded import CodedHTTPException from app.api.coded import CodedHTTPException
@@ -535,12 +535,13 @@ async def started_before(principal: Any, idem_key: str) -> str | None:
"""拥有者按这个幂等键用密钥发起过的实例(来源里记着 idempotency_key):处理到一半服务停了、同一个键重新认领时交回它。""" """拥有者按这个幂等键用密钥发起过的实例(来源里记着 idempotency_key):处理到一半服务停了、同一个键重新认领时交回它。"""
from app.db.base import SessionLocal from app.db.base import SessionLocal
from app.db.models import Run from app.db.models import Run
from app.db.portable import json_eq
async with SessionLocal() as session: async with SessionLocal() as session:
return await session.scalar(select(Run.id).where( return await session.scalar(select(Run.id).where(
Run.origin_kind == "app_entry", Run.started_by == principal.username, Run.origin_kind == "app_entry", Run.started_by == principal.username,
func.json_extract(Run.origin, "$.idempotency_key") == idem_key, json_eq(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, "via", "key")).order_by(Run.created_at.desc()).limit(1))
# ---------------------------------------------------------------- 请求记录(A6) # ---------------------------------------------------------------- 请求记录(A6)
+2 -1
View File
@@ -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: def _runs_of(app_id: str, graph_hash: str) -> Any:
from app.db.models import Run 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, 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: async def latest_refresh(session: AsyncSession, app_id: str, graph_hash: str, *, active_only: bool = False) -> Any:
+5 -4
View File
@@ -127,10 +127,9 @@ async def _discover(session: AsyncSession) -> dict[str, dict[str, Any]]:
async def sync_all(session: AsyncSession) -> int: async def sync_all(session: AsyncSession) -> int:
"""把内置工具、数据源、自定义工具、MCP 同步进 capabilities。返回变动的行数。启动时和各对象增删改时调。""" """把内置工具、数据源、自定义工具、MCP 同步进 capabilities。返回变动的行数。启动时和各对象增删改时调。"""
from sqlalchemy.dialects.sqlite import insert
from app.db.base import utcnow from app.db.base import utcnow
from app.db.models import Capability from app.db.models import Capability
from app.db.portable import insert_ignore
found = await _discover(session) found = await _discover(session)
existing = {row.name: row for row in (await session.execute(select(Capability))).scalars()} 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 row.synced_at = now
changed += 1 changed += 1
if fresh: 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 changed += done.rowcount or 0
if changed or fresh: if changed or fresh:
await session.commit() await session.commit()
+2 -1
View File
@@ -98,6 +98,7 @@ async def recent_questions(session: AsyncSession, principal: Any, *, limit: int)
"""最近问过的问题,从新到旧,最多 limit 条(取法见模块开头)。principal 用来去掉这个人读不了的只给拥有者的运行。""" """最近问过的问题,从新到旧,最多 limit 条(取法见模块开头)。principal 用来去掉这个人读不了的只给拥有者的运行。"""
from app.auth.access import hidden_runs from app.auth.access import hidden_runs
from app.db.models import Conversation, ConversationTurn, Run from app.db.models import Conversation, ConversationTurn, Run
from app.db.portable import json_is_string
from app.engine.origin import counted_runs from app.engine.origin import counted_runs
hidden = hidden_runs(principal) hidden = hidden_runs(principal)
@@ -112,7 +113,7 @@ async def recent_questions(session: AsyncSession, principal: Any, *, limit: int)
runs = (await session.execute( runs = (await session.execute(
select(Run.id, Run.input, Run.workflow_name, Run.created_at, Run.origin_kind) 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), .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) .order_by(Run.created_at.desc()).limit(limit)
)).all() )).all()
out = [Question(text=question.strip(), at=at, run_id=run_id, turn_id=turn_id, workflow=workflow, origin="chat") out = [Question(text=question.strip(), at=at, run_id=run_id, turn_id=turn_id, workflow=workflow, origin="chat")
+47 -2
View File
@@ -47,6 +47,17 @@ class Settings(BaseSettings):
# --- 存储 --- # --- 存储 ---
data_dir: Path = PROJECT_DIR / "data" 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),页面和接口同源、只起一个进程。 # 设了就由服务端一并托管页面(见 app/web.py),页面和接口同源、只起一个进程。
@@ -195,6 +206,23 @@ class Settings(BaseSettings):
return True return True
return value 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") @field_validator("web_dist", mode="before")
@classmethod @classmethod
def _blank_is_unset(cls, value: object) -> object: def _blank_is_unset(cls, value: object) -> object:
@@ -223,13 +251,30 @@ class Settings(BaseSettings):
@property @property
def db_url(self) -> str: 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 @property
def checkpoint_path(self) -> Path: def checkpoint_path(self) -> Path:
"""LangGraph checkpointer 用独立的库文件,避免和业务表抢写锁。""" """LangGraph checkpointer 用独立的库文件,避免和业务表抢写锁。只在 SQLite 时用。"""
return self.data_dir / "checkpoints.db" 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 @property
def workspace_dir(self) -> Path: def workspace_dir(self) -> Path:
"""沙箱与文件工具的可写根目录,一切文件访问都被限制在这里面。""" """沙箱与文件工具的可写根目录,一切文件访问都被限制在这里面。"""
+12 -4
View File
@@ -144,10 +144,18 @@ def normalize_fields(existing: Mapping[str, Any] | None, submitted: Any, *, at:
async def run_counts(session: AsyncSession) -> dict[str, int]: async def run_counts(session: AsyncSession) -> dict[str, int]:
"""每项能力被几次运行用到:{能力名: 运行次数},没用过的不出现。按运行发起时记下的能力治理快照数 """每项能力被几次运行用到:{能力名: 运行次数},没用过的不出现。按运行发起时记下的能力治理快照数
(Run.options["capabilities"],一次运行里一项能力只记一次)。只用 SQLite 的 json_each,一次查询。""" (Run.options["capabilities"],一次运行里一项能力只记一次)。一次查询:两种库都用 json_each(PostgreSQL 的碰到不是
rows = await session.execute(text( 对象的值会报错,先按类型挑出来)。"""
"SELECT j.key, COUNT(*) FROM runs, json_each(runs.options, '$.capabilities') AS j " from app.db.portable import dialect_of
"WHERE json_valid(runs.options) GROUP BY j.key"))
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()} return {str(name): int(count) for name, count in rows.all()}
+3 -2
View File
@@ -48,7 +48,7 @@ import logging
import time import time
from typing import Any 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.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
@@ -60,6 +60,7 @@ from app.data.catalog_profile import (
) )
from app.data.engine import SnapshotTampered from app.data.engine import SnapshotTampered
from app.db.base import SessionLocal, new_id, utcnow 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 from app.db.models import CatalogDraftJob, CatalogProfileJob, CatalogProfileTask, DataSource
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -345,7 +346,7 @@ _SHOWN = (case((CatalogProfileTask.started_at.is_(None), 1), else_=0), *_NEWEST)
#: 只做了取值观察的那一项(observe 范围、maintenance 范围里只到期观察的表)在逐表一行的 summary 里带这个标记:它不是一次 #: 只做了取值观察的那一项(observe 范围、maintenance 范围里只到期观察的表)在逐表一行的 summary 里带这个标记:它不是一次
#: 剖析,不算「上次剖析」——不然表清单的「上次剖析」、给助手的剖析事实(行数、日期范围)会被一次只记了行数的观察顶掉 #: 剖析,不算「上次剖析」——不然表清单的「上次剖析」、给助手的剖析事实(行数、日期范围)会被一次只记了行数的观察顶掉
OBSERVE_ONLY = "observe_only" 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: def _observe_only(checks: list[Any] | None) -> bool:
+2 -1
View File
@@ -90,6 +90,7 @@ from app.data.recipe_types import (
from app.data.xlsx_scan import WorkbookScan from app.data.xlsx_scan import WorkbookScan
from app.db.base import new_id, utcnow from app.db.base import new_id, utcnow
from app.db.models import DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe from app.db.models import DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe
from app.db.portable import unique_violation
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -2747,7 +2748,7 @@ async def commit_staging(session: AsyncSession, staging: ImportStaging, *, trial
raise _publish_refused(e) from e raise _publish_refused(e) from e
except IntegrityError as e: except IntegrityError as e:
await _back_to_drafting(session, staging_id) 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(409, "name_taken", f"名称「{name}」已被占用,请放弃这次导入后换一个名称") from e
raise ImportRefused(500, "publish_failed", "导入失败:写入版本记录时出错,当前版本未受影响") from e raise ImportRefused(500, "publish_failed", "导入失败:写入版本记录时出错,当前版本未受影响") from e
except BaseException: except BaseException:
+2 -1
View File
@@ -37,6 +37,7 @@ from app.data import schema_diff as sd
from app.data.names import name_key from app.data.names import name_key
from app.db.base import utcnow from app.db.base import utcnow
from app.db.models_sources import SourceChange from app.db.models_sources import SourceChange
from app.db.portable import json_eq
logger = logging.getLogger(__name__) 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( n = int((await session.execute(select(func.count()).select_from(SourceChange).where(
SourceChange.source_id == source_id, SourceChange.severity == "action", SourceChange.source_id == source_id, SourceChange.severity == "action",
SourceChange.status == "open", 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: if n:
from app.notify import service from app.notify import service
+5 -2
View File
@@ -540,10 +540,13 @@ _DB_SIDE_FILES = ("", "-wal", "-shm", "-journal")
def _app_db_files() -> list[Path]: def _app_db_files() -> list[Path]:
"""应用自己的库:元数据库、运行检查点库,以及实际连着的元数据库(测试里和 data_dir 下的不是同一个)。""" """应用自己的库:元数据库、运行检查点库,以及实际连着的元数据库(测试里和 data_dir 下的不是同一个)。
应用库是 PostgreSQL 时,数据目录里的 agentlab.db、checkpoints.db 照样算:导入 PostgreSQL 之前的那份还留在那里,
里面是同一批元数据。连着的库这时是库名不是文件,不算进来。"""
dbs = [settings.db_path, settings.checkpoint_path] dbs = [settings.db_path, settings.checkpoint_path]
try: 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 - 取不到就只按配置里的路径判断 except Exception: # noqa: BLE001 - 取不到就只按配置里的路径判断
live = None live = None
if live and live != ":memory:": if live and live != ":memory:":
+194 -37
View File
@@ -1,16 +1,23 @@
from __future__ import annotations from __future__ import annotations
import json
import logging
import math
import re
import uuid import uuid
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Any, AsyncIterator from typing import Any, AsyncIterator
from sqlalchemy import JSON, DateTime, MetaData, event from sqlalchemy import JSON, BigInteger, Column, DateTime, Integer, MetaData, event
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.pool import NullPool
from sqlalchemy.types import TypeDecorator from sqlalchemy.types import TypeDecorator
from app.core.config import settings from app.core.config import settings
logger = logging.getLogger(__name__)
NAMING_CONVENTION = { NAMING_CONVENTION = {
"ix": "ix_%(column_0_label)s", "ix": "ix_%(column_0_label)s",
"uq": "uq_%(table_name)s_%(column_0_name)s", "uq": "uq_%(table_name)s_%(column_0_name)s",
@@ -57,11 +64,31 @@ class UTCDateTime(TypeDecorator):
return value.astimezone(timezone.utc) 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): class Base(DeclarativeBase):
metadata = MetaData(naming_convention=NAMING_CONVENTION) metadata = MetaData(naming_convention=NAMING_CONVENTION)
# 图定义、节点配置、事件负载都是任意形状的 JSON,统一映射到 SQLite 的 JSON 列 # 图定义、节点配置、事件负载都是任意形状的 JSON,统一映射到 JSON 列(见 PortableJSON)
type_annotation_map = {dict[str, Any]: JSON, list[Any]: JSON, datetime: UTCDateTime} 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: class TimestampMixin:
@@ -101,16 +128,10 @@ def _stop_aiosqlite_once() -> None:
aiosqlite.Connection.stop = stop # type: ignore[method-assign] aiosqlite.Connection.stop = stop # type: ignore[method-assign]
# 应用库是 PostgreSQL 时也照样打:上传表格、手工登记的 SQLite 业务库(data/engine)还是走 aiosqlite 的连接池
_stop_aiosqlite_once() _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 - 驱动回调 def _set_sqlite_pragma(dbapi_conn, _record): # pragma: no cover - 驱动回调
"""WAL + 外键约束。WAL 让读写不互相阻塞,事件写入和前端轮询能并行。""" """WAL + 外键约束。WAL 让读写不互相阻塞,事件写入和前端轮询能并行。"""
cur = dbapi_conn.cursor() cur = dbapi_conn.cursor()
@@ -120,6 +141,92 @@ def _set_sqlite_pragma(dbapi_conn, _record): # pragma: no cover - 驱动回调
cur.close() 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) SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
@@ -129,8 +236,10 @@ async def get_session() -> AsyncIterator[AsyncSession]:
yield session yield session
# 轻量迁移:create_all 只建新表,不给已有表加列。SQLite 的 ADD COLUMN # 轻量迁移:create_all 只建新表,不给已有表加列。ADD COLUMN 便宜且带默认值回填,够用;等真需要改列类型再考虑 alembic。
# 便宜且带默认值回填,够用;等真需要改列类型再考虑 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]] = [ _COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [
("runs", "run_class", "TEXT DEFAULT 'exploratory'"), ("runs", "run_class", "TEXT DEFAULT 'exploratory'"),
("runs", "version", "INTEGER"), ("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]] = [ _UNIQUE_INDEX_MIGRATIONS: list[tuple[str, str, tuple[str, ...], str | None, str]] = [
("ux_app_hooks_app_name", "app_hooks", ("app_id", "name"), None, ("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 " "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)"), "(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 = 0", ("ux_conversations_app_owner", "conversations", ("app_id", "created_by"), "kind = 'app' AND archived = {false}",
"UPDATE conversations SET archived = 1 WHERE kind = 'app' AND archived = 0 AND created_by IS NOT NULL " "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 " "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 " "ORDER BY last_active_at DESC, id) AS rn FROM conversations WHERE kind = 'app' AND archived = {false} "
"AND created_by IS NOT NULL) WHERE rn = 1)"), "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 - 同步回调 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: 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: if cols and column not in cols:
conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}")) conn.execute(text(f"ALTER TABLE {table} {add} {column} {_column_ddl(dialect, table, column, ddl)}"))
for name, table, columns in _INDEX_MIGRATIONS: cols.add(column)
wanted = (columns,) if isinstance(columns, str) else columns for name, table, wanted in _INDEX_MIGRATIONS:
cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))} wanted = (wanted,) if isinstance(wanted, str) else wanted
cols = columns.get(table) or set()
if all(c in cols for c in wanted): if all(c in cols for c in wanted):
conn.execute(text(f"CREATE INDEX IF NOT EXISTS {name} ON {table} ({', '.join(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: for name, table, wanted, where, dedupe in _UNIQUE_INDEX_MIGRATIONS:
cols = {row[1] for row in conn.execute(text(f"PRAGMA table_info({table})"))} cols = columns.get(table) or set()
exists = conn.execute(text("SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = :n"), {"n": name}).first() if not all(c in cols for c in wanted):
if exists or not all(c in cols for c in columns):
continue continue
conn.execute(text(dedupe)) if any(ix["name"] == name for ix in inspect(conn).get_indexes(table)):
conn.execute(text(f"CREATE UNIQUE INDEX {name} ON {table} ({', '.join(columns)})" continue
+ (f" WHERE {where}" if where else ""))) 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: async def init_db() -> None:
from app.db import models # noqa: F401 确保模型已注册到 metadata logger.info("应用库:%s", describe_database())
await create_schema(engine)
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
await conn.run_sync(_migrate)
+6 -3
View File
@@ -362,7 +362,8 @@ class CatalogDraftJob(Base, TimestampMixin):
__table_args__ = ( __table_args__ = (
Index("ix_catalog_draft_jobs_source_created", "source_id", "created_at"), Index("ix_catalog_draft_jobs_source_created", "source_id", "created_at"),
Index("ux_catalog_draft_jobs_active", "source_id", unique=True, 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) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id)
@@ -450,7 +451,8 @@ class CatalogProfileJob(Base, TimestampMixin):
__table_args__ = ( __table_args__ = (
Index("ix_catalog_profile_jobs_source_created", "source_id", "created_at"), Index("ix_catalog_profile_jobs_source_created", "source_id", "created_at"),
Index("ux_catalog_profile_jobs_active", "source_id", unique=True, 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) 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 补建。 # 一个人在一个应用里只有一段(未归档的)助手对话(D-e)。存量库由 db/base.py 的 _UNIQUE_INDEX_MIGRATIONS 补建。
# created_by 为空(关登录又没填署名)的不受约束:SQLite 的唯一索引里 NULL 互不相等,那种情况靠接口里的进程内锁 # created_by 为空(关登录又没填署名)的不受约束:SQLite 的唯一索引里 NULL 互不相等,那种情况靠接口里的进程内锁
__table_args__ = (Index("ux_conversations_app_owner", "app_id", "created_by", unique=True, __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) id: Mapped[str] = mapped_column(String(32), primary_key=True, default=new_id)
title: Mapped[str] = mapped_column(String(200), default="") title: Mapped[str] = mapped_column(String(200), default="")
+6 -5
View File
@@ -19,10 +19,10 @@ from __future__ import annotations
from datetime import datetime from datetime import datetime
from typing import Any 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 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): class SourceObject(Base):
@@ -92,7 +92,8 @@ class IntrospectJob(Base):
__table_args__ = ( __table_args__ = (
Index("ix_introspect_jobs_source_created", "source_id", "created_at"), Index("ix_introspect_jobs_source_created", "source_id", "created_at"),
Index("ux_introspect_jobs_active", "source_id", unique=True, 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) 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) seen_count: Mapped[int] = mapped_column(Integer, default=1)
detected_by: Mapped[dict[str, Any]] = mapped_column(default=dict) detected_by: Mapped[dict[str, Any]] = mapped_column(default=dict)
basis: Mapped[dict[str, Any] | None] = mapped_column(default=None) basis: Mapped[dict[str, Any] | None] = mapped_column(default=None)
before: Mapped[Any] = mapped_column(JSON, nullable=True, default=None) before: Mapped[Any] = mapped_column(PortableJSON, nullable=True, default=None)
after: Mapped[Any] = mapped_column(JSON, nullable=True, default=None) after: Mapped[Any] = mapped_column(PortableJSON, nullable=True, default=None)
affects: Mapped[dict[str, Any] | None] = mapped_column(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_before: Mapped[int | None] = mapped_column(Integer, default=None)
struct_version_after: Mapped[int | None] = mapped_column(Integer, default=None) struct_version_after: Mapped[int | None] = mapped_column(Integer, default=None)
+235
View File
@@ -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)
+76
View File
@@ -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)
+57
View File
@@ -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}"
+87
View File
@@ -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
+16 -11
View File
@@ -10,7 +10,7 @@ from collections.abc import Callable, Sequence
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Any from typing import Any
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.errors import GraphRecursionError from langgraph.errors import GraphRecursionError
from langgraph.types import Command from langgraph.types import Command
from sqlalchemy import case, func, select, update from sqlalchemy import case, func, select, update
@@ -71,7 +71,7 @@ class RunManager:
def __init__(self) -> None: def __init__(self) -> None:
self._tasks: dict[str, asyncio.Task[None]] = {} self._tasks: dict[str, asyncio.Task[None]] = {}
self._checkpointer: AsyncSqliteSaver | None = None self._checkpointer: BaseCheckpointSaver | None = None
self._cm: Any = None self._cm: Any = None
self._semaphore = asyncio.Semaphore(settings.max_concurrent_runs) self._semaphore = asyncio.Semaphore(settings.max_concurrent_runs)
#: 正在关停。此时收到的取消来自关停而不是用户,运行要记成可恢复的中断 #: 正在关停。此时收到的取消来自关停而不是用户,运行要记成可恢复的中断
@@ -86,7 +86,10 @@ class RunManager:
async def setup(self) -> None: async def setup(self) -> None:
self._closing = False self._closing = False
settings.ensure_dirs() 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__() self._checkpointer = await self._cm.__aenter__()
# 上次进程被强杀时留下的记录。running 和 queued 要分开处理: # 上次进程被强杀时留下的记录。running 和 queued 要分开处理:
# running 至少跑过一步、有 checkpoint,可以续;queued 还在排队, # running 至少跑过一步、有 checkpoint,可以续;queued 还在排队,
@@ -163,7 +166,7 @@ class RunManager:
await self._cm.__aexit__(None, None, None) await self._cm.__aexit__(None, None, None)
@property @property
def checkpointer(self) -> AsyncSqliteSaver: def checkpointer(self) -> BaseCheckpointSaver:
if self._checkpointer is None: if self._checkpointer is None:
raise RuntimeError("RunManager 还没初始化") raise RuntimeError("RunManager 还没初始化")
return self._checkpointer return self._checkpointer
@@ -889,6 +892,15 @@ class RunManager:
try: try:
await self._finalize(run_id, status, error, final_state, elapsed, await self._finalize(run_id, status, error, final_state, elapsed,
terminal=terminal, error_node=error_node, end_reason=end_reason) 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: finally:
# 收尾被打断(关停时的取消、库连接出错)也得把自己摘掉。以前写在 # 收尾被打断(关停时的取消、库连接出错)也得把自己摘掉。以前写在
# _finalize 后面,一被打断就漏掉,留下的任务让下一次 shutdown 的 # _finalize 后面,一被打断就漏掉,留下的任务让下一次 shutdown 的
@@ -896,13 +908,6 @@ class RunManager:
self._finalizing.discard(run_id) self._finalizing.discard(run_id)
self._tasks.pop(run_id, None) self._tasks.pop(run_id, None)
self._stopped_by.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 都是新的, # 运行到终态就把沙箱会话收掉。thread_id 每个 run 都是新的,
# 不收的话每跑一次带代码节点的图就多一台常驻 microVM(几百 MB), # 不收的话每跑一次带代码节点的图就多一台常驻 microVM(几百 MB),
# 而且没有任何自动回收路径会碰它——闲置回收只在"下次有人执行 # 而且没有任何自动回收路径会碰它——闲置回收只在"下次有人执行
+1 -1
View File
@@ -294,7 +294,7 @@ def accept_problem(setting: dict[str, Any] | None, *, name: str, mime: str, size
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
#: 还没绑到记录上:JSON 列显式写过 None 的存的是 JSON 的 null,不是 SQL 的 NULL,两种都算 #: 还没绑到记录上: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]: def file_fields(form: Any) -> list[str]:
+3 -2
View File
@@ -44,6 +44,7 @@ from sqlalchemy import func, select
from app.db.base import SessionLocal from app.db.base import SessionLocal
from app.db.models import Run, Workflow from app.db.models import Run, Workflow
from app.db.portable import json_eq
logger = logging.getLogger(__name__) 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]: async def issued_runs(session: Any, run_id: str) -> list[Run]:
"""这次运行按口径出具过的正式运行(origin.source_run 是它),新的在前。origin_kind 有索引,先按它筛再看 origin """这次运行按口径出具过的正式运行(origin.source_run 是它),新的在前。origin_kind 有索引,先按它筛再看 origin
(json_extract,同 auth/keys 认幂等键的写法)。""" (db/portable.json_eq,同 auth/keys 认幂等键的写法)。"""
rows = (await session.execute(select(Run).where( 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() .order_by(Run.created_at.desc()))).scalars()
return list(rows) return list(rows)
+13 -14
View File
@@ -66,11 +66,12 @@ from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from typing import Any from typing import Any
from sqlalchemy import func, select from sqlalchemy import func, not_, select
from app.core.events import EventType from app.core.events import EventType
from app.db.base import SessionLocal, utcnow from app.db.base import SessionLocal, utcnow
from app.db.models import Run, RunEvent, Workflow 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 from app.ledger import uses as ledger_uses
logger = logging.getLogger(__name__) 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]: async def recall_runs(session: Any, workflow_id: str, new_version: int | None = None) -> list[Run]:
"""以 W(新版本 N)发起过的召回运行(试算的不算),新的在前。origin_kind 有索引,先按它筛再看 origin。""" """以 W(新版本 N)发起过的召回运行(试算的不算),新的在前。origin_kind 有索引,先按它筛再看 origin。"""
stmt = select(Run).where(Run.origin_kind == ORIGIN, stmt = select(Run).where(Run.origin_kind == ORIGIN,
func.json_extract(Run.origin, "$.workflow_id") == workflow_id, json_eq(Run.origin, "workflow_id", workflow_id), not_(json_flag(Run.origin, "probe")))
func.coalesce(func.json_extract(Run.origin, "$.probe"), 0) == 0)
if new_version is not None: 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()) 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) state, _ = await outcome_of(session, run)
changed += state == "changed" changed += state == "changed"
as_old = int(await session.scalar(select(func.count()).select_from(Run).where( 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, Run.origin_kind == ORIGIN, json_eq(Run.origin, "workflow_id", workflow_id),
func.json_extract(Run.origin, "$.from_version") == int(version), json_eq(Run.origin, "from_version", int(version)), not_(json_flag(Run.origin, "probe")))) or 0)
func.coalesce(func.json_extract(Run.origin, "$.probe"), 0) == 0)) or 0)
return {"groups": len(as_new), "changed": changed, "recalled_from": as_old} 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) select(Task, Run.origin).join(Run, Run.id == Task.run_id)
.where(Task.kind == TASK_KIND, Task.status == "pending", .where(Task.kind == TASK_KIND, Task.status == "pending",
(Task.run_id == run_id) | ((Run.origin_kind == ORIGIN) (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() .order_by(Task.created_at, Task.id))).all()
mine = await _assignments(session, [t.id for t, _ in candidates]) mine = await _assignments(session, [t.id for t, _ in candidates])
tasks: list[dict[str, Any]] = [] 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) 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", .where(Artifact.kind == QUERY_SNAPSHOT_KIND, Run.run_class == "exploratory", Run.status == "succeeded",
Run.manifest_hash.is_not(None), counted_runs(), Run.manifest_hash.is_not(None), counted_runs(),
or_(func.json_extract(Artifact.meta, "$.source_id") == source.id, or_(json_eq(Artifact.meta, "source_id", source.id),
func.json_extract(Artifact.meta, "$.source") == source.name)) json_eq(Artifact.meta, "source", source.name)))
.order_by(Run.finished_at.desc(), Run.id, Artifact.id)) .order_by(Run.finished_at.desc(), Run.id, Artifact.id))
cut = _aware(before) cut = _aware(before)
offset = 0 offset = 0
@@ -1473,10 +1472,10 @@ async def _probe_calibers(change: Mapping[str, Any]) -> list[str]:
continue continue
async with SessionLocal() as session: async with SessionLocal() as session:
done = await session.scalar(select(func.count()).select_from(Run).where( done = await session.scalar(select(func.count()).select_from(Run).where(
Run.origin_kind == ORIGIN, func.json_extract(Run.origin, "$.probe") == 1, Run.origin_kind == ORIGIN, json_flag(Run.origin, "probe"),
func.json_extract(Run.origin, "$.workflow_id") == caliber.workflow_id, json_eq(Run.origin, "workflow_id", caliber.workflow_id),
func.json_extract(Run.origin, "$.version") == caliber.version, json_eq(Run.origin, "version", caliber.version),
func.json_extract(Run.origin, "$.change") == change.get("id"))) json_eq(Run.origin, "change", change.get("id"))))
if done: if done:
continue continue
latest = (await session.execute(select(Run.input).where( latest = (await session.execute(select(Run.input).where(
+2 -2
View File
@@ -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]: async def reconciled_runs(session: Any, run_id: str) -> list[Run]:
"""对这次运行对过的账(origin.source_run 是它),新的在前。""" """对这次运行对过的账(origin.source_run 是它),新的在前。"""
from sqlalchemy import func from app.db.portable import json_eq
rows = (await session.execute(select(Run).where( 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() .order_by(Run.created_at.desc()))).scalars()
return list(rows) return list(rows)
+2 -1
View File
@@ -35,6 +35,7 @@ from app.db.base import SessionLocal, utcnow
from app.db.models import ( from app.db.models import (
App, Approval, AppVersion, AuditEvent, Group, GroupMember, Run, Task, TaskAssignment, User, 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.engine.lint_process import normalize_policy, parse_duration
from app.process.events import record from app.process.events import record
@@ -913,7 +914,7 @@ async def tick(now: datetime | None = None) -> list[str]:
"""调度循环每一拍:提醒、升级、超时驳回(永不自动批准)。交回这一拍处理过的待办 id(生效可查:办了哪几条, """调度循环每一拍:提醒、升级、超时驳回(永不自动批准)。交回这一拍处理过的待办 id(生效可查:办了哪几条,
不只是几条;测试按自己造的那条断言,不受同一个库里别的待办影响)。试运行的待办不处理(只显示时限)。""" 不只是几条;测试按自己造的那条断言,不受同一个库里别的待办影响)。试运行的待办不处理(只显示时限)。"""
now = now or utcnow() 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) trial = select(Run.id).where(Run.origin_kind == TRIAL_ORIGIN)
async with SessionLocal() as session: async with SessionLocal() as session:
due = list((await session.execute(select(Task.id).where( due = list((await session.execute(select(Task.id).where(
+12
View File
@@ -20,6 +20,9 @@ dependencies = [
# 所以上限卡到下一个小版本——升级要主动改这里、重新锁、跑一遍测试。 # 所以上限卡到下一个小版本——升级要主动改这里、重新锁、跑一遍测试。
# 精确版本在 requirements.lock 里 # 精确版本在 requirements.lock 里
"langgraph>=1.2.11,<1.3", "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", "langgraph-checkpoint-sqlite>=3.1.1,<3.2",
"langchain-core>=1.6.3,<1.7", "langchain-core>=1.6.3,<1.7",
"langchain-anthropic>=1.7.2,<1.8", "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 原生且免装数据库客户端—— # 数据库接入。三个驱动都是 asyncio 原生且免装数据库客户端——
# oracledb 走 thin 模式,省掉了 Oracle Instant Client 这个历来最劝退的一步。 # oracledb 走 thin 模式,省掉了 Oracle Instant Client 这个历来最劝退的一步。
db = ["aiomysql>=0.2", "oracledb>=2.0", "asyncpg>=0.30"] 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] [build-system]
requires = ["hatchling"] requires = ["hatchling"]
+2
View File
@@ -3,3 +3,5 @@
asyncio_mode = auto asyncio_mode = auto
testpaths = tests testpaths = tests
addopts = -q --tb=short addopts = -q --tb=short
markers =
sqlite_only: 只对 SQLite 应用库成立的测试(数据目录里的库文件、PRAGMA);AGENTLAB_TEST_DATABASE_URL 跑 PostgreSQL 时跳过
+19
View File
@@ -136,9 +136,13 @@ langgraph==1.2.11
# via agentlab-backend (pyproject.toml) # via agentlab-backend (pyproject.toml)
langgraph-checkpoint==4.2.0 langgraph-checkpoint==4.2.0
# via # via
# agentlab-backend (pyproject.toml)
# langgraph # langgraph
# langgraph-checkpoint-postgres
# langgraph-checkpoint-sqlite # langgraph-checkpoint-sqlite
# langgraph-prebuilt # langgraph-prebuilt
langgraph-checkpoint-postgres==3.1.2
# via agentlab-backend (pyproject.toml)
langgraph-checkpoint-sqlite==3.1.1 langgraph-checkpoint-sqlite==3.1.1
# via agentlab-backend (pyproject.toml) # via agentlab-backend (pyproject.toml)
langgraph-prebuilt==1.1.0 langgraph-prebuilt==1.1.0
@@ -167,6 +171,7 @@ oracledb==26.0.0
# via agentlab-backend (pyproject.toml) # via agentlab-backend (pyproject.toml)
orjson==3.12.0 orjson==3.12.0
# via # via
# langgraph-checkpoint-postgres
# langgraph-sdk # langgraph-sdk
# langsmith # langsmith
ormsgpack==1.12.2 ormsgpack==1.12.2
@@ -180,6 +185,16 @@ pillow==12.3.0
# via python-pptx # via python-pptx
pluggy==1.6.0 pluggy==1.6.0
# via pytest # 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' pycparser==3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy'
# via cffi # via cffi
pydantic==2.13.5 pydantic==2.13.5
@@ -294,6 +309,8 @@ typing-extensions==4.16.0
# mcp # mcp
# openai # openai
# oracledb # oracledb
# psycopg
# psycopg-pool
# pydantic # pydantic
# pydantic-core # pydantic-core
# pytest-asyncio # pytest-asyncio
@@ -309,6 +326,8 @@ typing-inspection==0.4.4
# mcp # mcp
# pydantic # pydantic
# pydantic-settings # pydantic-settings
tzdata==2026.5 ; sys_platform == 'win32'
# via psycopg
urllib3==2.7.0 urllib3==2.7.0
# via requests # via requests
uuid-utils==0.17.1 uuid-utils==0.17.1
+77
View File
@@ -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}"])
+58 -1
View File
@@ -12,12 +12,58 @@ app.db.base 导入时就按 db_url 建好了。conftest 早于测试模块加载
from __future__ import annotations from __future__ import annotations
import atexit
import os import os
import tempfile import tempfile
from pathlib import Path from pathlib import Path
_TMP = Path(tempfile.mkdtemp(prefix="agentlab-tests-")) _TMP = Path(tempfile.mkdtemp(prefix="agentlab-tests-"))
os.environ["AGENTLAB_DATA_DIR"] = str(_TMP) 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 带进测试环境 # 别把真实的 key 带进测试环境
os.environ.setdefault("AGENTLAB_SECRET_KEY", "test-only-not-a-real-key") os.environ.setdefault("AGENTLAB_SECRET_KEY", "test-only-not-a-real-key")
# 定时任务的调度循环不在测试里起:走 lifespan 的测试不该到点替人发起运行。要测调度的直接调它 # 定时任务的调度循环不在测试里起:走 lifespan 的测试不该到点替人发起运行。要测调度的直接调它
@@ -42,7 +88,18 @@ def _prepare_schema():
def pytest_report_header(config) -> str: 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") @pytest.fixture(scope="session")
+123
View File
@@ -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
+8 -1
View File
@@ -556,9 +556,16 @@ async def test_startup_locks_legacy_sources_registered_under_other_spellings(cli
def _app_db() -> str: 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 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) return str(db_base.engine.url.database)
+357
View File
@@ -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"]
+2 -7
View File
@@ -446,7 +446,7 @@ async def test_a_continue_that_wins_the_row_between_check_and_write_gets_no_even
"""核对过状态、还没写入的那一下,接着跑的条件更新(runner._claim)先提交了:追加判定的条件更新落空, """核对过状态、还没写入的那一下,接着跑的条件更新(runner._claim)先提交了:追加判定的条件更新落空,
照样不记——核对和写入是同一行上的条件更新,不是先读后写。""" 照样不记——核对和写入是同一行上的条件更新,不是先读后写。"""
from app.core.bus import bus from app.core.bus import bus
from app.core.config import settings from tests import appdb
Models(monkeypatch, verdicts={"新客": "contradicted"}) Models(monkeypatch, verdicts={"新客": "contradicted"})
row = await finish(weekly(claims="judge")) 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): def claimed_in_between(run_id, value):
# _append_judged 读完运行、分配序号之前正好调用它:在这里用另一条连接把运行占成 running # _append_judged 读完运行、分配序号之前正好调用它:在这里用另一条连接把运行占成 running
if run_id == row.id: if run_id == row.id:
db = sqlite3.connect(settings.db_path, timeout=5) appdb.execute(("UPDATE runs SET status = 'running' WHERE id = ?", (row.id,)))
try:
db.execute("UPDATE runs SET status = 'running' WHERE id = ?", (row.id,))
db.commit()
finally:
db.close()
set_seq(run_id, value) set_seq(run_id, value)
monkeypatch.setattr(bus, "set_seq", claimed_in_between) monkeypatch.setattr(bus, "set_seq", claimed_in_between)
+9 -34
View File
@@ -16,7 +16,6 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import sqlite3
from collections import Counter from collections import Counter
import pytest import pytest
@@ -25,7 +24,6 @@ from pydantic import BaseModel
from sqlalchemy import select from sqlalchemy import select
from app.core.bus import bus from app.core.bus import bus
from app.core.config import settings
from app.core.events import EventType from app.core.events import EventType
from app.db.base import SessionLocal from app.db.base import SessionLocal
from app.db.models import Approval, Run, RunEvent 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.runner import RunManager, run_manager, verify_manifest
from app.engine.schema import NodeType from app.engine.schema import NodeType
from app.tools.registry import register from app.tools.registry import register
from tests import appdb
#: Agent 场景里工具真正执行的次数。注册表是进程级的,工具只注册一次,计数按测试清零 #: Agent 场景里工具真正执行的次数。注册表是进程级的,工具只注册一次,计数按测试清零
CALLS: list[str] = [] CALLS: list[str] = []
@@ -162,12 +161,7 @@ async def _pending(run_id: str) -> list[Approval]:
def _sql(statement: str, *params) -> None: def _sql(statement: str, *params) -> None:
"""另开一条连接改库:模拟别处(恢复、继续运行、另一个写入者)抢在中间提交。""" """另开一条连接改库:模拟别处(恢复、继续运行、另一个写入者)抢在中间提交。"""
db = sqlite3.connect(settings.db_path, timeout=5) appdb.execute((statement, params))
try:
db.execute(statement, params)
db.commit()
finally:
db.close()
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
@@ -226,12 +220,7 @@ async def test_a_resume_that_wins_the_row_leaves_nothing_behind(monkeypatch):
def resumed_in_between(run_id, value): def resumed_in_between(run_id, value):
if run_id == paused.id: if run_id == paused.id:
db = sqlite3.connect(settings.db_path, timeout=5) appdb.execute(("UPDATE runs SET status = 'running' WHERE id = ?", (paused.id,)))
try:
db.execute("UPDATE runs SET status = 'running' WHERE id = ?", (paused.id,))
db.commit()
finally:
db.close()
set_seq(run_id, value) set_seq(run_id, value)
monkeypatch.setattr(bus, "set_seq", resumed_in_between) 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): def written_in_between(run_id, value):
if run_id == paused.id: if run_id == paused.id:
db = sqlite3.connect(settings.db_path, timeout=5) appdb.execute(("INSERT INTO run_events (id, run_id, seq, type, node_id, ts, data) "
try: "VALUES ('elsewhere-1', ?, ?, 'log', NULL, 0, '{}')", (paused.id, taken)),
db.execute("INSERT INTO run_events (id, run_id, seq, type, node_id, ts, data) " ("UPDATE runs SET last_seq = ? WHERE id = ?", (taken, paused.id)))
"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()
set_seq(run_id, value) set_seq(run_id, value)
monkeypatch.setattr(bus, "set_seq", written_in_between) 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): async def stored_when_published(event):
if event.type == EventType.APPROVAL_INQUIRY: if event.type == EventType.APPROVAL_INQUIRY:
db = sqlite3.connect(settings.db_path, timeout=5) [(count,)] = appdb.query("SELECT count(*) FROM run_events WHERE run_id = ? AND seq = ?",
try: event.run_id, event.seq)
(count,) = db.execute("SELECT count(*) FROM run_events WHERE run_id = ? AND seq = ?",
(event.run_id, event.seq)).fetchone()
finally:
db.close()
published.append((event.seq, count)) published.append((event.seq, count))
await publish(event) 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): def changed_in_between(run_id, value):
if run_id == final.id: if run_id == final.id:
db = sqlite3.connect(settings.db_path, timeout=5) appdb.execute((change, (final.id,)))
try:
db.execute(change, (final.id,))
db.commit()
finally:
db.close()
set_seq(run_id, value) set_seq(run_id, value)
monkeypatch.setattr(bus, "set_seq", changed_in_between) monkeypatch.setattr(bus, "set_seq", changed_in_between)
+6 -4
View File
@@ -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.db.base import engine
from app.engine.launch import launch_run from app.engine.launch import launch_run
from tests.appdb import refuse_inserts
lib = await upstream_with_two_versions() lib = await upstream_with_two_versions()
wf, _ = await workflow(weekly(lib, upgrade_policy="dual"), published=1, versions=(1,)) wf, _ = await workflow(weekly(lib, upgrade_policy="dual"), published=1, versions=(1,))
before = await _run_count() before = await _run_count()
# 口径升版处置那条写不进去。走异步引擎建删触发器:同步连接等写锁时会卡住事件循环,而持锁的正是这个循环里的写入 # 口径升版处置那条写不进去。走异步引擎建删触发器:同步连接等写锁时会卡住事件循环,而持锁的正是这个循环里的写入
create, drop = refuse_inserts("run_events", "NEW.type = 'caliber.upgrade'", name="refuse_upgrade_note")
async with engine.begin() as conn: async with engine.begin() as conn:
await conn.execute(text( for sql in create:
"CREATE TRIGGER refuse_upgrade_note BEFORE INSERT ON run_events WHEN NEW.type = 'caliber.upgrade' " await conn.execute(text(sql))
"BEGIN SELECT RAISE(ABORT, 'refused by test'); END"))
try: try:
async with SessionLocal() as session: async with SessionLocal() as session:
with pytest.raises(IntegrityError): 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, "没写进治理事件,却留下了一条运行" assert await _run_count() == before, "没写进治理事件,却留下了一条运行"
finally: finally:
async with engine.begin() as conn: 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(): async def test_an_exploratory_launch_takes_the_saved_graph_and_the_run_defaults():
+4 -2
View File
@@ -112,9 +112,11 @@ def test_a_missing_table_is_left_alone():
async def test_new_databases_index_schedule_id(): async def test_new_databases_index_schedule_id():
from sqlalchemy import inspect
async with app_engine.connect() as conn: async with app_engine.connect() as conn:
names = {row[1] for row in (await conn.execute(text("PRAGMA index_list(runs)"))).all()} indexes = await conn.run_sync(lambda sync: inspect(sync).get_indexes("runs"))
assert "ix_runs_schedule_id" in names 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(): async def test_a_migrated_row_reads_back_empty_options():
+3 -7
View File
@@ -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): async def test_an_approval_closed_just_before_it_is_marked_answered(slow, monkeypatch):
"""放弃恰好提交在「读到这条审批还是 pending」和「把它记成已回复」之间。""" """放弃恰好提交在「读到这条审批还是 pending」和「把它记成已回复」之间。"""
import sqlite3
from sqlalchemy import event from sqlalchemy import event
from app.db.base import engine from app.db.base import engine
from tests import appdb
run = await run_manager.start(graph=_graph(gate=True), input_payload={"q": "x"}) run = await run_manager.start(graph=_graph(gate=True), input_payload={"q": "x"})
await _until(lambda: _is(run.id, "interrupted")) 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): def abandon_first(conn, cursor, statement, params, context, executemany):
if armed["on"] and statement.lstrip().upper().startswith("UPDATE APPROVALS"): if armed["on"] and statement.lstrip().upper().startswith("UPDATE APPROVALS"):
armed["on"] = False armed["on"] = False
other = sqlite3.connect(engine.url.database, timeout=5) appdb.execute(("UPDATE approvals SET status='cancelled' WHERE run_id=?", (run.id,)),
other.execute("UPDATE approvals SET status='cancelled' WHERE run_id=?", (run.id,)) ("UPDATE runs SET status='cancelled' WHERE id=?", (run.id,)))
other.execute("UPDATE runs SET status='cancelled' WHERE id=?", (run.id,))
other.commit()
other.close()
event.listen(engine.sync_engine, "before_cursor_execute", abandon_first) event.listen(engine.sync_engine, "before_cursor_execute", abandon_first)
try: try:
+5 -1
View File
@@ -23,6 +23,7 @@ from typing import Any
import pytest import pytest
from sqlalchemy import create_engine, select, text from sqlalchemy import create_engine, select, text
from sqlalchemy.exc import IntegrityError
from app.core.config import settings from app.core.config import settings
from app.data import raw_store, recipe_parsers, recipe_types, table_versions 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 ( from app.db.models import (
DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe, DataSource, ImportStaging, SourceSnapshot, TableBuild, TableImport, TableRecipe,
) )
from app.db.portable import unique_violation
@pytest.fixture(autouse=True) @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)) session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256=RECIPE_SHA))
await session.commit() await session.commit()
session.add(TableRecipe(source_id=sid, seq=1, recipe={}, recipe_sha256=RECIPE_SHA)) 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() await session.commit()
# 撞的正是 (source_id, seq) 这一条唯一约束(两种库的报错写法不同,按表和列认)
assert unique_violation(caught.value, "table_recipes", "source_id", "seq")
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
+53
View File
@@ -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[@]}"