Files
AgentLab/backend/app/data/engine.py
T
ThearchyHelios fcdcc30d4e 可点击证据第二期:取数链路打通、问数据的答案逐段可点;发布前自动修复
一、可点击证据第二期

取数链路
- 智能体、工具节点、代码节点、知识检索取数时都记进证据台账;报告目录按台账编出 Q1… / K1…
- 报告可以引用快照里的一格([[v:Q1.r0.gmv]]),也可以由系统从快照生成整张表([[table:Q1 cols=… rows=0-4]]),
  表里每一格都是可点的片段;取值一律以快照为准,取回快照时复验哈希
- 口径卡表达式新增 cell(x, row, column),lineage 精确到单元格
- 智能体开 cite_fields(配合 output_schema)后,循环结束做一次结构化抽取,每个字段标出取自哪次查询的哪一格,
  再到快照里核对:一致记 verified;不一致以快照为准并告警;查不到的记为空值,永远不兜底成 0
- 代码节点新增 evidence_role(source / compute)和 code_sha;沙箱算出来的数不能被报告直接引用,要进口径卡
- 查询快照记下列类型,文本列按文本渲染(年份 "2026" 不再显示成「2,026」)
- 升级前发起的运行(停在审批上的)一律走旧逻辑:新的 task 只追加在循环之后,恢复时不会错号

证据接口与面板
- 片段链接上查询步骤:SQL、被引用的行加前后各 2 行、高亮到格;数据源可设 mask_columns 遮罩列,
  遮罩在查询时记进快照,数据源事后改名或删除也照样遮罩(在有身份体系之前,遮罩只减少暴露,不是安全边界)
- 输入来源写明核对结果(与快照一致 / 模型报的值和快照不一致 / 没查到记为空);口径卡缺输入时也给出完整的链
- 受管出具要在契约里声明 cells 才能引用单元格;受管与否记在所钉版本上,改画布不再让受管版本的正式运行降级
- 口径卡可以用 caliber_from 钉住另一个已发布工作流里的口径卡,上游有新版本时按 upgrade_policy 处置,
  和子工作流同一套规则;子工作流节点也补上了升版处置的配置项

Copilot、校验与报错
- 问数据的答案里有数字时,搭 input → 智能体 → 报告 → 出口,报告直接引用查过的单元格
- 写 SQL 前核对数据源的表结构,SQL 里出现表结构里没有的列会被打回
- 不再搭「智能体文字 → 转换节点解析 JSON」这种图;转换节点的 JSON 报错分清是上游模型写坏了、
  模板里裸插了模型文字(给出 | json 的改法),还是模板本身写错了
- 工具节点遇到「查询失败」「SQL 被拒绝」判为失败;发起运行前检查绑定的工具是否存在,不存在时直接拒绝并给出入口

二、发布前自动修复

- 问题带上机器编号(code)和修复 id;新增只读接口 POST /api/workflows/{id}/publish-check 和 /autofix
- 修复分三类:答案唯一、不降低要求的一键修复;要人拿主意的给选项(不替人选);结构性问题交给 Copilot
- 一律先预览,不自动保存、不自动发布;删节点、删契约、放宽审批、关掉 strict、改级别等降低要求的改动一律拒绝,
  Copilot 的方案也要过这一关;应用后重跑检查,错误必须变少且不能冒出新的
- 发布弹窗打开就先检查;问题面板新增「发布前检查」,平时就能看到、就地修掉
- 门禁修正:契约有 report_from 时不再要求 narrative(模板⑨可以发成受管);全图默认审批为「全部放行」时同样拦下
- 修掉画布保存会抹掉全图默认(defaults)的老问题:以前任何一次保存都会把它清空,审批 always 因此被悄悄放宽

验证(沙箱):tsc 干净,生产构建通过,pytest 2268 passed(两遍),check-all 17 项检查、3915 项全过
(新增 check-publish);关键行为做了变异验证、全部抓到。
2026-09-28 21:20:04 +08:00

