yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2026-01-11 06:36:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: fix method type and add comments (#5235)

Browse Source 
* fix method type and add comments

* fix tests
This commit is contained in:
Massimiliano Pippi 2023-06-30 11:55:52 +02:00 committed by GitHub
parent 037e4f24ce
commit cb638af0ff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
haystack/document_stores/elasticsearch/es7.py
View File

@ -125,10 +125,12 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
memory issues, decrease the batch_size.
"""
# Ensure all the required inputs were successful
es_import.check()
# Base constructor might need the client to be ready, create it first
client = self._init_elastic_client(
# Let the base class trap the right exception from the specific client
self._RequestError = RequestError
# Initiate the Elasticsearch client for version 7.x
client = ElasticsearchDocumentStore._init_elastic_client(
host=host,
port=port,
username=username,
@ -168,9 +170,6 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
batch_size=batch_size,
)
# Let the base class trap the right exception from the elasticpy client
self._RequestError = RequestError
def _do_bulk(self, *args, **kwargs):
"""Override the base class method to use the Elasticsearch client"""
return bulk(*args, **kwargs)
@ -179,9 +178,8 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
"""Override the base class method to use the Elasticsearch client"""
return scan(*args, **kwargs)
@classmethod
@staticmethod
def _init_elastic_client(
cls,
host: Union[str, List[str]],
port: Union[int, List[int]],
username: str,

14
haystack/document_stores/elasticsearch/es8.py
View File

@ -141,10 +141,12 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
memory issues, decrease the batch_size.
"""
# Ensure all the required inputs were successful
es_import.check()
# Base constructor might need the client to be ready, create it first
client = self._init_elastic_client(
# Let the base class trap the right exception from the specific client
self._RequestError = RequestError
# Initiate the Elasticsearch client for version 8.x
client = ElasticsearchDocumentStore._init_elastic_client(
host=host,
port=port,
username=username,
@ -183,9 +185,6 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
batch_size=batch_size,
)
# Let the base class trap the right exception from the elasticpy client
self._RequestError = RequestError
def _do_bulk(self, *args, **kwargs):
"""Override the base class method to use the Elasticsearch client"""
return bulk(*args, **kwargs)
@ -194,9 +193,8 @@ class ElasticsearchDocumentStore(_ElasticsearchDocumentStore):
"""Override the base class method to use the Elasticsearch client"""
return scan(*args, **kwargs)
@classmethod
@staticmethod
def _init_elastic_client(
cls,
host: Union[str, List[str]],
port: Union[int, List[int]],
username: str,

13
test/document_stores/test_elasticsearch.py
View File

@ -56,13 +56,14 @@ class TestElasticsearchDocumentStore(DocumentStoreBaseTestAbstract, SearchEngine
ElasticsearchDocumentStore equipped with a mocked client
"""
class DSMock(ElasticsearchDocumentStore):
# We mock a subclass to avoid messing up the actual class object
pass
with patch(f"{ElasticsearchDocumentStore.__module__}.ElasticsearchDocumentStore._init_elastic_client"):
DSMock._init_elastic_client = MagicMock()
DSMock.client = MagicMock()
return DSMock()
class DSMock(ElasticsearchDocumentStore):
# We mock a subclass to avoid messing up the actual class object
pass
DSMock.client = MagicMock()
yield DSMock()
@pytest.mark.integration
def test___init__(self):