From a962de90e4b3c3183a8913aa9400af05b3d37593 Mon Sep 17 00:00:00 2001 From: antejavor Date: Fri, 17 Jul 2026 09:28:01 +0200 Subject: [PATCH 1/2] Init fix for roundtrip. --- .../src/lightrag_memgraph/docstatus_impl.py | 48 +++++++++++-------- .../src/lightrag_memgraph/kv_impl.py | 32 +++++++++---- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/integrations/lightrag-memgraph/src/lightrag_memgraph/docstatus_impl.py b/integrations/lightrag-memgraph/src/lightrag_memgraph/docstatus_impl.py index c570e2b..c40d12c 100644 --- a/integrations/lightrag-memgraph/src/lightrag_memgraph/docstatus_impl.py +++ b/integrations/lightrag-memgraph/src/lightrag_memgraph/docstatus_impl.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio import json import os from dataclasses import dataclass @@ -276,26 +277,35 @@ async def get_docs_paginated( params.update(filter_params) driver = await get_driver() - async with driver.session(database=get_database(), default_access_mode="READ") as session: - count_result = await session.run( - f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total", - **filter_params, - ) - count_record = await count_result.single() - await count_result.consume() - total = count_record["total"] if count_record else 0 - page_result = await session.run( - f""" - MATCH (n:`{self._label}`) {where} - RETURN n.id AS id, n.data AS data - ORDER BY {sort_expr} {order} - SKIP $skip LIMIT $limit - """, - **params, - ) - records = [record async for record in page_result] - await page_result.consume() + async def _count() -> int: + async with driver.session(database=get_database(), default_access_mode="READ") as session: + result = await session.run( + f"MATCH (n:`{self._label}`) {where} RETURN count(n) AS total", + **filter_params, + ) + record = await result.single() + await result.consume() + return record["total"] if record else 0 + + async def _page() -> list: + async with driver.session(database=get_database(), default_access_mode="READ") as session: + result = await session.run( + f""" + MATCH (n:`{self._label}`) {where} + RETURN n.id AS id, n.data AS data + ORDER BY {sort_expr} {order} + SKIP $skip LIMIT $limit + """, + **params, + ) + records = [record async for record in result] + await result.consume() + return records + + # No data dependency between the count and the page -- run them concurrently + # instead of two sequential round trips. + total, records = await asyncio.gather(_count(), _page()) docs: list[tuple[str, DocProcessingStatus]] = [] for record in records: diff --git a/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py b/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py index 0caf31a..5b0206b 100644 --- a/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py +++ b/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py @@ -121,20 +121,31 @@ async def upsert(self, data: dict[str, dict[str, Any]]) -> None: if not data: return current_time = int(time.time()) - # create_time only set for genuinely new records (matches JsonKVStorage semantics). - existing = set(data.keys()) - await self.filter_keys(set(data.keys())) + # Precompute both the "new record" and "existing record" JSON payloads per entry; + # MERGE picks the right one via ON CREATE/ON MATCH in a single round trip instead of + # a separate filter_keys() query to decide in Python (which was also a TOCTOU race). entries = [] for k, v in data.items(): value = dict(v) if self.namespace.endswith("text_chunks") and "llm_cache_list" not in value: value["llm_cache_list"] = [] - if k in existing: - value["update_time"] = current_time - else: - value["create_time"] = current_time - value["update_time"] = current_time value["_id"] = k - entries.append({"id": k, "data": json.dumps(value, ensure_ascii=False), "ts": current_time}) + + new_value = dict(value) + new_value["create_time"] = current_time + new_value["update_time"] = current_time + + update_value = dict(value) + update_value["update_time"] = current_time + + entries.append( + { + "id": k, + "new_data": json.dumps(new_value, ensure_ascii=False), + "update_data": json.dumps(update_value, ensure_ascii=False), + "ts": current_time, + } + ) driver = await get_driver() async with driver.session(database=get_database()) as session: @@ -143,8 +154,9 @@ async def upsert(self, data: dict[str, dict[str, Any]]) -> None: f""" UNWIND $entries AS e MERGE (n:`{self._label}` {{id: e.id}}) - ON CREATE SET n.created_at = e.ts - SET n.data = e.data, n.updated_at = e.ts + ON CREATE SET n.data = e.new_data, n.created_at = e.ts + ON MATCH SET n.data = e.update_data + SET n.updated_at = e.ts """, entries=entries, ) From 35d8d5bea690a4814f1bedf3823f1bd10784f882 Mon Sep 17 00:00:00 2001 From: antejavor Date: Fri, 17 Jul 2026 13:06:54 +0200 Subject: [PATCH 2/2] Update comment. --- .../lightrag-memgraph/src/lightrag_memgraph/kv_impl.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py b/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py index 5b0206b..606f451 100644 --- a/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py +++ b/integrations/lightrag-memgraph/src/lightrag_memgraph/kv_impl.py @@ -121,9 +121,7 @@ async def upsert(self, data: dict[str, dict[str, Any]]) -> None: if not data: return current_time = int(time.time()) - # Precompute both the "new record" and "existing record" JSON payloads per entry; - # MERGE picks the right one via ON CREATE/ON MATCH in a single round trip instead of - # a separate filter_keys() query to decide in Python (which was also a TOCTOU race). + # MERGE picks new vs. existing payload via ON CREATE/ON MATCH -- no separate filter_keys() round trip. entries = [] for k, v in data.items(): value = dict(v)