elasticsearch - allow passing headers (#1156)

deepset-ai · Oct 28, 2024 · 2602bc1 · 2602bc1
1 parent 5034f96
commit 2602bc1
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/...s/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py b/...s/elasticsearch/src/haystack_integrations/document_stores/elasticsearch/document_store.py
@@ -105,9 +105,12 @@ def __init__(
     @property
     def client(self) -> Elasticsearch:
         if self._client is None:
+            headers = self._kwargs.pop("headers", {})
+            headers["user-agent"] = f"haystack-py-ds/{haystack_version}"
+
             client = Elasticsearch(
                 self._hosts,
-                headers={"user-agent": f"haystack-py-ds/{haystack_version}"},
+                headers=headers,
                 **self._kwargs,
             )
             # Check client connection, this will raise if not connected

diff --git a/integrations/elasticsearch/tests/test_document_store.py b/integrations/elasticsearch/tests/test_document_store.py
@@ -22,6 +22,20 @@ def test_init_is_lazy(_mock_es_client):
     _mock_es_client.assert_not_called()
 
 
+@patch("haystack_integrations.document_stores.elasticsearch.document_store.Elasticsearch")
+def test_headers_are_supported(_mock_es_client):
+    _ = ElasticsearchDocumentStore(hosts="testhost", headers={"header1": "value1", "header2": "value2"}).client
+
+    assert _mock_es_client.call_count == 1
+    _, kwargs = _mock_es_client.call_args
+
+    headers_found = kwargs["headers"]
+    assert headers_found["header1"] == "value1"
+    assert headers_found["header2"] == "value2"
+
+    assert headers_found["user-agent"].startswith("haystack-py-ds/")
+
+
 @patch("haystack_integrations.document_stores.elasticsearch.document_store.Elasticsearch")
 def test_to_dict(_mock_elasticsearch_client):
     document_store = ElasticsearchDocumentStore(hosts="some hosts")