diff --git a/backend/app/data/engine.py b/backend/app/data/engine.py index b419f2a1..0f4cbf3d 100644 --- a/backend/app/data/engine.py +++ b/backend/app/data/engine.py @@ -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, diff --git a/backend/app/engine/merge_query.py b/backend/app/engine/merge_query.py index 52a54d59..87d5d6c7 100644 --- a/backend/app/engine/merge_query.py +++ b/backend/app/engine/merge_query.py @@ -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: diff --git a/backend/tests/test_merge_engine.py b/backend/tests/test_merge_engine.py index 9acf20c9..65306842 100644 --- a/backend/tests/test_merge_engine.py +++ b/backend/tests/test_merge_engine.py @@ -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="秒"):