DQ0-1170: Add new datasource - ElasticSearch

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
+```

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',
+]

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)

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,

requirements-elastic_search.txt

@@ -0,0 +1 @@
+elasticsearch==8.2.0

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

setup.py

@@ -51,13 +51,19 @@
 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:
     SNOWFLAKE = [line for line in f.read().splitlines() if line and not line.startswith('#')]
 
 EXTRA_REQUIRE = {
     "big_query" : BIG_QUERY,
+    "elastic_search" : ELASTIC_SEARCH,
     "snowflake" : SNOWFLAKE,
 }