From 03b17cf74d732d2c3eee62463616fca7179c7dd1 Mon Sep 17 00:00:00 2001
From: Daniel Mitterdorfer <daniel@mitterdorfer.name>
Date: Wed, 28 Jun 2017 14:37:44 +0200
Subject: [PATCH] Allow to pass arbitrary request parameters for ES queries

Closes #290
---
 docs/track.rst             |  5 +++++
 esrally/driver/runner.py   | 13 +++++++++++--
 esrally/track/params.py    |  2 ++
 tests/track/params_test.py | 31 +++++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/docs/track.rst b/docs/track.rst
index 3b2325e8c..cadab97f5 100644
--- a/docs/track.rst
+++ b/docs/track.rst
@@ -191,6 +191,7 @@ With the operation type ``search`` you can execute `request body searches <http:
 * ``index`` (optional): An `index pattern <https://www.elastic.co/guide/en/elasticsearch/reference/current/multi-index.html>`_ that defines which indices should be targeted by this query. Only needed if the ``index`` section contains more than one index. Otherwise, Rally will automatically derive the index to use. If you have defined multiple indices and want to query all of them, just specify ``"index": "_all"``.
 * ``type`` (optional): Defines the type within the specified index for this query.
 * ``cache`` (optional): Whether to use the query request cache. By default, Rally will define no value thus the default depends on the benchmark candidate settings and Elasticsearch version.
+* ``request-params`` (optional): A structure containing arbitrary request parameters. The supported parameters names are documented in the `Python ES client API docs <http://elasticsearch-py.readthedocs.io/en/master/api.html#elasticsearch.Elasticsearch.search>`_. Parameters that are implicitly set by Rally (e.g. `body` or `request_cache`) are not supported (i.e. you should not try to set them and if so expect unspecified behavior).
 * ``body`` (mandatory): The query body.
 * ``pages`` (optional): Number of pages to retrieve. If this parameter is present, a scroll query will be executed. If you want to retrieve all result pages, use the value "all".
 * ``results-per-page`` (optional):  Number of documents to retrieve per page for scroll queries.
@@ -204,6 +205,10 @@ Example::
         "query": {
           "match_all": {}
         }
+      },
+      "request-params": {
+        "_source_include": "some_field",
+        "analyze_wildcard": false
       }
     }
 
diff --git a/esrally/driver/runner.py b/esrally/driver/runner.py
index d386a7705..20da3b721 100644
--- a/esrally/driver/runner.py
+++ b/esrally/driver/runner.py
@@ -393,10 +393,17 @@ def __call__(self, es, params):
             return self.request_body_query(es, params)
 
     def request_body_query(self, es, params):
-        es.search(index=params["index"], doc_type=params["type"], request_cache=params["use_request_cache"], body=params["body"])
+        request_params = params.get("request_params", {})
+        es.search(
+            index=params["index"],
+            doc_type=params["type"],
+            request_cache=params["use_request_cache"],
+            body=params["body"],
+            **request_params)
         return 1, "ops"
 
     def scroll_query(self, es, params):
+        request_params = params.get("request_params", {})
         hits = 0
         retrieved_pages = 0
         self.es = es
@@ -412,7 +419,9 @@ def scroll_query(self, es, params):
                     sort="_doc",
                     scroll="10s",
                     size=params["items_per_page"],
-                    request_cache=params["use_request_cache"])
+                    request_cache=params["use_request_cache"],
+                    **request_params
+                )
                 # This should only happen if we concurrently create an index and start searching
                 self.scroll_id = r.get("_scroll_id", None)
             else:
diff --git a/esrally/track/params.py b/esrally/track/params.py
index 27e6709f8..466517c24 100644
--- a/esrally/track/params.py
+++ b/esrally/track/params.py
@@ -130,11 +130,13 @@ def __init__(self, indices, params):
         query_body = params.get("body", None)
         pages = params.get("pages", None)
         items_per_page = params.get("results-per-page", None)
+        request_params = params.get("request-params", {})
 
         self.query_params = {
             "index": index_name,
             "type": type_name,
             "use_request_cache": request_cache,
+            "request_params": request_params,
             "body": query_body
         }
 
diff --git a/tests/track/params_test.py b/tests/track/params_test.py
index b985d5d12..71de1f513 100644
--- a/tests/track/params_test.py
+++ b/tests/track/params_test.py
@@ -666,3 +666,34 @@ def test_can_register_class_as_param_source(self):
         self.assertEqual({"class-key": 42}, source.params())
 
         params._unregister_param_source_for_name(source_name)
+
+
+class SearchParamSourceTests(TestCase):
+    def test_passes_request_parameters(self):
+        type1 = track.Type("type1", mapping_file="", number_of_documents=3)
+        index1 = track.Index(name="index1", auto_managed=True, types=[type1])
+
+        source = params.SearchParamSource(indices=[index1], params={
+            "request-params": {
+                "_source_include": "some_field"
+            },
+            "body": {
+                "query": {
+                    "match_all": {}
+                }
+            }
+        })
+        p = source.params()
+
+        self.assertEqual(5, len(p))
+        self.assertEqual("index1", p["index"])
+        self.assertEqual("type1", p["type"])
+        self.assertEqual({
+            "_source_include": "some_field"
+        }, p["request_params"])
+        self.assertFalse(p["use_request_cache"])
+        self.assertEqual({
+            "query": {
+                "match_all": {}
+            }
+        }, p["body"])