feat: arcadedb - add closing methods (#3653)

Author: anakin87

integrations/arcadedb/src/haystack_integrations/components/retrievers/arcadedb/embedding_retriever.py

@@ -122,6 +122,12 @@ def to_dict(self) -> dict[str, Any]:
             filter_policy=self._filter_policy.value,
         )
 
+    def close(self) -> None:
+        """
+        Release the synchronous resources of the underlying Document Store.
+        """
+        self._document_store.close()
+
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "ArcadeDBEmbeddingRetriever":
         """

integrations/arcadedb/src/haystack_integrations/document_stores/arcadedb/document_store.py

@@ -5,6 +5,7 @@
 """ArcadeDB DocumentStore for Haystack 2.x — document storage + vector search via HTTP/JSON API."""
 
 import logging
+from contextlib import suppress
 from http import HTTPStatus
 from typing import Any, ClassVar
 
@@ -95,7 +96,7 @@ def __init__(
         self._recreate_type = recreate_type
         self._create_database = create_database
 
-        self._session = requests.Session()
+        self._session: requests.Session | None = None
         self._initialized = False
 
     def to_dict(self) -> dict[str, Any]:
@@ -134,10 +135,25 @@ def from_dict(cls, data: dict[str, Any]) -> "ArcadeDBDocumentStore":
                 init_params[key] = Secret.from_dict(init_params[key])
         return default_from_dict(cls, data)
 
+    def close(self) -> None:
+        """
+        Release the associated synchronous resources.
+        """
+        if self._session is not None:
+            with suppress(Exception):
+                self._session.close()
+            self._session = None
+
     # ------------------------------------------------------------------
     # HTTP helpers
     # ------------------------------------------------------------------
 
+    def _ensure_session(self) -> requests.Session:
+        """Lazily create the HTTP session"""
+        if self._session is None:
+            self._session = requests.Session()
+        return self._session
+
     def _auth(self) -> tuple[str, str] | None:
         user = self._username.resolve_value() if self._username else None
         pwd = self._password.resolve_value() if self._password else None
@@ -152,7 +168,7 @@ def _command(self, sql: str, *, positional_params: list[Any] | None = None) -> l
         if positional_params:
             payload["params"] = positional_params
 
-        resp = self._session.post(url, json=payload, auth=self._auth())
+        resp = self._ensure_session().post(url, json=payload, auth=self._auth())
         if resp.status_code >= HTTPStatus.BAD_REQUEST:
             msg = f"ArcadeDB command failed ({resp.status_code}): {resp.text}"
             raise RuntimeError(msg)
@@ -163,7 +179,7 @@ def _command(self, sql: str, *, positional_params: list[Any] | None = None) -> l
     def _server_command(self, command: str) -> dict[str, Any]:
         """Execute a server-level command (e.g. CREATE DATABASE)."""
         url = f"{self._url}/api/v1/server"
-        resp = self._session.post(url, json={"command": command}, auth=self._auth())
+        resp = self._ensure_session().post(url, json={"command": command}, auth=self._auth())
         if resp.status_code >= HTTPStatus.BAD_REQUEST:
             msg = f"ArcadeDB server command failed ({resp.status_code}): {resp.text}"
             raise RuntimeError(msg)

integrations/arcadedb/tests/test_document_store.py

@@ -5,7 +5,7 @@
 import dataclasses
 import datetime
 import os
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 from haystack import Document
@@ -338,6 +338,20 @@ def test_embedding_retrieval_without_filter_returns_all(self, store):
         docs = store._embedding_retrieval([0.0] * 4)
         assert [d.id for d in docs] == ["a"]
 
+    def test_close(self, store):
+        store._session = Mock()
+        session = store._session
+        store.close()
+        session.close.assert_called_once_with()
+        assert store._session is None
+        store.close()
+
+    def test_close_is_exception_safe(self, store):
+        store._session = Mock()
+        store._session.close.side_effect = RuntimeError
+        store.close()
+        assert store._session is None
+
 
 @pytest.mark.skipif(
     not os.environ.get("ARCADEDB_PASSWORD"),
@@ -388,6 +402,15 @@ def assert_documents_are_equal(self, received: list[Document], expected: list[Do
             expected_clean = dataclasses.replace(expected_doc, embedding=None)
             assert actual == expected_clean
 
+    def test_close_and_reopen(self, document_store: ArcadeDBDocumentStore):
+        document_store.write_documents([Document(id="1")])
+        assert document_store.count_documents() == 1
+
+        document_store.close()
+        assert document_store._session is None
+
+        assert document_store.count_documents() == 1
+
     def test_write_documents(self, document_store: ArcadeDBDocumentStore):
         """Override mixin: test default write_documents and duplicate fail behaviour."""
         docs = [Document(id="1")]

integrations/arcadedb/tests/test_embedding_retriever.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from unittest.mock import Mock
+
+from haystack.dataclasses import Document
+from haystack.document_stores.types import FilterPolicy
+
+from haystack_integrations.components.retrievers.arcadedb import ArcadeDBEmbeddingRetriever
+from haystack_integrations.document_stores.arcadedb import ArcadeDBDocumentStore
+
+
+class TestEmbeddingRetriever:
+    def test_init_default(self):
+        mock_store = Mock(spec=ArcadeDBDocumentStore)
+        retriever = ArcadeDBEmbeddingRetriever(document_store=mock_store)
+        assert retriever._document_store == mock_store
+        assert retriever._filters is None
+        assert retriever._top_k == 10
+        assert retriever._filter_policy == FilterPolicy.REPLACE
+
+    def test_init(self):
+        mock_store = Mock(spec=ArcadeDBDocumentStore)
+        retriever = ArcadeDBEmbeddingRetriever(
+            document_store=mock_store, filters={"field": "value"}, top_k=5, filter_policy=FilterPolicy.MERGE
+        )
+        assert retriever._filters == {"field": "value"}
+        assert retriever._top_k == 5
+        assert retriever._filter_policy == FilterPolicy.MERGE
+
+    def test_to_dict_from_dict(self):
+        store = ArcadeDBDocumentStore(url="http://localhost:2480", database="test", create_database=False)
+        retriever = ArcadeDBEmbeddingRetriever(document_store=store, filters={"field": "value"}, top_k=5)
+
+        restored = ArcadeDBEmbeddingRetriever.from_dict(retriever.to_dict())
+
+        assert isinstance(restored._document_store, ArcadeDBDocumentStore)
+        assert restored._document_store._database == "test"
+        assert restored._filters == {"field": "value"}
+        assert restored._top_k == 5
+        assert restored._filter_policy == FilterPolicy.REPLACE
+
+    def test_close(self):
+        mock_store = Mock(spec=ArcadeDBDocumentStore)
+        retriever = ArcadeDBEmbeddingRetriever(document_store=mock_store)
+
+        retriever.close()
+
+        mock_store.close.assert_called_once()
+        assert retriever._document_store is mock_store
+
+    def test_run(self):
+        mock_store = Mock(spec=ArcadeDBDocumentStore)
+        doc = Document(content="Test doc", embedding=[0.1, 0.2])
+        mock_store._embedding_retrieval.return_value = [doc]
+
+        retriever = ArcadeDBEmbeddingRetriever(document_store=mock_store)
+        res = retriever.run(query_embedding=[0.3, 0.5])
+
+        mock_store._embedding_retrieval.assert_called_once_with(query_embedding=[0.3, 0.5], filters=None, top_k=10)
+        assert res == {"documents": [doc]}