From 2a1fa1dec39861734c1221ddeaa6368748025fca Mon Sep 17 00:00:00 2001
From: Daniel Mitterdorfer <daniel@mitterdorfer.name>
Date: Fri, 3 Apr 2020 13:13:04 +0200
Subject: [PATCH] Allow to disable HTTP compression for all queries

With this commit we expand the support added in #941 so that not only
scroll queries can take advantage but also regular queries.

Relates #941
---
 docs/track.rst              |  2 +-
 esrally/driver/runner.py    | 20 ++++++-------
 tests/driver/runner_test.py | 56 ++++++++++++++++++++++---------------
 3 files changed, 44 insertions(+), 34 deletions(-)

diff --git a/docs/track.rst b/docs/track.rst
index c2449104e..e2df75f59 100644
--- a/docs/track.rst
+++ b/docs/track.rst
@@ -402,10 +402,10 @@ With the operation type ``search`` you can execute `request body searches <http:
         2. Rally will not attempt to serialize the parameters and pass them as is. Always use "true" / "false" strings for boolean parameters (see example below).
 
 * ``body`` (mandatory): The query body.
+* ``response-compression-enabled`` (optional, defaults to ``true``): Allows to disable HTTP compression of responses. As these responses are sometimes large and decompression may be a bottleneck on the client, it is possible to turn off response compression.
 * ``detailed-results`` (optional, defaults to ``false``): Records more detailed meta-data about queries. As it analyzes the corresponding response in more detail, this might incur additional overhead which can skew measurement results. This flag is ineffective for scroll queries.
 * ``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.
-* ``response-compression-enabled`` (optional, defaults to ``true``): Allows to disable HTTP compression of scroll responses. As these responses are sometimes large and decompression may be a bottleneck on the client, it is possible to turn off response compression. This option is ineffective for regular queries.
 
 If ``detailed-results`` is set to ``true``, the following meta-data properties will be determined and stored:
 
diff --git a/esrally/driver/runner.py b/esrally/driver/runner.py
index 68a62d87e..ace19a204 100644
--- a/esrally/driver/runner.py
+++ b/esrally/driver/runner.py
@@ -772,15 +772,12 @@ async def request_body_query(self, es, params):
         body = mandatory(params, "body", self)
         doc_type = params.get("type")
         detailed_results = params.get("detailed-results", False)
-        params = request_params
+        headers = self._headers(params)
 
         # disable eager response parsing - responses might be huge thus skewing results
         es.return_raw_response()
 
-        if doc_type is not None:
-            r = await self._raw_search(es, doc_type, index, body, params)
-        else:
-            r = await es.search(index=index, body=body, params=params)
+        r = await self._raw_search(es, doc_type, index, body, request_params, headers)
 
         if detailed_results:
             props = parse(r, ["hits.total", "hits.total.value", "hits.total.relation", "timed_out", "took"])
@@ -815,11 +812,7 @@ async def scroll_query(self, es, params):
         # explicitly convert to int to provoke an error otherwise
         total_pages = sys.maxsize if params["pages"] == "all" else int(params["pages"])
         size = params.get("results-per-page")
-        # reduces overhead due to decompression of very large responses
-        if params.get("response-compression-enabled", True):
-            headers = None
-        else:
-            headers = {"Accept-Encoding": "identity"}
+        headers = self._headers(params)
         scroll_id = None
 
         # disable eager response parsing - responses might be huge thus skewing results
@@ -896,6 +889,13 @@ def _default_request_params(self, params):
             request_params["request_cache"] = str(cache).lower()
         return request_params
 
+    def _headers(self, params):
+        # reduces overhead due to decompression of very large responses
+        if params.get("response-compression-enabled", True):
+            return None
+        else:
+            return {"Accept-Encoding": "identity"}
+
     def __repr__(self, *args, **kwargs):
         return "query"
 
diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py
index 840dc4eba..57a16a2a0 100644
--- a/tests/driver/runner_test.py
+++ b/tests/driver/runner_test.py
@@ -1085,7 +1085,7 @@ async def test_query_match_only_request_body_defined(self, es):
                 ]
             }
         }
-        es.search.return_value = as_future(io.StringIO(json.dumps(search_response)))
+        es.transport.perform_request.return_value = as_future(io.StringIO(json.dumps(search_response)))
 
         query_runner = runner.Query()
 
@@ -1110,10 +1110,12 @@ async def test_query_match_only_request_body_defined(self, es):
         self.assertEqual(5, result["took"])
         self.assertFalse("error-type" in result)
 
-        es.search.assert_called_once_with(
-            index="_all",
+        es.transport.perform_request.assert_called_once_with(
+            "GET",
+            "/_all/_search",
+            params={"request_cache": "true"},
             body=params["body"],
-            params={"request_cache": "true"}
+            headers=None
         )
         es.clear_scroll.assert_not_called()
 
@@ -1139,7 +1141,7 @@ async def test_query_match_using_request_params(self, es):
                 ]
             }
         }
-        es.search.return_value = as_future(io.StringIO(json.dumps(response)))
+        es.transport.perform_request.return_value = as_future(io.StringIO(json.dumps(response)))
 
         query_runner = runner.Query()
         params = {
@@ -1162,13 +1164,15 @@ async def test_query_match_using_request_params(self, es):
         self.assertEqual(62, result["took"])
         self.assertFalse("error-type" in result)
 
-        es.search.assert_called_once_with(
-            index="_all",
-            body=params["body"],
+        es.transport.perform_request.assert_called_once_with(
+            "GET",
+            "/_all/_search",
             params={
                 "request_cache": "false",
                 "q": "user:kimchy"
-            }
+            },
+            body=params["body"],
+            headers=None
         )
         es.clear_scroll.assert_not_called()
 
@@ -1194,7 +1198,7 @@ async def test_query_no_detailed_results(self, es):
                 ]
             }
         }
-        es.search.return_value = as_future(io.StringIO(json.dumps(response)))
+        es.transport.perform_request.return_value = as_future(io.StringIO(json.dumps(response)))
 
         query_runner = runner.Query()
         params = {
@@ -1216,10 +1220,12 @@ async def test_query_no_detailed_results(self, es):
         self.assertNotIn("took", result)
         self.assertNotIn("error-type", result)
 
-        es.search.assert_called_once_with(
-            index="_all",
+        es.transport.perform_request.assert_called_once_with(
+            "GET",
+            "/_all/_search",
+            params={"q": "user:kimchy"},
             body=params["body"],
-            params={"q": "user:kimchy"}
+            headers=None
         )
         es.clear_scroll.assert_not_called()
 
@@ -1241,7 +1247,7 @@ async def test_query_hits_total_as_number(self, es):
                 ]
             }
         }
-        es.search.return_value = as_future(io.StringIO(json.dumps(search_response)))
+        es.transport.perform_request.return_value = as_future(io.StringIO(json.dumps(search_response)))
 
         query_runner = runner.Query()
 
@@ -1266,12 +1272,14 @@ async def test_query_hits_total_as_number(self, es):
         self.assertEqual(5, result["took"])
         self.assertFalse("error-type" in result)
 
-        es.search.assert_called_once_with(
-            index="_all",
-            body=params["body"],
+        es.transport.perform_request.assert_called_once_with(
+            "GET",
+            "/_all/_search",
             params={
                 "request_cache": "true"
-            }
+            },
+            body=params["body"],
+            headers=None
         )
         es.clear_scroll.assert_not_called()
 
@@ -1296,14 +1304,14 @@ async def test_query_match_all(self, es):
                 ]
             }
         }
-        es.search.return_value = as_future(io.StringIO(json.dumps(search_response)))
+        es.transport.perform_request.return_value = as_future(io.StringIO(json.dumps(search_response)))
 
         query_runner = runner.Query()
 
         params = {
             "index": "unittest",
             "detailed-results": True,
-            "cache": None,
+            "response-compression-enabled": False,
             "body": {
                 "query": {
                     "match_all": {}
@@ -1322,10 +1330,12 @@ async def test_query_match_all(self, es):
         self.assertEqual(5, result["took"])
         self.assertFalse("error-type" in result)
 
-        es.search.assert_called_once_with(
-            index="unittest",
+        es.transport.perform_request.assert_called_once_with(
+            "GET",
+            "/unittest/_search",
+            params={},
             body=params["body"],
-            params={}
+            headers={"Accept-Encoding": "identity"}
         )
         es.clear_scroll.assert_not_called()