diff --git a/mempalace/hallways.py b/mempalace/hallways.py index 907527a..9bee6c6 100644 --- a/mempalace/hallways.py +++ b/mempalace/hallways.py @@ -179,7 +179,12 @@ def compute_hallways_for_wing( Args: wing: wing name to scan. - col: ChromaDB collection — must support ``.get(where=..., include=...)``. + col: ChromaDB collection — must support ``.count()`` and paginated + ``.get(limit=..., offset=..., include=...)``. The fetch is filtered + to ``wing`` client-side rather than via ``.get(where={"wing": ...})``, + which binds one SQL variable per matched id and overflows SQLite's + ``SQLITE_MAX_VARIABLE_NUMBER`` on large wings (#1619). Fake + collections and alternate backends must implement this shape. If ``None``, returns ``[]`` (caller didn't supply a backing store, so nothing to compute against). Tests pass a controlled MagicMock. @@ -198,16 +203,32 @@ def compute_hallways_for_wing( min_count = max(1, int(min_count)) - # 1. Query drawers for this wing. + # 1. Query drawers for this wing. Paginate the fetch and filter to the + # wing client-side: a single get(where={"wing": wing}) binds one SQL + # variable per matched id and overflows SQLite's + # SQLITE_MAX_VARIABLE_NUMBER (32766) on wings > ~32k drawers (#1619). + # Mirrors the established pagination in miner.status / palace / + # palace_graph. + metadatas: list = [] try: - results = col.get(where={"wing": wing}, include=["metadatas"]) + total = col.count() + batch_size = 5000 + offset = 0 + while offset < total: + batch = col.get(limit=batch_size, offset=offset, include=["metadatas"]) + batch_metas = (batch or {}).get("metadatas") or [] + if not batch_metas: + break + metadatas.extend( + m for m in batch_metas if isinstance(m, dict) and m.get("wing") == wing + ) + offset += len(batch_metas) except Exception: logger.warning( - "compute_hallways_for_wing: collection.get failed for %s", wing, exc_info=True + "compute_hallways_for_wing: collection fetch failed for %s", wing, exc_info=True ) return [] - metadatas = (results or {}).get("metadatas") or [] if not metadatas: return [] diff --git a/tests/test_hallways.py b/tests/test_hallways.py index 94b7ed2..92ba186 100644 --- a/tests/test_hallways.py +++ b/tests/test_hallways.py @@ -27,11 +27,21 @@ def _use_tmp_hallway_file(monkeypatch, tmp_path): def _fake_collection(drawers): - """Build a MagicMock collection whose .get() returns the given drawer set.""" + """Build a MagicMock collection over ``drawers`` that supports the paginated + fetch (``count()`` + ``get(limit=, offset=)``) that compute_hallways_for_wing + uses to stay under SQLite's variable limit (#1619).""" col = MagicMock() - metadatas = [d for d in drawers] - ids = [f"drawer_{i}" for i in range(len(drawers))] - col.get.return_value = {"ids": ids, "metadatas": metadatas} + metas = [d for d in drawers] + col.count.return_value = len(metas) + + def _get(limit=None, offset=0, include=None, where=None, ids=None, **kwargs): + page = metas[offset : offset + limit] if limit is not None else metas + return { + "ids": [f"drawer_{i}" for i in range(offset, offset + len(page))], + "metadatas": page, + } + + col.get.side_effect = _get return col diff --git a/tests/test_hallways_pagination.py b/tests/test_hallways_pagination.py new file mode 100644 index 0000000..8847071 --- /dev/null +++ b/tests/test_hallways_pagination.py @@ -0,0 +1,56 @@ +"""Regression test for #1619. + +``compute_hallways_for_wing`` must fetch drawers by paginating +(``count()`` + ``get(limit=, offset=)``) and filtering the wing client-side, +NOT with a single ``get(where={"wing": wing})`` — the latter binds one SQL +variable per matched id and overflows SQLite's ``SQLITE_MAX_VARIABLE_NUMBER`` +(32766) on wings larger than ~32k drawers, silently leaving the hallway graph +unbuilt on exactly the large wings that benefit most. +""" + +from unittest.mock import MagicMock, patch + +with patch.dict("sys.modules", {"chromadb": MagicMock()}): + from mempalace import hallways as hallways_mod + + +def _use_tmp_hallway_file(monkeypatch, tmp_path): + monkeypatch.setattr(hallways_mod, "_HALLWAY_FILE", str(tmp_path / "hallways.json")) + + +def _collection_that_rejects_where_get(drawers): + """count() + paginated get(limit,offset) work; a where-get raises, exactly + as ChromaDB does when the bound-variable count overflows on a big wing.""" + col = MagicMock() + col.count.return_value = len(drawers) + + def _get(limit=None, offset=0, include=None, where=None, ids=None, **kw): + if where is not None and limit is None: + raise RuntimeError("Error executing plan: too many SQL variables") + filtered_drawers = drawers + if where and "wing" in where: + target_wing = where["wing"] + filtered_drawers = [ + d for d in drawers if isinstance(d, dict) and d.get("wing") == target_wing + ] + page = filtered_drawers[offset : offset + limit] if limit is not None else filtered_drawers + return { + "ids": [f"d{i}" for i in range(offset, offset + len(page))], + "metadatas": page, + } + + col.get.side_effect = _get + return col + + +class TestComputeHallwaysPagination: + def test_large_wing_builds_hallways_via_pagination(self, tmp_path, monkeypatch): + _use_tmp_hallway_file(monkeypatch, tmp_path) + # 3 drawers all co-placing Alice+Bob → one hallway at min_count=2, + # but ONLY if the fetch paginates instead of the variable-bound where-get. + drawers = [{"wing": "wing_alpha", "room": "diary", "entities": "Alice;Bob"}] * 3 + col = _collection_that_rejects_where_get(drawers) + result = hallways_mod.compute_hallways_for_wing("wing_alpha", col=col) + assert any({h["entity_a"], h["entity_b"]} == {"Alice", "Bob"} for h in result), ( + "hallways came back empty — the where-get path crashed; the fetch must paginate (#1619)" + )