diff --git a/es_index.py b/es_index.py index 8cd9db6..1b26dfe 100644 --- a/es_index.py +++ b/es_index.py @@ -15,21 +15,16 @@ class ESIndex(PyOperator): Use bulk to insert docs into ElasticSearch index, using auto id generated. Args: - host (`str`): host to connect ElasticSearch client - port (`int`): port to connect ElasticSearch client - index_name (`str`): index name to index input docs - user (`str=None`): user name to connect ElasticSearch client, defaults to None - password (`str=None`): user password to connect ElasticSearch client, defaults to None - ca_certs (`str=None`): path to CA certificate, defaults to None + **connection_kwargs: For example, host """ - def __init__(self, host: str, port: int, index_name: str, user: str = None, password: str = None, ca_certs: str = None): + def __init__(self, index_name: str, **connection_kwargs): super().__init__() self.index_name = index_name + if 'port' in connection_kwargs: + assert 'host' in connection_kwargs, 'Missing port in connection kwargs but given port only.' + connection_kwargs['hosts'] = [f'https://{host}:{port}'] try: - self.client = Elasticsearch( - f'https://{host}:{port}', - ca_certs=ca_certs, - basic_auth=(user, password)) + self.client = Elasticsearch(**connection_kwargs) logger.info('Successfully connected to ElasticSearch client.') except Exception as e: logger.error('Failed to connect ElasticSearch client:\n', e)