646 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""数据源连接层。各家数据库的方言差异全部收敛在这个文件里。
上层(工具、探查、API)只认 DataSource 这个模型,不需要知道 Oracle 要 service_name、
MySQL 要 charset、asyncpg 不认 sslmode 这些事。哪天要加一种库,改这里一处。
"""
from __future__ import annotations
import asyncio
import json
import logging
import re
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncIterator, Awaitable, Callable
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
from sqlalchemy import text
from app.core.crypto import decrypt
from app.data import guard
from app.data.guard import QueryLimits, SqlRejected, check, is_write
logger = logging.getLogger(__name__)
# 各方言的 SQLAlchemy 驱动。都选 asyncio 原生驱动,没有一个需要装数据库客户端:
# oracledb 走 thin 模式(纯 Python 协议实现),省掉了 Oracle Instant Client 这个
# 历来最劝退的一步。
_DRIVERS = {
"mysql": "mysql+aiomysql",
"mariadb": "mysql+aiomysql",
"postgres": "postgresql+asyncpg",
"postgresql": "postgresql+asyncpg",
"oracle": "oracle+oracledb",
"sqlite": "sqlite+aiosqlite",
}
_DEFAULT_PORTS = {
"mysql": 3306, "mariadb": 3306,
"postgres": 5432, "postgresql": 5432,
"oracle": 1521,
}
SUPPORTED_KINDS = tuple(sorted(set(_DRIVERS)))
#: 数据源 options 里按源配置的查询时限(秒)。表单上叫「查询时限」
QUERY_TIMEOUT_OPTION = "query_timeout_s"
#: 按源配置时能填的范围。再长的查询不该让一个 agent 干等:该做成离线任务了
MIN_QUERY_TIMEOUT_S = 1
MAX_QUERY_TIMEOUT_S = 600
#: 数据源 options 里要遮罩的列(列名列表,或者逗号分隔的一段文字)。证据面板展示原始行时,
#: 这些列的值一律换成「已遮罩」。在有身份体系之前,遮罩只减少暴露,不是安全边界:
#: 完整快照仍能按工件 id 取到,SQL 里给列起个别名也能绕开
MASK_COLUMNS_OPTION = "mask_columns"
# options 里这些 key 是 AgentLab 自己的配置,不是驱动参数,拼 URL 时要摘掉。
# schema:探查哪个 schema(企业库里只读账号名下常常什么都没有,数据在别处)
# query_timeout_s:查询时限,由数据层按语句下发给数据库,见 _server_deadline
# mask_columns:证据面板遮罩的列,驱动不认识它,拼进连接串会被当成未知参数拒掉
_NON_DRIVER_OPTIONS = frozenset({"schema", QUERY_TIMEOUT_OPTION, MASK_COLUMNS_OPTION})
_MASK_SPLIT = re.compile(r"[,,、;;\n]+")
@dataclass
class QueryResult:
columns: list[str]
rows: list[list[Any]]
row_count: int
truncated: bool
elapsed_ms: int
sql: str
#: 列名 → number / text / date / datetime / time / boolean。拿不准的列(全是空值、类型混着)不在里面
column_types: dict[str, str] | None = None
def to_payload(self) -> dict[str, Any]:
payload: dict[str, Any] = {
"columns": self.columns,
"rows": self.rows,
"row_count": self.row_count,
"truncated": self.truncated,
"elapsed_ms": self.elapsed_ms,
"sql": self.sql,
}
if self.column_types:
payload["column_types"] = self.column_types
return payload
def build_url(source: Any, *, reveal: bool = False) -> str:
"""拼连接串。reveal=False 时把密码换成星号——日志和错误信息里用这个。
密码永远不该出现在日志、事件流、或者返回给前端的任何地方。默认遮掉,
真要连库时才显式 reveal=True。
"""
kind = (source.kind or "").lower()
driver = _DRIVERS.get(kind)
if not driver:
raise ValueError(f"不支持的数据库类型:{source.kind}(支持 {', '.join(SUPPORTED_KINDS)})")
options: dict[str, Any] = dict(source.options or {})
# options 混着两类东西:驱动参数(charset、service_name…)和我们自己的配置
# (schema 指探查哪个 schema)。后者不能拼进连接串,否则驱动会把它当成
# 未知的连接参数直接拒掉。
for key in _NON_DRIVER_OPTIONS:
options.pop(key, None)
if kind == "sqlite":
# SQLite 没有主机和账号,database 就是文件路径
return f"{driver}:///{source.database or ':memory:'}"
password = decrypt(source.password) or "" if reveal else "***"
from urllib.parse import quote_plus
user = quote_plus(source.username or "")
pw = quote_plus(password) if reveal else password
auth = f"{user}:{pw}@" if user else ""
port = source.port or _DEFAULT_PORTS.get(kind)
host = f"{source.host or 'localhost'}:{port}" if port else (source.host or "localhost")
if kind == "oracle":
# Oracle 用 service_name 或 sid,不是"数据库名"。这是最常见的配错点,
# 所以两种都认,并且在都没填时给出明确提示而不是让驱动报晦涩的 ORA 错误。
service = options.pop("service_name", None) or source.database
sid = options.pop("sid", None)
if not service and not sid:
raise ValueError("Oracle 需要在 options 里给 service_name 或 sid(也可以填在库名里)")
tail = f"?service_name={service}" if service else f"?sid={sid}"
query = "&".join(f"{k}={v}" for k, v in options.items())
return f"{driver}://{auth}{host}/{tail}{('&' + query) if query else ''}"
query = "&".join(f"{k}={v}" for k, v in options.items())
return f"{driver}://{auth}{host}/{source.database or ''}{('?' + query) if query else ''}"
def engine_args(source: Any) -> tuple[str, dict[str, Any]]:
"""真正连库用的连接串和额外参数。只读源在这里拿到**连接层**的只读。
守卫按关键字判定,而模型能写出什么 SQL 事前无法穷举——`WITH … DELETE`
就曾从首关键字白名单下面钻过去,在 SQLite 上真把数据删了(pysqlite 只在
INSERT/UPDATE/DELETE 开头的语句前隐式开事务,WITH 开头的直接自动提交)。
所以只读不能只靠猜,连接本身就得写不进去:
- SQLite:以 mode=ro 打开文件,写操作由 SQLite 自己拒绝。
- PostgreSQL:会话的默认事务只读。能关掉它的 SET / set_config() 守卫都拦。
- MySQL / MariaDB:会话事务只读。能关掉它的只有 SET,守卫拦。
- Oracle:没有会话级只读开关,只剩守卫这一道——只读数据源请配只读账号。
"""
url = build_url(source, reveal=True)
if not source.readonly:
return url, {}
kind = (source.kind or "").lower()
if kind == "sqlite":
path = source.database or ""
if path and path != ":memory:":
from urllib.parse import quote
url = f"{_DRIVERS['sqlite']}:///file:{quote(path)}?mode=ro&uri=true"
return url, {}
if kind in ("postgres", "postgresql"):
return url, {"connect_args": {"server_settings": {"default_transaction_read_only": "on"}}}
if kind in ("mysql", "mariadb"):
return url, {"connect_args": {"init_command": "SET SESSION TRANSACTION READ ONLY"}}
return url, {}
def _timeout_value(value: Any) -> float | None:
"""options 里的查询时限:表单存的是字符串,也认数字。填得不对返回 None。"""
if value is None or isinstance(value, bool):
return None
try:
seconds = float(str(value).strip())
except ValueError:
return None
if not MIN_QUERY_TIMEOUT_S <= seconds <= MAX_QUERY_TIMEOUT_S:
return None
return int(seconds) if seconds.is_integer() else seconds
def query_timeout_problem(options: dict[str, Any] | None) -> str | None:
"""保存数据源时查一下查询时限填得对不对;没填、或者填对了返回 None。"""
value = (options or {}).get(QUERY_TIMEOUT_OPTION)
if value is None or str(value).strip() == "":
return None
if _timeout_value(value) is None:
return (f"查询时限要填 {MIN_QUERY_TIMEOUT_S} 到 {MAX_QUERY_TIMEOUT_S} 之间的秒数,"
f"比如 60;现在填的是「{value}」")
return None
def masked_columns(options: dict[str, Any] | None) -> list[str]:
"""options.mask_columns → 列名列表(去空白、去重,保持顺序)。列表和「a, b」这样的文字都认。"""
raw = (options or {}).get(MASK_COLUMNS_OPTION)
items = _MASK_SPLIT.split(raw) if isinstance(raw, str) else raw if isinstance(raw, list) else []
return list(dict.fromkeys(str(c).strip() for c in items if isinstance(c, (str, int)) and str(c).strip()))
def mask_columns_problem(options: dict[str, Any] | None) -> str | None:
"""保存数据源时查一下遮罩列填得对不对;没填、或者填对了返回 None。"""
raw = (options or {}).get(MASK_COLUMNS_OPTION)
if raw is None or raw == "" or raw == []:
return None
if isinstance(raw, str) or (isinstance(raw, list) and all(isinstance(c, str) for c in raw)):
return None
return f"遮罩的列要填列名,用逗号分开,比如「phone, email」;现在填的是「{raw}」"
def query_timeout(source: Any) -> float:
"""这个数据源上一条查询最多跑多久(秒)。没配、或者库里存着的值不对,用缺省。
缺省按模块属性现取:只有 guard.QueryLimits 这一处定义。
"""
value = _timeout_value((source.options or {}).get(QUERY_TIMEOUT_OPTION))
return value if value is not None else guard.QueryLimits().timeout_seconds
def _seconds(value: float) -> str:
return str(int(value)) if float(value).is_integer() else f"{value:g}"
class EngineCache:
"""按数据源 id 缓存 engine。连接池的建立不便宜,不该每次查询都重来一遍。
配置变了要显式 invalidate——改了密码却还在用旧连接,排查起来很费神;
改了只读开关却还在用旧连接,只读就形同虚设。
"""
def __init__(self) -> None:
self._engines: dict[str, AsyncEngine] = {}
self._lock = asyncio.Lock()
async def get(self, source: Any) -> AsyncEngine:
key = str(source.id)
cached = self._engines.get(key)
if cached is not None:
return cached
async with self._lock:
if key in self._engines:
return self._engines[key]
url, extra = engine_args(source)
engine = create_async_engine(
url,
pool_size=3,
max_overflow=2,
pool_pre_ping=True, # 长时间空闲后连接会被数据库掐掉,先探活再用
pool_recycle=1800,
**extra,
)
self._engines[key] = engine
return engine
async def invalidate(self, source_id: str) -> None:
engine = self._engines.pop(str(source_id), None)
if engine is not None:
await engine.dispose()
async def close(self) -> None:
for engine in list(self._engines.values()):
await engine.dispose()
self._engines.clear()
engines = EngineCache()
async def test_connection(source: Any) -> dict[str, Any]:
"""测连接。失败时把驱动的原始错误带出来——这类问题九成靠错误信息定位。"""
started = time.perf_counter()
try:
engine = await engines.get(source)
async with engine.connect() as conn:
probe = "SELECT 1 FROM DUAL" if (source.kind or "").lower() == "oracle" else "SELECT 1"
await conn.execute(text(probe))
return {
"ok": True,
"elapsed_ms": int((time.perf_counter() - started) * 1000),
"url": build_url(source), # 遮掉密码的版本
}
except Exception as e: # noqa: BLE001
return {
"ok": False,
"error": f"{type(e).__name__}: {e}",
"elapsed_ms": int((time.perf_counter() - started) * 1000),
"url": build_url(source),
}
#: 数据库按时限停下语句之后,数据层自己的兜底再多等多久。要比引擎的宽限
#: (toolcalls.GRACE_S)短:数据库报上来的「超过 Ns 被中断」才是给人看的那句,
#: 引擎那边只在数据层也卡住时才出面
BACKSTOP_S = 0.5
#: SQLite 每执行多少条虚拟机指令问一次进度回调。回调只比一次时间,几毫秒一问足够准
_SQLITE_STEPS = 10_000
#: 各家数据库按时限停下语句时的报错。认出来就说「超时」,而不是把驱动原文当成
#: SQL 写错了交给模型——它会去改一条本来没错的 SQL。
#:
#: 先认错误码:服务器的报错文字会跟着它的语言设置走,中文环境的 PostgreSQL 报的是
#: 「由于语句执行超时,正在取消查询命令」,按英文原文找一个也认不出来
_PG_CANCELED = "57014" # PostgreSQL query_canceled(statement_timeout 到点)
_MYSQL_TIMEOUT_CODES = frozenset({3024, # MySQL:max_execution_time 到点
1969}) # MariaDB:max_statement_time 到点
_TIMEOUT_MARKERS = (
"statement timeout", # PostgreSQL(英文环境)
"maximum statement execution time", # MySQL 3024
"max_statement_time", # MariaDB 1969
"dpy-4024", "dpi-1067", "call timeout", # python-oracledb:call timeout of N ms exceeded(驱动自己的话,不随服务器语言变)
)
def _causes(e: BaseException) -> list[BaseException]:
"""异常连同它包着的原始驱动异常。SQLAlchemy 包一层(.orig),适配器再包一层(__cause__),
有的路径(服务端游标取行)又原样抛出驱动异常,所以几条链都走一遍。"""
out: list[BaseException] = []
todo: list[Any] = [e]
while todo and len(out) < 8:
err = todo.pop(0)
if isinstance(err, BaseException) and err not in out:
out.append(err)
todo += [getattr(err, "orig", None), err.__cause__, err.__context__]
return out
class _Deadline:
"""一条语句的时限。数据库那边按它停下;停下时报的错据此认成超时。
从取连接之前起算,和后端兜底同一个起点:取连接花掉的也算在时限里,数据库只拿到
剩下的那段(remaining_ms)。否则取连接一慢,兜底就抢在数据库前面取消语句。
"""
def __init__(self, seconds: float) -> None:
self.seconds = seconds
self.ends = time.monotonic() + seconds
#: SQLite 的进度回调真的掐过
self.fired = False
#: 语句已经结束。回调这时要是还挂在连接上(摘除没成功),也不能再掐别人的语句
self.over = False
def remaining_ms(self) -> int:
return max(1, int((self.ends - time.monotonic()) * 1000))
def remaining_seconds(self) -> float:
return max(0.001, round(self.ends - time.monotonic(), 3))
def expired(self) -> bool:
return time.monotonic() >= self.ends
def sqlite_progress(self) -> int:
# 在 SQLite 执行语句的那个线程里被调用;返回非零,SQLite 就中止当前语句
if not self.over and time.monotonic() >= self.ends:
self.fired = True
return 1
return 0
def stopped(self, e: BaseException) -> bool:
"""这个报错是不是数据库按时限停下语句造成的。"""
if self.fired:
return True
for err in _causes(e):
if _PG_CANCELED in (getattr(err, "sqlstate", None), getattr(err, "pgcode", None)):
return True
args = getattr(err, "args", None) or ()
if args and isinstance(args[0], int) and args[0] in _MYSQL_TIMEOUT_CODES:
return True
low = str(err).lower()
if any(marker in low for marker in _TIMEOUT_MARKERS):
return True
# 各家的措辞不一(interrupted / canceled,或者干脆只在类名里);
# 过了时限才报出来的,就是时限掐的
if self.expired() and any(w in f"{type(err).__name__} {low}".lower()
for w in ("interrupt", "cancel")):
return True
return False
async def _driver(conn: AsyncConnection) -> Any:
return (await conn.get_raw_connection()).driver_connection
@asynccontextmanager
async def _server_deadline(conn: AsyncConnection, kind: str, deadline: _Deadline) -> AsyncIterator[None]:
"""让数据库自己在时限处停下这条语句。
只靠后端这边 wait_for 放弃等待是不够的:数据库那边的语句照样在跑,占着连接池
里的一个连接直到跑完(Oracle 实测多占了 60 秒);aiosqlite 更糟,语句在线程里
跑,取消之后关连接要排在它后面,一条不收敛的查询能让调用永远回不来。
会话级的设置用完要恢复:连接会回到池子里,探查结构这类长操作也用它。
下发失败(老版本的 MySQL 没有这个变量)不挡查询,只剩后端这边的兜底。
"""
if kind in ("postgres", "postgresql"):
# SET LOCAL 只管这一个事务:连接还回池子时回滚,设置随之作废
await conn.execute(text(f"SET LOCAL statement_timeout = {deadline.remaining_ms()}"))
yield
return
if kind in ("mysql", "mariadb"):
# MySQL 的 max_execution_time 以毫秒计、只管 SELECT;MariaDB 没有它,
# 用 max_statement_time(秒)。按连上的服务器认,不按表单里选的类型
if getattr(conn.dialect, "is_mariadb", False):
var, value = "max_statement_time", _seconds(deadline.remaining_seconds())
else:
var, value = "max_execution_time", str(deadline.remaining_ms())
try:
await conn.execute(text(f"SET SESSION {var} = {value}"))
except Exception as e: # noqa: BLE001
logger.warning("MySQL 语句时限没下发成功(%s),只剩后端兜底", e)
yield
return
cancelled = False
try:
yield
except asyncio.CancelledError:
# 连接还卡在那条语句上,这时再发一条只会排在它后面
cancelled = True
raise
finally:
if not cancelled:
try:
await conn.execute(text(f"SET SESSION {var} = DEFAULT"))
except Exception: # noqa: BLE001 - 连接已经坏了的话,连接池会把它作废
pass
return
if kind == "oracle":
driver = await _driver(conn)
before = getattr(driver, "call_timeout", 0)
try:
driver.call_timeout = deadline.remaining_ms()
except Exception as e: # noqa: BLE001
logger.warning("Oracle call_timeout 没设上(%s),只剩后端兜底", e)
yield
return
try:
yield
finally:
try:
driver.call_timeout = before
except Exception: # noqa: BLE001 - 超时后连接可能已经被驱动关掉
pass
return
if kind == "sqlite":
driver = await _driver(conn)
await driver.set_progress_handler(deadline.sqlite_progress, _SQLITE_STEPS)
try:
yield
finally:
deadline.over = True
try:
await driver.set_progress_handler(None, 0)
except Exception: # noqa: BLE001 - over 已经让它不再掐人
pass
return
yield
async def _bounded(
opener: Callable[[], Any], kind: str, limits: QueryLimits, over: str,
work: Callable[[AsyncConnection], Awaitable[Any]],
) -> Any:
"""在 opener() 开的连接上跑 work,时限交给数据库执行,后端这边再留一道兜底。
数据库按时限停下语句时报的错,翻成 over 这句话:「超过 Ns 被中断」。
两道时限从同一刻(取连接之前)算起,数据库那道先到,兜底晚 BACKSTOP_S。
"""
deadline: _Deadline | None = None
async def _exec() -> Any:
nonlocal deadline
deadline = _Deadline(limits.timeout_seconds)
async with opener() as conn:
# 光等连接就把时限用完了:语句不必再发,发出去也只剩 1ms
if deadline.expired():
raise SqlRejected(over)
async with _server_deadline(conn, kind, deadline):
return await work(conn)
try:
return await asyncio.wait_for(_exec(), timeout=limits.timeout_seconds + BACKSTOP_S)
except asyncio.TimeoutError as e:
raise SqlRejected(over) from e
except Exception as e:
if deadline is not None and deadline.stopped(e):
raise SqlRejected(over) from e
raise
async def run_query(
source: Any, sql: str, *, limits: QueryLimits | None = None
) -> QueryResult:
"""执行一条查询。守卫先过,再谈执行。
注意执行的是 guard.check 的返回值而不是原始输入——守卫会把尾分号之类
规范掉,执行原文等于绕过了守卫。
时限没指定时按数据源自己的配置(query_timeout),由数据库执行(_server_deadline)。
"""
limits = limits or QueryLimits(timeout_seconds=query_timeout(source))
statement = check(sql, readonly=bool(source.readonly), source_name=source.name)
kind = (source.kind or "").lower()
engine = await engines.get(source)
started = time.perf_counter()
if not source.readonly and is_write(statement):
return await _run_write(engine, kind, statement, limits, started)
async def _read(conn: AsyncConnection) -> tuple[list[str], list[list[Any]], bool, dict[str, str]]:
cursor = await conn.stream(text(statement))
columns = list(cursor.keys())
rows: list[list[Any]] = []
seen = _TypeTally(len(columns))
truncated = False
size = 0
async for row in cursor:
# 逐行累加,边收边判上限:一次性 fetchall 一个亿级表就晚了
seen.add(row)
values = [_jsonable(v) for v in row]
rows.append(values)
size += len(json.dumps(values, ensure_ascii=False, default=str))
if len(rows) >= limits.max_rows or size >= limits.max_bytes:
truncated = True
break
return columns, rows, truncated, seen.types(columns)
columns, rows, truncated, types = await _bounded(
engine.connect, kind, limits,
f"查询超过 {_seconds(limits.timeout_seconds)}s 被中断。加上 WHERE 条件或 LIMIT 缩小范围。",
_read,
)
return QueryResult(
columns=columns,
rows=rows,
row_count=len(rows),
truncated=truncated,
elapsed_ms=int((time.perf_counter() - started) * 1000),
sql=statement,
column_types=types,
)
class _TypeTally:
"""边读边记每列见过的值类型。只看驱动交回的原始值:_jsonable 之后 Decimal 就成了字符串,
小数位为 0 的 DECIMAL("45678")和文本列的 "2026" 再也分不开。"""
def __init__(self, width: int) -> None:
self._kinds: list[set[str]] = [set() for _ in range(width)]
def add(self, row: Any) -> None:
for kinds, value in zip(self._kinds, row):
if value is not None:
kinds.add(_value_kind(value))
def types(self, columns: list[str]) -> dict[str, str]:
return {str(name): kind for name, kinds in zip(columns, self._kinds)
if len(kinds) == 1 and (kind := next(iter(kinds))) != "other"}
def _value_kind(value: Any) -> str:
from datetime import date, datetime, time as dtime
from decimal import Decimal
if isinstance(value, bool):
return "boolean"
if isinstance(value, (int, float, Decimal)):
return "number"
if isinstance(value, datetime):
return "datetime"
if isinstance(value, date):
return "date"
if isinstance(value, dtime):
return "time"
return "text" if isinstance(value, str) else "other"
def column_types(columns: list[str], rows: list[Any]) -> dict[str, str]:
"""一批原始行里每列的类型:number / text / date / datetime / time / boolean。
全是空值、类型混着(同一列里有数也有文本)的列不记——拿不准就不说,下游照旧按值猜。
"""
tally = _TypeTally(len(columns))
for row in rows:
tally.add(row)
return tally.types(columns)
async def _run_write(
engine: AsyncEngine, kind: str, statement: str, limits: QueryLimits, started: float
) -> QueryResult:
"""可写源上的写操作:真的提交。
以前写和查询走同一条 stream 路径:拿不到结果集就抛 ResourceClosedError,
连接关闭时再回滚——可写源上的写从来没落过库,调用方只看到一个报错。
能走到这里的写都已经过了审批(registry.call_is_dangerous → 各节点的审批关卡),
只读源上的写在 check 里就被拒了。
"""
async def _write(conn: AsyncConnection) -> tuple[list[str], list[list[Any]]]:
result = await conn.execute(text(statement))
if result.returns_rows: # RETURNING、PG 的数据修改 CTE
return (list(result.keys()),
[[_jsonable(v) for v in row] for row in result.fetchmany(limits.max_rows)])
return ["affected_rows"], [[result.rowcount]]
# engine.begin():正常退出提交,出错(包括数据库按时限停下)回滚
columns, rows = await _bounded(
engine.begin, kind, limits,
f"写操作超过 {_seconds(limits.timeout_seconds)}s 被中断,已回滚。", _write,
)
return QueryResult(
columns=columns,
rows=rows,
row_count=len(rows),
truncated=len(rows) >= limits.max_rows,
elapsed_ms=int((time.perf_counter() - started) * 1000),
sql=statement,
)
def _jsonable(value: Any) -> Any:
"""把驱动返回的类型转成能进 JSON 的形态。
Decimal、date、datetime、bytes 这些直接 json.dumps 会炸,而它们在真实业务
数据里遍地都是——金额是 Decimal、时间是 datetime,不处理等于不可用。
"""
from datetime import date, datetime, time as dtime
from decimal import Decimal
if value is None or isinstance(value, (str, int, float, bool)):
return value
if isinstance(value, Decimal):
# 转 float 会丢精度,而金额恰恰不能丢。字符串保真,下游要算再自己转。
return str(value)
if isinstance(value, (datetime, date, dtime)):
return value.isoformat()
if isinstance(value, (bytes, bytearray, memoryview)):
raw = bytes(value)
try:
return raw.decode()
except UnicodeDecodeError:
return f"<binary {len(raw)} bytes>"
return str(value)