memgraph-toolbox 0.2.1__tar.gz → 0.2.2__tar.gz

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memgraph-toolbox
-Version: 0.2.1
+Version: 0.2.2
 Summary: Memgraph toolbox library for Memgraph AI tools and utilities
 Project-URL: Homepage, https://github.com/memgraph/ai-toolkit
 Project-URL: Source, https://github.com/memgraph/ai-toolkit

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "memgraph-toolbox"
-version = "0.2.1"
+version = "0.2.2"
 description = "Memgraph toolbox library for Memgraph AI tools and utilities"
 readme = "README.md"
 authors = [

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/api/memgraph.py

@@ -1,4 +1,5 @@
 import os
+import time
 from collections.abc import Mapping
 from typing import Any, cast
 
@@ -78,6 +79,39 @@ def memgraph_env(
     }
 
 
+def merge_with_retry(db: Any, query: str, params: dict | None = None, *, attempts: int = 3) -> list[dict[str, Any]]:
+    """Run *query*, a MERGE of a uniqueness-constrained node, through *db*, surviving a concurrent create.
+
+    Two transactions that MERGE the same new node both find it missing and
+    both create it; the later commit fails the uniqueness constraint. Run
+    again, its MERGE matches the node the other one committed. Writers that
+    race this way are typically one process per event (agent hooks).
+
+    Args:
+        db: Anything with ``query(query, params)``, normally a :class:`Memgraph`.
+        query: The Cypher to run; it must be safe to run more than once.
+        params: Its parameters.
+        attempts: Runs before giving up, the first included.
+
+    Returns:
+        The rows of the run that committed.
+
+    Raises:
+        neo4j.exceptions.ClientError: a uniqueness violation on the last
+            attempt, or any other client error at once.
+    """
+    from neo4j.exceptions import ClientError
+
+    for attempt in range(attempts):
+        try:
+            return db.query(query, params=params)
+        except ClientError as exc:
+            if "unique constraint" not in str(exc) or attempt == attempts - 1:
+                raise
+            time.sleep(0.05 * (attempt + 1))
+    raise ValueError(f"attempts must be at least 1, got {attempts}")
+
+
 class Memgraph:
     """
     Base Memgraph client for interacting with Memgraph database.

memgraph_toolbox-0.2.2/src/memgraph_toolbox/tests/test_merge_with_retry.py

@@ -0,0 +1,69 @@
+"""Tests for merge_with_retry against a live Memgraph; they skip when none is reachable."""
+
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+from neo4j.exceptions import ClientError
+
+from ..api.memgraph import Memgraph, merge_with_retry
+
+_LABEL = "MergeWithRetryProbe"
+
+
+@pytest.fixture
+def db():
+    try:
+        client = Memgraph()
+    except Exception as exc:
+        pytest.skip(f"No Memgraph server reachable at the configured MEMGRAPH_URL: {exc}")
+    client.query(f"MATCH (n:{_LABEL}) DETACH DELETE n")
+    client.query(f"CREATE CONSTRAINT ON (n:{_LABEL}) ASSERT n.key IS UNIQUE;")
+    yield client
+    client.query(f"MATCH (n:{_LABEL}) DETACH DELETE n")
+    client.query(f"DROP CONSTRAINT ON (n:{_LABEL}) ASSERT n.key IS UNIQUE;")
+    client.close()
+
+
+def test_concurrent_merges_of_new_nodes_converge_on_one_node_each(db):
+    keys = [f"k{i}" for i in range(10)]
+
+    def merge(key: str) -> None:
+        merge_with_retry(db, f"MERGE (n:{_LABEL} {{key: $key}})", {"key": key})
+
+    with ThreadPoolExecutor(max_workers=16) as pool:
+        list(pool.map(merge, keys * 16))
+
+    rows = db.query(f"MATCH (n:{_LABEL}) RETURN n.key AS key, count(*) AS c ORDER BY key")
+    assert rows == [{"key": key, "c": 1} for key in sorted(keys)]
+
+
+def test_returns_the_rows_of_the_run_that_committed(db):
+    rows = merge_with_retry(db, f"MERGE (n:{_LABEL} {{key: 'a'}}) RETURN n.key AS key")
+    assert rows == [{"key": "a"}]
+
+
+def test_any_other_client_error_is_raised_without_retrying(db):
+    calls = []
+
+    class Counting:
+        def query(self, query, params=None):
+            calls.append(query)
+            return db.query(query, params)
+
+    with pytest.raises(ClientError):
+        merge_with_retry(Counting(), "MERGE (n:")
+    assert len(calls) == 1
+
+
+def test_a_conflict_on_every_attempt_is_raised_after_the_last():
+    """A conflict that never clears can't be produced on demand by a real server."""
+    calls = []
+
+    class AlwaysConflicting:
+        def query(self, query, params=None):
+            calls.append(query)
+            raise ClientError("Unable to commit due to unique constraint violation on :X(key)")
+
+    with pytest.raises(ClientError, match="unique constraint"):
+        merge_with_retry(AlwaysConflicting(), "MERGE (n:X {key: 1})", attempts=2)
+    assert len(calls) == 2