Elasticsearch: headers support for retrieving data (#1110)

* headers support for retrieving data

* modify docstring

* formatted connection error message

* lint fix

Author: kvignesh1420

tensorflow_io/core/kernels/elasticsearch_kernels.cc

@@ -33,11 +33,12 @@ class ElasticsearchReadableResource : public ResourceBase {
   Status Init(const std::string& healthcheck_url,
               const std::string& healthcheck_field,
               const std::string& request_url,
+              const std::vector<string>& headers,
               std::function<Status(const TensorShape& columns_shape,
                                    Tensor** columns, Tensor** dtypes)>
                   allocate_func) {
     // Perform healthcheck before proceeding
-    Healthcheck(healthcheck_url, healthcheck_field);
+    Healthcheck(healthcheck_url, healthcheck_field, headers);
 
     // Make the request API call and set the metadata based on a sample of
     // data returned. The request_url will have the "scroll" param set with
@@ -46,7 +47,7 @@ class ElasticsearchReadableResource : public ResourceBase {
     base_dtypes_.clear();
     base_columns_.clear();
     rapidjson::Document response_json;
-    MakeAPICall(request_url, &response_json);
+    MakeAPICall(request_url, &response_json, headers);
 
     // Validate the presence of the _scroll_id in the response.
     // The _scroll_id keeps might change in subsequent calls, thus not
@@ -121,9 +122,9 @@ class ElasticsearchReadableResource : public ResourceBase {
           data_allocate_func) {
     rapidjson::Document response_json;
     if (scroll_id == "") {
-      MakeAPICall(request_url, &response_json);
+      MakeAPICall(request_url, &response_json, headers_);
     } else {
-      MakeAPICall(scroll_request_url, &response_json);
+      MakeAPICall(scroll_request_url, &response_json, headers_);
     }
 
     if (response_json.HasMember("_scroll_id")) {
@@ -172,10 +173,11 @@ class ElasticsearchReadableResource : public ResourceBase {
 
  protected:
   Status Healthcheck(const std::string& healthcheck_url,
-                     const std::string& healthcheck_field) {
+                     const std::string& healthcheck_field,
+                     const std::vector<string>& headers) {
     // Make the healthcheck API call and get the response json
     rapidjson::Document response_json;
-    MakeAPICall(healthcheck_url, &response_json);
+    MakeAPICall(healthcheck_url, &response_json, headers);
 
     if (response_json.HasMember(healthcheck_field.c_str())) {
       // LOG(INFO) << "cluster health: "
@@ -186,8 +188,8 @@ class ElasticsearchReadableResource : public ResourceBase {
     return Status::OK();
   }
 
-  Status MakeAPICall(const std::string& url,
-                     rapidjson::Document* response_json) {
+  Status MakeAPICall(const std::string& url, rapidjson::Document* response_json,
+                     const std::vector<string>& headers) {
     HttpRequest* request = http_request_factory_.Create();
 
     if (scroll_id != "") {
@@ -200,7 +202,15 @@ class ElasticsearchReadableResource : public ResourceBase {
     }
 
     // LOG(INFO) << "Setting the headers";
-    request->AddHeader("Content-Type", "application/json; charset=utf-8");
+    for (size_t i = 0; i < headers.size(); ++i) {
+      std::string header = headers[i];
+      std::vector<string> parts = str_util::Split(header, "=");
+      if (parts.size() != 2) {
+        return errors::InvalidArgument("invalid header configuration: ",
+                                       header);
+      }
+      request->AddHeader(parts[0], parts[1]);
+    }
 
     // LOG(INFO) << "Setting the response buffer";
     std::vector<char> response;
@@ -231,6 +241,9 @@ class ElasticsearchReadableResource : public ResourceBase {
           "Invalid JSON response. The response should be an object");
     }
 
+    // Store the default headers if the response is valid
+    headers_ = headers;
+
     return Status::OK();
   }
 
@@ -242,6 +255,7 @@ class ElasticsearchReadableResource : public ResourceBase {
   std::vector<DataType> base_dtypes_;
   std::vector<string> base_columns_;
   std::string scroll_id = "";
+  std::vector<string> headers_;
 };
 
 class ElasticsearchReadableInitOp
@@ -271,17 +285,24 @@ class ElasticsearchReadableInitOp
     OP_REQUIRES_OK(context, context->input("request_url", &request_url_tensor));
     const string& request_url = request_url_tensor->scalar<tstring>()();
 
+    const Tensor* headers_tensor;
+    OP_REQUIRES_OK(context, context->input("headers", &headers_tensor));
+    std::vector<string> headers;
+    for (int64 i = 0; i < headers_tensor->NumElements(); i++) {
+      headers.push_back(headers_tensor->flat<tstring>()(i));
+    }
+
     OP_REQUIRES_OK(
-        context,
-        resource_->Init(healthcheck_url, healthcheck_field, request_url,
-                        [&](const TensorShape& columns_shape, Tensor** columns,
-                            Tensor** dtypes) -> Status {
-                          TF_RETURN_IF_ERROR(context->allocate_output(
-                              1, columns_shape, columns));
-                          TF_RETURN_IF_ERROR(context->allocate_output(
-                              2, columns_shape, dtypes));
-                          return Status::OK();
-                        }));
+        context, resource_->Init(
+                     healthcheck_url, healthcheck_field, request_url, headers,
+                     [&](const TensorShape& columns_shape, Tensor** columns,
+                         Tensor** dtypes) -> Status {
+                       TF_RETURN_IF_ERROR(
+                           context->allocate_output(1, columns_shape, columns));
+                       TF_RETURN_IF_ERROR(
+                           context->allocate_output(2, columns_shape, dtypes));
+                       return Status::OK();
+                     }));
   }
 
   Status CreateResource(ElasticsearchReadableResource** resource)

tensorflow_io/core/ops/elasticsearch_ops.cc

@@ -25,6 +25,7 @@ REGISTER_OP("IO>ElasticsearchReadableInit")
     .Input("healthcheck_url: string")
     .Input("healthcheck_field: string")
     .Input("request_url: string")
+    .Input("headers: string")
     .Output("resource: resource")
     .Output("columns: string")
     .Output("dtypes: string")

tensorflow_io/core/python/experimental/elasticsearch_dataset_ops.py

@@ -25,10 +25,11 @@ class _ElasticsearchHandler:
         session data.
     """
 
-    def __init__(self, nodes, index, doc_type):
+    def __init__(self, nodes, index, doc_type, headers_dict):
         self.nodes = nodes
         self.index = index
         self.doc_type = doc_type
+        self.headers_dict = headers_dict
         self.prepare_base_urls()
         self.prepare_connection_data()
 
@@ -57,8 +58,6 @@ def prepare_base_urls(self):
             base_url = "{}://{}".format(url_obj.scheme, url_obj.netloc)
             self.base_urls.append(base_url)
 
-        return self.base_urls
-
     def prepare_connection_data(self):
         """Prepares the healthcheck and resource urls from the base_urls"""
 
@@ -75,7 +74,17 @@ def prepare_connection_data(self):
                 )
             self.request_urls.append(request_url)
 
-        return self.healthcheck_urls, self.request_urls
+        self.headers = ["Content-Type=application/json"]
+        if self.headers_dict is not None:
+            if isinstance(self.headers_dict, dict):
+                for key, value in self.headers_dict.items():
+                    if key.lower() == "content-type":
+                        continue
+                    self.headers.append("{}={}".format(key, value))
+            else:
+                raise ValueError(
+                    "Headers should be a dict of key:value pairs. Got: ", self.headers
+                )
 
     def get_healthy_resource(self):
         """Retrieve the resource which is connected to a healthy node"""
@@ -88,6 +97,7 @@ def get_healthy_resource(self):
                     healthcheck_url=healthcheck_url,
                     healthcheck_field="status",
                     request_url=request_url,
+                    headers=self.headers,
                 )
                 print("Connection successful: {}".format(healthcheck_url))
                 dtypes = []
@@ -102,11 +112,13 @@ def get_healthy_resource(self):
                         dtypes.append(tf.string)
                 return resource, columns.numpy(), dtypes, request_url
             except Exception:
-                print("Skipping host: {}".format(healthcheck_url))
+                print("Skipping node: {}".format(healthcheck_url))
                 continue
         else:
             raise ConnectionError(
-                "No healthy node available for this index, check the cluster status and index"
+                "No healthy node available for the index: {}, please check the cluster config".format(
+                    self.index
+                )
             )
 
     def get_next_batch(self, resource, request_url):
@@ -155,21 +167,25 @@ def parse_json(self, raw_item, columns, dtypes):
 class ElasticsearchIODataset(tf.compat.v2.data.Dataset):
     """Represents an elasticsearch based tf.data.Dataset"""
 
-    def __init__(self, nodes, index, doc_type=None, internal=True):
+    def __init__(self, nodes, index, doc_type=None, headers=None, internal=True):
         """Prepare the ElasticsearchIODataset.
 
         Args:
             nodes: A `tf.string` tensor containing the hostnames of nodes
                 in [protocol://hostname:port] format.
                 For example: ["http://localhost:9200"]
             index: A `tf.string` representing the elasticsearch index to query.
-            doc_type: A `tf.string` representing the type of documents in the index
+            doc_type: (Optional) A `tf.string` representing the type of documents in the index
                 to query.
+            headers: (Optional) A dict of headers. For example:
+                {'Content-Type': 'application/json'}    
         """
         with tf.name_scope("ElasticsearchIODataset"):
             assert internal
 
-            handler = _ElasticsearchHandler(nodes=nodes, index=index, doc_type=doc_type)
+            handler = _ElasticsearchHandler(
+                nodes=nodes, index=index, doc_type=doc_type, headers_dict=headers
+            )
             resource, columns, dtypes, request_url = handler.get_healthy_resource()
 
             dataset = tf.data.experimental.Counter()

tests/test_elasticsearch/elasticsearch_test.sh

@@ -21,14 +21,20 @@ action=$1
 
 if [ "$action" == "start" ]; then
 
+echo ""
 echo "Preparing the environment variables file..."
+echo ""
 cat >> .env-vars << "EOF"
 cluster.name=tfio-es-cluster
 bootstrap.memory_lock=true
 discovery.type=single-node
+ELASTIC_PASSWORD=default_password
+xpack.security.enabled=true
 EOF
 
+echo ""
 echo "Starting the tfio elasticsearch docker container..."
+echo ""
 ELASTICSEARCH_IMAGE="docker.elastic.co/elasticsearch/elasticsearch:7.4.0"
 
 docker run -d --rm --name=tfio-elasticsearch \
@@ -37,22 +43,36 @@ docker run -d --rm --name=tfio-elasticsearch \
 --ulimit memlock=-1:-1 \
 ${ELASTICSEARCH_IMAGE}
 
+echo ""
 echo "Waiting for the elasticsearch cluster to be up and running..."
+echo ""
 sleep 20
 
+echo ""
 echo "Checking the base REST-API endpoint..."
-curl localhost:9200/
+echo ""
+# The Authorization header contains the base64 encoded value of "elastic:default_password"
+# As per the environment variable set while starting the container.
+curl -X GET localhost:9200/ --header 'Authorization: Basic ZWxhc3RpYzpkZWZhdWx0X3Bhc3N3b3Jk'
 
+echo ""
 echo "Checking the healthcheck REST-API endpoint..."
-curl localhost:9200/_cluster/health
+echo ""
+curl -X GET localhost:9200/_cluster/health --header 'Authorization: Basic ZWxhc3RpYzpkZWZhdWx0X3Bhc3N3b3Jk'
 
+echo ""
 echo "Clean up..."
+echo ""
 rm -rf ./.env-vars
 
 elif [ "$action" == "stop" ]; then
+echo ""
 echo "Removing the tfio elasticsearch container..."
+echo ""
 docker rm -f tfio-elasticsearch
 
 else
+echo ""
 echo "Invalid value: Use 'start' to run the container and 'stop' to remove it."
+echo ""
 fi

tests/test_elasticsearch_eager.py

@@ -29,7 +29,10 @@
 NODE = "http://localhost:9200"
 INDEX = "people"
 DOC_TYPE = "survivors"
-HEADERS = {"Content-Type": "application/json"}
+HEADERS = {
+    "Content-Type": "application/json",
+    "Authorization": "Basic ZWxhc3RpYzpkZWZhdWx0X3Bhc3N3b3Jk",
+}
 ATTRS = ["name", "gender", "age", "fare", "survived"]
 
 
@@ -51,7 +54,7 @@ def test_create_index():
     """Create an index in the cluster"""
 
     create_index_url = "{}/{}".format(NODE, INDEX)
-    res = requests.put(create_index_url)
+    res = requests.put(create_index_url, headers=HEADERS)
     assert res.status_code == 200
 
 
@@ -87,7 +90,7 @@ def test_elasticsearch_io_dataset():
     """Test the functionality of the ElasticsearchIODataset"""
 
     dataset = tfio.experimental.elasticsearch.ElasticsearchIODataset(
-        nodes=[NODE], index=INDEX, doc_type=DOC_TYPE
+        nodes=[NODE], index=INDEX, doc_type=DOC_TYPE, headers=HEADERS
     )
 
     assert issubclass(type(dataset), tf.data.Dataset)
@@ -97,13 +100,31 @@ def test_elasticsearch_io_dataset():
             assert attr in item
 
 
+@pytest.mark.skipif(not is_container_running(), reason="The container is not running")
+def test_elasticsearch_io_dataset_no_auth():
+    """Test the functionality of the ElasticsearchIODataset when basic auth is
+    required but the associated header is not passed.
+    """
+
+    try:
+        dataset = tfio.experimental.elasticsearch.ElasticsearchIODataset(
+            nodes=[NODE], index=INDEX, doc_type=DOC_TYPE
+        )
+    except ConnectionError as e:
+        assert str(
+            e
+        ) == "No healthy node available for the index: {}, please check the cluster config".format(
+            INDEX
+        )
+
+
 @pytest.mark.skipif(not is_container_running(), reason="The container is not running")
 def test_elasticsearch_io_dataset_batch():
     """Test the functionality of the ElasticsearchIODataset"""
 
     BATCH_SIZE = 2
     dataset = tfio.experimental.elasticsearch.ElasticsearchIODataset(
-        nodes=[NODE], index=INDEX, doc_type=DOC_TYPE
+        nodes=[NODE], index=INDEX, doc_type=DOC_TYPE, headers=HEADERS
     ).batch(BATCH_SIZE)
 
     assert issubclass(type(dataset), tf.data.Dataset)
@@ -119,5 +140,5 @@ def test_cleanup():
     """Clean up the index"""
 
     delete_index_url = "{}/{}".format(NODE, INDEX)
-    res = requests.delete(delete_index_url)
+    res = requests.delete(delete_index_url, headers=HEADERS)
     assert res.status_code == 200