diff --git a/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py b/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py index fdbf95eb0..1dea6e08b 100644 --- a/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py +++ b/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py @@ -85,6 +85,7 @@ def __init__( "Set the ASTRA_DB_API_ENDPOINT environment variable (recommended) or pass it explicitly." ) raise ValueError(msg) + self.resolved_api_endpoint = resolved_api_endpoint resolved_token = token.resolve_value() if resolved_token is None: @@ -93,6 +94,7 @@ def __init__( "Set the ASTRA_DB_APPLICATION_TOKEN environment variable (recommended) or pass it explicitly." ) raise ValueError(msg) + self.resolved_token = resolved_token self.api_endpoint = api_endpoint self.token = token @@ -101,15 +103,20 @@ def __init__( self.duplicates_policy = duplicates_policy self.similarity = similarity self.namespace = namespace - - self.index = AstraClient( - resolved_api_endpoint, - resolved_token, - self.collection_name, - self.embedding_dimension, - self.similarity, - namespace, - ) + self._index: Optional[AstraClient] = None + + @property + def index(self) -> AstraClient: + if self._index is None: + self._index = AstraClient( + self.resolved_api_endpoint, + self.resolved_token, + self.collection_name, + self.embedding_dimension, + self.similarity, + self.namespace, + ) + return self._index @classmethod def from_dict(cls, data: Dict[str, Any]) -> "AstraDocumentStore": diff --git a/integrations/astra/tests/test_document_store.py b/integrations/astra/tests/test_document_store.py index c475e6ae2..72945bcb8 100644 --- a/integrations/astra/tests/test_document_store.py +++ b/integrations/astra/tests/test_document_store.py @@ -20,13 +20,19 @@ def mock_auth(monkeypatch): monkeypatch.setenv("ASTRA_DB_APPLICATION_TOKEN", "test_token") +@mock.patch("haystack_integrations.document_stores.astra.astra_client.AstraDB") +def test_init_is_lazy(_mock_client): + _ = AstraDocumentStore() + _mock_client.assert_not_called() + + def test_namespace_init(mock_auth): # noqa with mock.patch("haystack_integrations.document_stores.astra.astra_client.AstraDB") as client: - AstraDocumentStore() + _ = AstraDocumentStore().index assert "namespace" in client.call_args.kwargs assert client.call_args.kwargs["namespace"] is None - AstraDocumentStore(namespace="foo") + _ = AstraDocumentStore(namespace="foo").index assert "namespace" in client.call_args.kwargs assert client.call_args.kwargs["namespace"] == "foo"