Source code for now.data_loading.elasticsearch.connector

import logging
from typing import Dict, Generator, List, Optional

from elasticsearch import Elasticsearch


[docs]class ElasticsearchConnector: def __init__( self, connection_str: str = 'http://localhost:9200', connection_args: Optional[Dict] = None, ): """ Provides an interface to an Elasticsearch database. :param connection_str: A connection string for the ES instance. Usually, it includes url, port, username, password, etc. Typically, it has the form: 'https://{user_name}:{password}@{host}:{port}' :param connection_args: Dictionary with additional connection arguments, e.g., information about certificates """ self._connection_str = connection_str self._connection_args = ( connection_args if connection_args else {'verify_certs': False} ) = Elasticsearch(self._connection_str, **self._connection_args) def __enter__(self) -> 'ElasticsearchConnector': return self def __exit__(self, type, value, traceback) -> None: self.close()
[docs] def get_documents_by_query( self, query: Dict, index_name: str, page_size: Optional[int] = 10 ) -> Generator[List[Dict], None, None]: """ Executes an Elasticsearch query on a given index and returns a generator which yields pages of documents from the query results. :param query: Elasticsearch query :param index_name: Name of an Elasticsearch index :param page_size: Number of documents per page :return: Generator which yields one page of documents on each call. """ resp = **query, index=index_name, scroll='2m', size=page_size, source=False, ) documents = [ {**doc['_source'], **{'id': doc['_id']}} for doc in resp['hits']['hits'] ] scroll_id = resp['_scroll_id'] scroll_size = len(documents) while scroll_size > 0: yield documents resp =, scroll='2m') scroll_id = resp['_scroll_id'] documents = [ {**doc['_source'], **{'id': doc['_id']}} for doc in resp['hits']['hits'] ] scroll_size = len(documents)
[docs] def close(self) -> None: """ Closes Elasticsearch connection. """