From a9ce95dbf6f65068d1fa42d781710ac43fc66026 Mon Sep 17 00:00:00 2001 From: Artur Susdorf Date: Wed, 11 May 2022 16:05:33 +0200 Subject: [PATCH] added ElasticSearch datasource class --- DEVELOPMENT.md | 14 +++ dq0/sdk/data/clients/__init__.py | 11 ++ dq0/sdk/data/clients/elastic_search.py | 146 +++++++++++++++++++++++++ dq0/sdk/data/data_source_factory.py | 3 +- requirements-elastic_search.txt | 1 + requirements.txt | 2 +- setup.py | 6 + 7 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 dq0/sdk/data/clients/__init__.py create mode 100644 dq0/sdk/data/clients/elastic_search.py create mode 100644 requirements-elastic_search.txt diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 4c8ccd96..1cc0fb87 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -18,6 +18,11 @@ pip install -r requirements.txt pip install -r requirements-dev.txt ``` +Additional dependencies can be installed, too: +```bash +pip install -r requirements-elastic_search.txt +``` + Helper to resolve dependency conflicts: ```bash # pip install pipdeptree @@ -53,6 +58,7 @@ python -m pytest -c /dev/null tests/test_dummies.py::test_slow Use autopep8 to auto format code: ```bash +# pip install autopep8==1.5.5 python -m autopep8 --in-place --aggressive --aggressive -r dq0 ``` @@ -94,3 +100,11 @@ Production / binary installation TBD python setup.py sdist bdist_wheel ``` it will create at least two files in "dist". + +## Helpers + +Remove pycache folders and generated .pyc files +```bash +# on BSD (macOS & Linux) +find . -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete +``` \ No newline at end of file diff --git a/dq0/sdk/data/clients/__init__.py b/dq0/sdk/data/clients/__init__.py new file mode 100644 index 00000000..8c67f31f --- /dev/null +++ b/dq0/sdk/data/clients/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +"""DQ0 SDK Data Sources Client package. + +This package contains all client based data source implementation. +""" + +from .elastic_search import ElasticSearch + +__all__ = [ + 'ElasticSearch', +] diff --git a/dq0/sdk/data/clients/elastic_search.py b/dq0/sdk/data/clients/elastic_search.py new file mode 100644 index 00000000..d35c40e1 --- /dev/null +++ b/dq0/sdk/data/clients/elastic_search.py @@ -0,0 +1,146 @@ +# -*- coding: utf-8 -*- +"""Data Source for ElasticSearch clients. + +This source class provides access to data from an ElasticSearch instance and exposes it as pandas data frames. + +Copyright 2020, Gradient Zero +All rights reserved +""" + +from dq0.sdk.data.source import Source + +try: + from elasticsearch import Elasticsearch + elastic_search_available = True +except ImportError: + elastic_search_available = False + +import pandas as pd +from pandas.io.json import json_normalize + + +class ElasticSearch(Source): + """Data Source for ElasticSearch data. + + Provides function to read data from ElasticSearch. + + ElasticSearch connection string: 'https://[username]:[password]@[host]:[port]/' + + Usage: + from dq0.sdk.data.clients.elastic_search import ElasticSearch + es = ElasticSearch(username='elastic', password='password') + df = es.search(index='my_index', query={"match_all":{}}) + + Args: + protocol (:obj:`str`): Protocol of the ElasticSearch instance (default: 'https') + host (:obj:`str`): Host of the ElasticSearch instance (default: 'localhost') + username (:obj:`str`): Username used to connect to the ElasticSearch instance + password (:obj:`str`): Password used to connect to the ElasticSearch instance + port (:obj:`int`): Port of the running ElasticSearch instance (default: 9200) + """ + + def __init__( + self, + protocol='https', + host='localhost', + username='', + password='', + port=9200, + **kwargs + ): + + password_sep = ':' if 0 < len(password) else '' + user_sep = '@' if 0 < len(username) else '' + port_sep = ':' if port is not None else '' + connection_string = f"{protocol}://{username}{password_sep}{password}{user_sep}{host}{port_sep}{port}/" + self.connection_string = connection_string + super().__init__(connection_string, **kwargs) + self.type = 'elasticsearch' + + # Construct a Elasticsearch client object. + if not elastic_search_available: + raise ImportError('elasticsearch dependencies must be installed first') + + self.client = Elasticsearch( + hosts=self.connection_string, + **kwargs + ) + + # ping instance + if self.client.ping() is not True: + # use self.client.info() ? + print(self.client.info()) + raise ImportError('elasticsearch ping failed') + + def search(self, index=None, query=None, **kwargs): + """Search ElasticSearch using index and query + + Args: + index: Index + query: Query + + Returns: + ElasticSearch search result as pandas dataframe + """ + # checks + if index is None: + raise ValueError('you need to pass a index parameter') + if query is None: + raise ValueError('you need to pass a query parameter') + + # make an API request + resp = self.client.search(index=index, query=query, **kwargs) + + # sample result format: + ''' + { + "took" : 28, + "timed_out" : false, + "_shards" : { + "total" : 5, + "successful" : 5, + "skipped" : 0, + "failed" : 0 + }, + "hits" : { + "total" : 1, + "max_score" : 0.8630463, + "hits" : [ + { + "_index" : "logs", + "_type" : "my_app", + "_id" : "ZsWdJ2EBir6MIbMWSMyF", + "_score" : 0.8630463, + "_source" : { + "timestamp" : "2018-01-24 12:34:56", + "message" : "User logged in", + "user_id" : 4, + "admin" : false + } + } + ] + } + } + ''' + # results are stored in _source inside hits.hits + df = None + if 'hits' in resp: + if 'hits' in resp['hits']: + df = json_normalize(resp['hits']['hits']) + + # return pandas dataframe + return df + + def read(self, **kwargs): + """Read json data sources + + Args: + kwargs: keyword arguments + + Returns: + json data as pandas dataframe + + Raises: + IOError: if file was not found + """ + return pd.read_json(self.path, **kwargs) diff --git a/dq0/sdk/data/data_source_factory.py b/dq0/sdk/data/data_source_factory.py index f950bb37..a60ac5b3 100644 --- a/dq0/sdk/data/data_source_factory.py +++ b/dq0/sdk/data/data_source_factory.py @@ -6,7 +6,7 @@ Copyright 2020, Gradient Zero All rights reserved """ -from dq0.sdk.data import binary, image, sql, text +from dq0.sdk.data import binary, clients, image, sql, text data_source_classes = { @@ -19,6 +19,7 @@ 'sas': binary.sas.SAS, 'spss': binary.spss.SPSS, 'stata': binary.stata.Stata, + 'elasticsearch': clients.elastic_search.ElasticSearch, 'image': image.Image, 'bigquery': sql.big_query.BigQuery, 'drill': sql.drill.Drill, diff --git a/requirements-elastic_search.txt b/requirements-elastic_search.txt new file mode 100644 index 00000000..46d6fd9f --- /dev/null +++ b/requirements-elastic_search.txt @@ -0,0 +1 @@ +elasticsearch==8.2.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d9208fb2..1b02f939 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -requests==2.23.0 +requests==2.27.1 absl-py==0.11.0 pandas==1.2.2 scipy==1.6.0 diff --git a/setup.py b/setup.py index 48383c74..d433d779 100644 --- a/setup.py +++ b/setup.py @@ -51,6 +51,11 @@ with open('requirements-big_query.txt') as f: BIG_QUERY = [line for line in f.read().splitlines() if line and not line.startswith('#')] +# Extras requirements (elastic_search) +ELASTIC_SEARCH = [''] +with open('requirements-elastic_search.txt') as f: + ELASTIC_SEARCH = [line for line in f.read().splitlines() if line and not line.startswith('#')] + # Extras requirements (snowflake) SNOWFLAKE = [''] with open('requirements-snowflake.txt') as f: @@ -58,6 +63,7 @@ EXTRA_REQUIRE = { "big_query" : BIG_QUERY, + "elastic_search" : ELASTIC_SEARCH, "snowflake" : SNOWFLAKE, }