Skip to content

Commit

Permalink
Allow to pass arbitrary request parameters for ES queries
Browse files Browse the repository at this point in the history
Closes #290
  • Loading branch information
danielmitterdorfer committed Jun 28, 2017
1 parent 57c4c2d commit 03b17cf
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 2 deletions.
5 changes: 5 additions & 0 deletions docs/track.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -204,6 +205,10 @@ Example::
"query": {
"match_all": {}
}
},
"request-params": {
"_source_include": "some_field",
"analyze_wildcard": false
}
}

Expand Down
13 changes: 11 additions & 2 deletions esrally/driver/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions esrally/track/params.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
31 changes: 31 additions & 0 deletions tests/track/params_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])

0 comments on commit 03b17cf

Please sign in to comment.