数据目录 F3 一致性:查询结果的截断判定抽成 data/engine.py 的 RowCollector,查询层和合并查询都调它,不再各写一遍

This commit is contained in:
ThearchyHelios
2026-10-04 01:08:42 +08:00
parent c12f0f2852
commit 4c7cfa62b2
3 changed files with 70 additions and 26 deletions
+31 -13
View File
@@ -879,6 +879,33 @@ async def _bounded(
raise
class RowCollector:
"""逐行收结果、边收边判上限(行数 max_rows、字节 max_bytes)。截断的判定只有这一份:查询层(run_query)和
合并查询(engine/merge_query._read)都调它,不各写一遍。
到了上限之后确实还有下一行才记 truncated。以前一到上限就标截断,恰好 1000 行的完整结果也算截断——截断的结果
整组拿去算的指标要降档出具、截断的合并结果下游会拒收,误报就是把完整的结果降了档、拒了。字节数按收下的行累加,
第一行总是收下:单行就超过字节上限也交回这一行。
"""
def __init__(self, limits: QueryLimits) -> None:
self.limits = limits
self.rows: list[list[Any]] = []
self.size = 0
self.truncated = False
self._full = False
def take(self, values: list[Any]) -> bool:
"""收一行(已转成能进 JSON 的值)。已经取满还来一行:记截断、不收,返回 False,调用方就此停止读取。"""
if self._full:
self.truncated = True
return False
self.rows.append(values)
self.size += len(json.dumps(values, ensure_ascii=False, default=str))
self._full = len(self.rows) >= self.limits.max_rows or self.size >= self.limits.max_bytes
return True
async def run_query(
source: Any, sql: str, *, limits: QueryLimits | None = None
) -> QueryResult:
@@ -903,23 +930,14 @@ async def run_query(
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]] = []
collected = RowCollector(limits)
seen = _TypeTally(len(columns))
truncated = full = False
size = 0
# 逐行累加,边收边判上限:一次性 fetchall 一个亿级表就晚了
async for row in cursor:
if full:
# 到了上限之后还有下一行:确实没取完。以前一到上限就标截断,恰好 1000 行的完整结果
# 也算截断——截断的结果整组拿去算的指标要降档出具,误报就是把完整的结果降了档
truncated = True
if not collected.take([_jsonable(v) for v in row]):
break
# 逐行累加,边收边判上限:一次性 fetchall 一个亿级表就晚了
seen.add(row)
values = [_jsonable(v) for v in row]
rows.append(values)
size += len(json.dumps(values, ensure_ascii=False, default=str))
full = len(rows) >= limits.max_rows or size >= limits.max_bytes
return columns, rows, truncated, seen.types(columns)
return columns, collected.rows, collected.truncated, seen.types(columns)
columns, rows, truncated, types = await _bounded(
engine.connect, kind, limits,
+6 -13
View File
@@ -40,7 +40,7 @@ from dataclasses import dataclass, field
from typing import Any
from app.data import guard
from app.data.engine import _jsonable, column_types, open_memory_sqlite
from app.data.engine import RowCollector, _jsonable, column_types, open_memory_sqlite
from app.data.guard import QueryLimits, SqlRejected
from app.data.names import name_key
from app.engine.direct_select import _kw, _lex, _Lex
@@ -343,24 +343,17 @@ def _authorize(action: int, *_args: Any) -> int:
def _read(conn: sqlite3.Connection, stmt: str, limits: QueryLimits) -> tuple[list[str], list[list[Any]], bool]:
"""逐行读、边读边判上限:和查询层(data.engine.run_query)同一个口径——到了行数或字节上限之后
确实还有下一行才记 truncated,恰好取满不算(截断的合并结果下游会拒收,误报就是拒了完整的结果)。"""
"""逐行读、边读边判上限。判定调查询层的 RowCollector(data/engine.py),和 run_query 是同一份:到了行数或
字节上限之后确实还有下一行才记 truncated,恰好取满不算(截断的合并结果下游会拒收,误报就是拒了完整的结果)。"""
cursor = conn.execute(stmt)
columns = [str(d[0]) for d in cursor.description or []]
if not columns:
raise MergeError("合并 SQL 没有返回任何列")
rows: list[list[Any]] = []
size = 0
truncated = full = False
collected = RowCollector(limits)
for row in cursor:
if full:
truncated = True
if not collected.take([_jsonable(v) for v in row]):
break
values = [_jsonable(v) for v in row]
rows.append(values)
size += len(json.dumps(values, ensure_ascii=False, default=str))
full = len(rows) >= limits.max_rows or size >= limits.max_bytes
return columns, rows, truncated
return columns, collected.rows, collected.truncated
def _sqlite_reason(e: sqlite3.Error, limits: QueryLimits) -> str:
+33
View File
@@ -225,6 +225,39 @@ def test_byte_limit_matches_the_query_layer():
assert out.truncated is True and 0 < len(out.rows) < 16
def test_row_collector_is_the_one_cutoff_rule():
"""截断判定只有一份(data.engine.RowCollector):到了行数或字节上限之后确实还有下一行才算截断,恰好取满不算;
第一行总是收下(单行就超过字节上限也得交回点什么)。"""
from app.data.engine import RowCollector
full = RowCollector(QueryLimits(max_rows=2))
assert full.take([1]) and full.take([2])
assert full.rows == [[1], [2]] and full.truncated is False # 恰好取满、后面没有了
assert full.take([3]) is False # 取满之后还有一行:截断,这一行不收
assert full.rows == [[1], [2]] and full.truncated is True
tiny = RowCollector(QueryLimits(max_bytes=5))
assert tiny.take(["一行就超过字节上限"]) is True and tiny.take(["下一行"]) is False
assert len(tiny.rows) == 1 and tiny.truncated is True
def test_merge_reads_through_the_query_layer_cutoff(monkeypatch):
"""合并查询调用查询层的判定,不另写一遍:两边的口径不会再各改各的。"""
import app.engine.merge_query as merge_query
from app.data.engine import RowCollector
made: list[RowCollector] = []
class Spy(RowCollector):
def __init__(self, limits):
super().__init__(limits)
made.append(self)
monkeypatch.setattr(merge_query, "RowCollector", Spy)
out = execute([sales(), visits()], "SELECT * FROM s CROSS JOIN v", limits=QueryLimits(max_rows=5))
assert len(made) == 1 and made[0].truncated is True and len(out.rows) == len(made[0].rows) == 5
def test_runaway_query_is_stopped():
endless = "WITH RECURSIVE c(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM c) SELECT count(*) FROM c"
with pytest.raises(MergeError, match="秒"):