feat(vector-stores): add Databricks provider to TypeScript OSS SDK (#5824)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

tests/vector_stores/test_neptune_analytics.py

@@ -275,9 +275,116 @@ class TestNeptuneAnalyticsVectorInitValidation:
     def test_rejects_injection_payload_in_init(self, payload, monkeypatch):
         from mem0.vector_stores.neptune_analytics import NeptuneAnalyticsVector
         monkeypatch.setattr("mem0.vector_stores.neptune_analytics.NeptuneAnalyticsGraph", lambda *args, **kwargs: None)
-        
+
         with pytest.raises(ValueError, match="Invalid collection_name"):
             NeptuneAnalyticsVector(
                 endpoint="neptune-graph://test",
                 collection_name=payload
             )
+
+
+class _FakeNeptuneGraph:
+    """Minimal stand-in for `NeptuneAnalyticsGraph.query()` so update()'s compensation
+    path can be exercised without a real Neptune Analytics endpoint."""
+
+    def __init__(self):
+        self.nodes = {}
+        self.fail_next_upsert = False
+        self.soft_fail_next_upsert = False
+        self.get_call_count = 0
+
+    def query(self, query_string, params=None):
+        params = params or {}
+
+        if "UNWIND $rows" in query_string:
+            rows = params["rows"]
+            if "CALL neptune.algo.vectors.upsert" in query_string:
+                return [{"success": True} for _ in rows]
+            for row in rows:
+                self.nodes[row["node_id"]] = dict(row["properties"])
+            return []
+
+        if "CALL neptune.algo.vectors.upsert" in query_string:
+            if self.fail_next_upsert:
+                self.fail_next_upsert = False
+                raise RuntimeError("simulated Neptune upsert failure")
+            if self.soft_fail_next_upsert:
+                self.soft_fail_next_upsert = False
+                return [{"success": False}]
+            return [{"success": True}]
+
+        if "SET n = $properties" in query_string:
+            self.nodes[params["vector_id"]] = dict(params["properties"])
+            return []
+
+        if "RETURN n" in query_string and "node_id" in params:
+            self.get_call_count += 1
+            vector_id = params["node_id"]
+            if vector_id not in self.nodes:
+                return []
+            return [{"n": {"~id": vector_id, "~properties": dict(self.nodes[vector_id])}}]
+
+        if "DETACH DELETE n" in query_string:
+            self.nodes.pop(params.get("node_id"), None)
+            return []
+
+        return []
+
+
+class TestNeptuneAnalyticsUpdateRollback:
+    """update() must not leave a payload committed against a stale embedding when the
+    vector upsert step fails. See the compensation logic in `NeptuneAnalyticsVector.update()`."""
+
+    def _make_vec(self, monkeypatch):
+        monkeypatch.setattr("mem0.vector_stores.neptune_analytics.NeptuneAnalyticsGraph", lambda *args, **kwargs: None)
+        vec = NeptuneAnalyticsVector(endpoint="neptune-graph://test", collection_name="rollback")
+        vec.graph = _FakeNeptuneGraph()
+        return vec
+
+    def test_restores_prior_payload_when_upsert_fails(self, monkeypatch):
+        vec = self._make_vec(monkeypatch)
+        vec.insert(vectors=[[0.1, 0.2]], ids=["A"], payloads=[{"data": "alpha", "user_id": "u1"}])
+
+        vec.graph.fail_next_upsert = True
+        with pytest.raises(RuntimeError):
+            vec.update("A", vector=[0.9, 0.9], payload={"data": "beta", "user_id": "u1"})
+
+        restored = vec.get("A")
+        assert restored.payload["data"] == "alpha"
+        assert restored.payload["user_id"] == "u1"
+
+    def test_does_not_snapshot_prior_state_for_a_vector_only_update(self, monkeypatch):
+        """Only a combined payload+vector update can desync -- a vector-only update has
+        nothing to roll back to, so it must skip the extra get() snapshot entirely."""
+        vec = self._make_vec(monkeypatch)
+        vec.insert(vectors=[[0.1, 0.2]], ids=["A"], payloads=[{"data": "alpha", "user_id": "u1"}])
+
+        vec.graph.fail_next_upsert = True
+        calls_before = vec.graph.get_call_count
+        with pytest.raises(RuntimeError):
+            vec.update("A", vector=[0.9, 0.9])
+
+        assert vec.graph.get_call_count == calls_before
+
+    def test_succeeds_normally_when_upsert_does_not_fail(self, monkeypatch):
+        vec = self._make_vec(monkeypatch)
+        vec.insert(vectors=[[0.1, 0.2]], ids=["A"], payloads=[{"data": "alpha", "user_id": "u1"}])
+
+        vec.update("A", vector=[0.9, 0.9], payload={"data": "beta", "user_id": "u1"})
+
+        updated = vec.get("A")
+        assert updated.payload["data"] == "beta"
+
+    def test_rolls_back_on_soft_upsert_failure(self, monkeypatch):
+        """A soft {"success": False} row desyncs the payload from the embedding just as much as a
+        thrown error, so update() must treat it as a failure and roll the payload back too."""
+        vec = self._make_vec(monkeypatch)
+        vec.insert(vectors=[[0.1, 0.2]], ids=["A"], payloads=[{"data": "alpha", "user_id": "u1"}])
+
+        vec.graph.soft_fail_next_upsert = True
+        with pytest.raises(RuntimeError):
+            vec.update("A", vector=[0.9, 0.9], payload={"data": "beta", "user_id": "u1"})
+
+        restored = vec.get("A")
+        assert restored.payload["data"] == "alpha"
+        assert restored.payload["user_id"] == "u1"