diff --git a/src/surreal_memory/storage/surrealdb/store.py b/src/surreal_memory/storage/surrealdb/store.py index d901e651..e7c0aea7 100755 --- a/src/surreal_memory/storage/surrealdb/store.py +++ b/src/surreal_memory/storage/surrealdb/store.py @@ -199,6 +199,20 @@ def _is_auth_error(exc: Exception) -> bool: # measured 2026-09-13 on a copy of a production brain (18 657 state rows). _DIRECT_STATE_FETCH_LIMIT = 64 +# The alphabet emitted by _to_surreal_id, in SurrealDB's ASCII record-name +# ordering. Keyset windows use adjacent leading-character ranges so a +# RecordIdScan has an upper bound instead of walking the table's full suffix. +_SAFE_RECORD_ID_CHARS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz" + + +def _next_surreal_id_prefix(character: str) -> str | None: + """Return the next sanitized record-name initial in lexical order.""" + index = _SAFE_RECORD_ID_CHARS.find(character) + if index < 0: + return None + next_index = index + 1 + return _SAFE_RECORD_ID_CHARS[next_index] if next_index < len(_SAFE_RECORD_ID_CHARS) else None + def _prefer_ws_transport(url: str) -> str: """Rewrite http(s):// URLs to ws(s):// for the SDK connection. @@ -1434,10 +1448,17 @@ async def find_neurons_after_id( brain_id = self._get_brain_id() conditions = [f"brain_id = {_brain_literal(brain_id)}"] params: dict[str, Any] = {} - - if cursor_id is not None: - params["cursor_id"] = _to_surreal_id(cursor_id) + range_cursor = _to_surreal_id(cursor_id) if cursor_id is not None else None + + if range_cursor is not None: + # The explicit exclusive predicate preserves the public keyset + # contract. The lower-bounded record range lets SurrealDB 3.2 + # seek to the cursor instead of scanning the brain index and + # filtering every earlier id. _to_surreal_id is the single + # choke-point for the inlined record name; quoting keeps numeric- + # only and underscore-prefixed IDs as string record keys. conditions.append("id > type::record('neuron', $cursor_id)") + params["cursor_id"] = range_cursor if created_before is not None: params["created_before"] = created_before conditions.append("(created_at IS NONE OR created_at <= $created_before)") @@ -1447,11 +1468,26 @@ async def find_neurons_after_id( page_limit = min(max(int(limit), 1), 2000) projection = "SELECT *" if include_embedding else "SELECT * OMIT embedding_vec" - rows = await self._query( - f"{projection} FROM neuron WHERE {' AND '.join(conditions)} " - f"ORDER BY id ASC LIMIT {page_limit}", - **params, - ) + rows: list[dict[str, Any]] = [] + next_prefix = _next_surreal_id_prefix(range_cursor[0]) if range_cursor else None + range_start = range_cursor or None + while True: + remaining = page_limit - len(rows) + upper_clause = f"`{next_prefix}`" if next_prefix is not None else "" + if range_start is None: + from_clause = "neuron" + else: + from_clause = f"neuron:`{range_start}`..{upper_clause}" + batch = await self._query( + f"{projection} FROM {from_clause} WHERE {' AND '.join(conditions)} " + f"ORDER BY id ASC LIMIT {remaining}", + **params, + ) + rows.extend(batch) + if len(batch) >= remaining or next_prefix is None: + break + range_start = next_prefix + next_prefix = _next_surreal_id_prefix(next_prefix) return [_row_to_neuron(row) for row in rows] async def find_neurons_ranked( @@ -2204,10 +2240,13 @@ async def get_synapses_after_id( brain_id = self._get_brain_id() conditions = ["brain_id = $brain_id"] params: dict[str, Any] = {"brain_id": brain_id} + range_cursor = _to_surreal_id(cursor_id) if cursor_id is not None else None - if cursor_id is not None: - params["cursor_id"] = _to_surreal_id(cursor_id) + if range_cursor is not None: + # See find_neurons_after_id: pair the exclusive predicate with a + # quoted, sanitized lower record bound so the planner can seek. conditions.append("id > type::record('synapse', $cursor_id)") + params["cursor_id"] = range_cursor if created_before is not None: # The frozen run reference excludes records created after this scan began. # Preserve legacy rows without a timestamp; they cannot be classified as @@ -2216,11 +2255,26 @@ async def get_synapses_after_id( conditions.append("(created_at IS NONE OR created_at <= $created_before)") page_limit = min(max(int(limit), 1), 2000) - rows = await self._query( - f"SELECT * FROM synapse WHERE {' AND '.join(conditions)} " - f"ORDER BY id ASC LIMIT {page_limit}", - **params, - ) + rows: list[dict[str, Any]] = [] + next_prefix = _next_surreal_id_prefix(range_cursor[0]) if range_cursor else None + range_start = range_cursor or None + while True: + remaining = page_limit - len(rows) + upper_clause = f"`{next_prefix}`" if next_prefix is not None else "" + if range_start is None: + from_clause = "synapse" + else: + from_clause = f"synapse:`{range_start}`..{upper_clause}" + batch = await self._query( + f"SELECT * FROM {from_clause} WHERE {' AND '.join(conditions)} " + f"ORDER BY id ASC LIMIT {remaining}", + **params, + ) + rows.extend(batch) + if len(batch) >= remaining or next_prefix is None: + break + range_start = next_prefix + next_prefix = _next_surreal_id_prefix(next_prefix) return [_row_to_synapse(row) for row in rows] async def get_synapse_prune_page( diff --git a/tests/integration/test_surrealdb_semantic_keyset_ranges.py b/tests/integration/test_surrealdb_semantic_keyset_ranges.py new file mode 100644 index 00000000..f5d06f3e --- /dev/null +++ b/tests/integration/test_surrealdb_semantic_keyset_ranges.py @@ -0,0 +1,224 @@ +"""Disposable SurrealDB tests for bounded semantic keyset ranges.""" + +from __future__ import annotations + +import ipaddress +import json +import os +import uuid +from datetime import timedelta +from urllib.parse import urlparse + +import pytest +import pytest_asyncio + +from surreal_memory.core.brain import Brain +from surreal_memory.core.neuron import Neuron, NeuronType +from surreal_memory.core.synapse import Synapse, SynapseType +from surreal_memory.storage.surrealdb._ids import _record_id_part, _to_public_id +from surreal_memory.storage.surrealdb.store import SurrealDBStorage +from surreal_memory.utils.timeutils import utcnow + +SURREALDB_URL = os.getenv("SURREALDB_URL") +SURREALDB_USER = os.getenv("SURREALDB_USER", "root") +SURREALDB_PASS = os.getenv("SURREALDB_PASS", "root") +SURREALDB_NS = os.getenv("SURREALDB_NS", "smem_keyset_it") + + +def _is_loopback_url(url: str | None) -> bool: + if not url: + return False + try: + hostname = urlparse(url).hostname + return hostname == "localhost" or bool( + hostname and ipaddress.ip_address(hostname).is_loopback + ) + except ValueError: + return False + + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + not _is_loopback_url(SURREALDB_URL), + reason="requires explicit loopback SURREALDB_URL for disposable SurrealDB 3.2.4", + ), +] + + +@pytest_asyncio.fixture +async def store(): + storage = SurrealDBStorage( + url=SURREALDB_URL, + user=SURREALDB_USER, + password=SURREALDB_PASS, + namespace=SURREALDB_NS, + database="it_" + uuid.uuid4().hex[:12], + ) + await storage.initialize() + brain = Brain.create(name="semantic-keyset-range-it") + await storage.save_brain(brain) + storage.set_brain(brain.id) + try: + yield storage + finally: + await storage.close() + + +def _captured_queries(store: SurrealDBStorage) -> list[tuple[str, dict[str, object]]]: + queries: list[tuple[str, dict[str, object]]] = [] + original_query = store._query + + async def capture(sql: str, **params): + queries.append((sql, params)) + return await original_query(sql, **params) + + store._query = capture # type: ignore[method-assign] + return queries + + +def _public_record_ids(rows: list[dict[str, object]]) -> list[str]: + return [_to_public_id(_record_id_part(str(row["id"]))) for row in rows] + + +async def _assert_record_scan( + store: SurrealDBStorage, query: tuple[str, dict[str, object]] +) -> None: + sql, params = query + plan = await SurrealDBStorage._query_response(store, sql + " EXPLAIN FULL", **params) + plan_text = json.dumps(plan, ensure_ascii=False, default=str) + assert "RecordIdScan" in plan_text, plan_text[:800] + assert ".." in plan_text, plan_text[:800] + + +@pytest.mark.asyncio +async def test_neuron_keyset_range_fills_page_and_matches_brain_scoped_reference_query( + store: SurrealDBStorage, +) -> None: + reference_time = utcnow() + timedelta(days=1) + future_time = reference_time + timedelta(days=1) + cursor = "1234567890000000" + brain_a_id = store._get_brain_id() + brain_a_ids = [cursor, "1234567890000001", "1234567890000002", "1234567890000003"] + for neuron_id in brain_a_ids: + await store.add_neuron( + Neuron( + id=neuron_id, + type=NeuronType.CONCEPT, + content=neuron_id, + created_at=future_time if neuron_id.endswith("2") else reference_time, + ) + ) + + other_brain = Brain.create(name="semantic-keyset-other-brain") + await store.save_brain(other_brain) + store.set_brain(other_brain.id) + await store.add_neuron( + Neuron.create(type=NeuronType.CONCEPT, content="foreign", neuron_id="1234567890000004") + ) + store.set_brain(brain_a_id) + queries = _captured_queries(store) + actual = await store.find_neurons_after_id(cursor, limit=2, created_before=reference_time) + sql, params = queries[0] + reference_rows = await SurrealDBStorage._query_response( + store, + "SELECT * FROM neuron WHERE brain_id = $brain_id AND id > type::record('neuron', $cursor_id) " + "AND (created_at IS NONE OR created_at <= $created_before) AND ephemeral = false " + "ORDER BY id ASC LIMIT 2", + brain_id=store._get_brain_id(), + cursor_id=cursor, + created_before=reference_time, + ) + + assert [row.id for row in actual] == _public_record_ids(reference_rows) + assert [row.id for row in actual] == ["1234567890000001", "1234567890000003"] + assert "FROM neuron:`1234567890000000`..`2`" in sql + assert params["cursor_id"] == cursor + await _assert_record_scan(store, queries[0]) + + # Underscore-prefixed record names remain quoted and cursor-exclusive. + for neuron_id in ("_123", "_124", "_125"): + await store.add_neuron( + Neuron.create(type=NeuronType.CONCEPT, content=neuron_id, neuron_id=neuron_id) + ) + underscore_page = await store.find_neurons_after_id( + "_123", limit=2, created_before=reference_time + ) + assert [row.id for row in underscore_page] == ["-124", "-125"] + + for neuron_id in ("7123", "7999", "8001", "8002"): + await store.add_neuron( + Neuron.create(type=NeuronType.CONCEPT, content=neuron_id, neuron_id=neuron_id) + ) + fill_offset = len(queries) + filled_page = await store.find_neurons_after_id("7123", limit=3, created_before=reference_time) + fill_queries = queries[fill_offset:] + assert [row.id for row in filled_page] == ["7999", "8001", "8002"] + assert len(fill_queries) == 2 + assert "FROM neuron:`7123`..`8`" in fill_queries[0][0] + assert "FROM neuron:`8`..`9`" in fill_queries[1][0] + + +@pytest.mark.asyncio +async def test_synapse_keyset_range_fills_page_and_keeps_other_brains_out( + store: SurrealDBStorage, +) -> None: + reference_time = utcnow() + timedelta(days=1) + cursor = "1234567891000000" + for synapse_id in (cursor, "1234567891000001", "1234567891000002", "1234567891000003"): + await store.add_synapse( + Synapse.create( + source_id="source", + target_id="target", + type=SynapseType.RELATED_TO, + synapse_id=synapse_id, + ) + ) + brain_id = store._get_brain_id() + other_brain = Brain.create(name="semantic-synapse-keyset-other") + await store.save_brain(other_brain) + store.set_brain(other_brain.id) + await store.add_synapse( + Synapse.create( + source_id="source", + target_id="target", + type=SynapseType.RELATED_TO, + synapse_id="1234567891000004", + ) + ) + store.set_brain(brain_id) + + queries = _captured_queries(store) + actual = await store.get_synapses_after_id(cursor, limit=2, created_before=reference_time) + sql, params = queries[0] + reference_rows = await SurrealDBStorage._query_response( + store, + "SELECT * FROM synapse WHERE brain_id = $brain_id AND id > type::record('synapse', $cursor_id) " + "AND (created_at IS NONE OR created_at <= $created_before) ORDER BY id ASC LIMIT 2", + brain_id=brain_id, + cursor_id=cursor, + created_before=reference_time, + ) + + assert [row.id for row in actual] == _public_record_ids(reference_rows) + assert [row.id for row in actual] == ["1234567891000001", "1234567891000002"] + assert "FROM synapse:`1234567891000000`..`2`" in sql + assert params["cursor_id"] == cursor + await _assert_record_scan(store, queries[0]) + + for synapse_id in ("7123", "7999", "8001", "8002"): + await store.add_synapse( + Synapse.create( + source_id="source", + target_id="target", + type=SynapseType.RELATED_TO, + synapse_id=synapse_id, + ) + ) + fill_offset = len(queries) + filled_page = await store.get_synapses_after_id("7123", limit=3, created_before=reference_time) + fill_queries = queries[fill_offset:] + assert [row.id for row in filled_page] == ["7999", "8001", "8002"] + assert len(fill_queries) == 2 + assert "FROM synapse:`7123`..`8`" in fill_queries[0][0] + assert "FROM synapse:`8`..`9`" in fill_queries[1][0] diff --git a/tests/unit/test_surrealdb_synapse_queries.py b/tests/unit/test_surrealdb_synapse_queries.py index 82bedc72..13fcb3f1 100644 --- a/tests/unit/test_surrealdb_synapse_queries.py +++ b/tests/unit/test_surrealdb_synapse_queries.py @@ -295,13 +295,27 @@ async def test_synapse_keyset_page_binds_frozen_reference_and_cursor(self): ) assert page == [] - sql = st._query.await_args.args[0] + sql = st._query.await_args_list[0].args[0] params = st._query.await_args.kwargs assert "id > type::record('synapse', $cursor_id)" in sql + assert "FROM synapse:`edge_10`..`f`" in sql assert "created_at IS NONE OR created_at <= $created_before" in sql assert "ORDER BY id ASC LIMIT 2000" in sql assert params["cursor_id"] == "edge_10" + @pytest.mark.asyncio + async def test_synapse_keyset_page_quotes_numeric_cursor_in_bounded_record_range(self): + st, _ = _store_with_mock_conn() + st._query = AsyncMock(return_value=[]) # type: ignore[method-assign] + + await st.get_synapses_after_id("1122334455667788", limit=25) + + sql = st._query.await_args_list[0].args[0] + params = st._query.await_args.kwargs + assert "FROM synapse:`1122334455667788`..`2`" in sql + assert "id > type::record('synapse', $cursor_id)" in sql + assert params["cursor_id"] == "1122334455667788" + @pytest.mark.asyncio async def test_prune_page_uses_projected_tuple_cursor_without_upper_bound(self): st, _ = _store_with_mock_conn() @@ -337,7 +351,7 @@ async def test_prune_first_page_uses_creation_time_order(self): await st.get_synapse_prune_page(None, None, limit=100) - sql = st._query.await_args.args[0] + sql = st._query.await_args_list[0].args[0] assert "brain_id = $brain_id" in sql assert "ORDER BY created_at ASC, id ASC LIMIT 100" in sql assert "SELECT * FROM synapse" not in sql @@ -355,9 +369,9 @@ async def test_neuron_keyset_page_omits_embeddings_and_is_brain_scoped(self): ) assert page == [] - sql = st._query.await_args.args[0] + sql = st._query.await_args_list[0].args[0] params = st._query.await_args.kwargs - assert "SELECT * OMIT embedding_vec FROM neuron" in sql + assert "SELECT * OMIT embedding_vec FROM neuron:`node_10`..`o`" in sql assert 'brain_id = "b1"' in sql assert "id > type::record('neuron', $cursor_id)" in sql assert "created_at IS NONE OR created_at <= $created_before" in sql @@ -365,6 +379,23 @@ async def test_neuron_keyset_page_omits_embeddings_and_is_brain_scoped(self): assert params["cursor_id"] == "node_10" assert params["ephemeral"] is False + @pytest.mark.asyncio + async def test_neuron_keyset_page_quotes_underscore_cursor_and_keeps_none_unbounded(self): + st, _ = _store_with_mock_conn() + st._query = AsyncMock(return_value=[]) # type: ignore[method-assign] + + await st.find_neurons_after_id("_123", limit=25) + + sql = st._query.await_args_list[0].args[0] + assert "FROM neuron:`_123`..`a`" in sql + assert "id > type::record('neuron', $cursor_id)" in sql + + st._query = AsyncMock(return_value=[]) # type: ignore[method-assign] + await st.find_neurons_after_id(None, limit=25) + sql = st._query.await_args.args[0] + assert "FROM neuron WHERE" in sql + assert "id > type::record('neuron'" not in sql + @pytest.mark.asyncio async def test_get_connected_neuron_ids_for_only_queries_supplied_endpoints(self): st, _ = _store_with_mock_conn()