mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
数据目录 F3 一致性:查询结果的截断判定抽成 data/engine.py 的 RowCollector,查询层和合并查询都调它,不再各写一遍
This commit is contained in:
+31
-13
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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="秒"):
|
||||
|
||||
Reference in New Issue
Block a user