Repository: Yelp/pyramid-hypernova

Issue:
create plugin lifecycle event that allows injecting headers for batch requests
We have a use-case where it would be nice to be able to pass some HTTP headers down from the inbound request to the batch endpoint. Let's add a new plugin lifecycle method that can be used to modify the request headers that are sent to the batch endpoint.

The hook should accept the inbound pyramid request and the job group as arguments, and should return a dictionary of headers to be sent with the batch request.

Discussion:

Patch:
diff --git a/pyramid_hypernova/batch.py b/pyramid_hypernova/batch.py
index 0b98a5b..9f19682 100644
--- a/pyramid_hypernova/batch.py
+++ b/pyramid_hypernova/batch.py
@@ -146,12 +146,14 @@ class BatchRequest(object):
             job_groups = create_job_groups(self.jobs, self.max_batch_size)
             queries = []
 
+            # Fido is asynchronous and Python2.7 is bad at asynchronous, incurring 10-30ms of overhead
+            # when calling and immediately waiting on an HTTP request. If we only have one request to
+            # make, use synchronous Requests instead to save a little time.
+            synchronous = len(job_groups) == 1
+
             for job_group in job_groups:
-                # Fido is asynchronous and Python2.7 is bad at asynchronous, incurring 10-30ms of overhead
-                # when calling and immediately waiting on an HTTP request. If we only have one request to
-                # make, use synchronous Requests instead to save a little time.
-                synchronous = len(job_groups) == 1
-                query = HypernovaQuery(job_group, self.batch_url, self.json_encoder, synchronous)
+                request_headers = self.plugin_controller.transform_request_headers({})
+                query = HypernovaQuery(job_group, self.batch_url, self.json_encoder, synchronous, request_headers)
                 query.send()
                 queries.append((job_group, query))
 
diff --git a/pyramid_hypernova/plugins.py b/pyramid_hypernova/plugins.py
index 216e6d2..a82308b 100644
--- a/pyramid_hypernova/plugins.py
+++ b/pyramid_hypernova/plugins.py
@@ -42,6 +42,19 @@ class PluginController(object):
             current_jobs = plugin.prepare_request(current_jobs, original_jobs)
         return current_jobs
 
+    def transform_request_headers(self, headers):
+        """A reducer type function that is called with the request headers that
+        will be sent to Huypernova. Plugins can use this to inject request
+        headers, e.g. to add distributed tracing headers (like Zipkin).
+
+        :type headers: Dict[str, str]
+        :returns: the modified headers to be submitted
+        :rtype: Dict[str, str]
+        """
+        for plugin in self.plugins:
+            headers = plugin.transform_request_headers(headers)
+        return headers
+
     def should_send_request(self, jobs):
         """An every type function. If one returns false then the request is
         canceled.
@@ -127,6 +140,17 @@ class BasePlugin(object):
         """
         return current_jobs
 
+    def transform_request_headers(self, headers):
+        """A reducer type function that is called with the request headers that
+        will be sent to Huypernova. Plugins can use this to inject request
+        headers, e.g. to add distributed tracing headers (like Zipkin).
+
+        :type headers: Dict[str, str]
+        :returns: the modified headers to be submitted
+        :rtype: Dict[str, str]
+        """
+        return headers
+
     def should_send_request(self, jobs):
         """An every type function. If one returns false then the request is
         canceled.
diff --git a/pyramid_hypernova/request.py b/pyramid_hypernova/request.py
index f72782a..a3f1467 100644
--- a/pyramid_hypernova/request.py
+++ b/pyramid_hypernova/request.py
@@ -23,39 +23,40 @@ class HypernovaQueryError(Exception):
 class HypernovaQuery(object):
     """ Abstract Hypernova query """
 
-    def __init__(self, job_group, url, json_encoder, synchronous):
+    def __init__(self, job_group, url, json_encoder, synchronous, request_headers):
         """
         Build a Hypernova query.
         :param job_group: A job group (see create_job_groups)
         :param url: the URL of the Hypernova server we should query
         :param json_encoder: A JSON encoder to encode the query with
-        :param synchronous: True to synchronously query CRS (faster), False to
-            query asynchronously (allows parallelization)
+        :param synchronous: True to synchronously query hypernova (faster),
+            False to query asynchronously (allows parallelization)
+        :param request_headers: dict of request headers to add
         """
         self.job_group = job_group
         self.url = url
         self.json_encoder = json_encoder
         self.synchronous = synchronous
+        self.request_headers = request_headers
 
     def send(self):
         """ Query Hypernova """
         job_str = self.json_encoder.encode(create_jobs_payload(self.job_group))
         job_bytes = job_str.encode('utf-8')
 
+        request_headers = dict(self.request_headers)
+        request_headers['Content-Type'] = 'application/json'
+
         if self.synchronous:
             self.response = requests.post(
                 url=self.url,
-                headers={
-                    'Content-Type': 'application/json',
-                },
+                headers=request_headers,
                 data=job_bytes,
             )
         else:
             self.response = fido.fetch(
                 url=self.url,
-                headers={
-                    'Content-Type': ['application/json'],
-                },
+                headers={key: [value] for key, value in request_headers.items()},
                 method='POST',
                 body=job_bytes,
             )


Test patch:
diff --git a/tests/batch_test.py b/tests/batch_test.py
index 8e921a1..b4ef595 100644
--- a/tests/batch_test.py
+++ b/tests/batch_test.py
@@ -170,7 +170,7 @@ class TestBatchRequest(object):
             max_batch_size = batch_request.max_batch_size
             batch_count = (jobs_count + (max_batch_size - 1)) // max_batch_size
             assert mock_hypernova_query.call_count == batch_count
-            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1)
+            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1, {})
 
         assert response == {
             token_1.identifier: JobResult(
@@ -237,7 +237,7 @@ class TestBatchRequest(object):
             max_batch_size = batch_request.max_batch_size
             batch_count = (jobs_count + (max_batch_size - 1)) // max_batch_size
             assert mock_hypernova_query.call_count == batch_count
-            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1)
+            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1, {})
 
         assert response == {
             token_1.identifier: JobResult(
@@ -286,7 +286,7 @@ class TestBatchRequest(object):
             max_batch_size = batch_request.max_batch_size
             batch_count = (jobs_count + (max_batch_size - 1)) // max_batch_size
             assert mock_hypernova_query.call_count == batch_count
-            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1)
+            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1, {})
 
         assert response == {
             token.identifier: JobResult(
@@ -322,7 +322,7 @@ class TestBatchRequest(object):
             max_batch_size = batch_request.max_batch_size
             batch_count = (jobs_count + (max_batch_size - 1)) // max_batch_size
             assert mock_hypernova_query.call_count == batch_count
-            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1)
+            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1, {})
 
         assert response == {
             token.identifier: JobResult(
diff --git a/tests/plugins_test.py b/tests/plugins_test.py
index f48be65..a715c76 100644
--- a/tests/plugins_test.py
+++ b/tests/plugins_test.py
@@ -74,6 +74,15 @@ class TestPluginController(object):
         jobs = mock.Mock()
         assert plugin_controller.should_send_request(jobs) is expected_value
 
+    def test_transform_request_headers(self, plugins, plugin_controller):
+        plugins[0].transform_request_headers.side_effect = lambda x: dict(x, header1='yes')
+        plugins[1].transform_request_headers.side_effect = lambda x: dict(x, header2='yes')
+        assert plugin_controller.transform_request_headers({'foo': 'bar'}) == {
+            'foo': 'bar',
+            'header1': 'yes',
+            'header2': 'yes',
+        }
+
     def test_will_send_request(self, plugins, plugin_controller):
         jobs = mock.Mock()
         plugin_controller.will_send_request(jobs)
@@ -128,6 +137,10 @@ class TestBasePlugin(object):
             original_jobs,
         ) == current_jobs
 
+    def test_transform_request_headers(self):
+        plugin = BasePlugin()
+        assert plugin.transform_request_headers({'foo': 'bar'}) == {'foo': 'bar'}
+
     def test_should_send_request(self):
         plugin = BasePlugin()
         jobs = mock.sentinel.jobs
diff --git a/tests/request_test.py b/tests/request_test.py
index ea40268..2ccf361 100644
--- a/tests/request_test.py
+++ b/tests/request_test.py
@@ -52,18 +52,22 @@ class TestHypernovaQuery(object):
     def test_successful_send_synchronous(self, mock_fido_fetch, mock_requests_post):
         mock_requests_post.return_value.json.return_value = 'ayy lmao'
 
-        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), True)
+        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), True, {'header1': 'value1'})
         query.send()
 
         mock_fido_fetch.assert_not_called()
-        mock_requests_post.assert_called_once()
+        mock_requests_post.assert_called_once_with(
+            url='google.com',
+            headers={'header1': 'value1', 'Content-Type': 'application/json'},
+            data=mock.ANY,
+        )
 
         assert query.json() == 'ayy lmao'
 
     def test_erroneous_send_synchronous(self, mock_fido_fetch, mock_requests_post):
         mock_requests_post.return_value.raise_for_status.side_effect = HTTPError('ayy lmao')
 
-        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), True)
+        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), True, {})
         query.send()
 
         mock_fido_fetch.assert_not_called()
@@ -77,10 +81,15 @@ class TestHypernovaQuery(object):
         mock_fido_fetch.return_value.wait.return_value.code = 200
         mock_fido_fetch.return_value.wait.return_value.json.return_value = 'ayy lmao'
 
-        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False)
+        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False, {'header1': 'value1'})
         query.send()
 
-        mock_fido_fetch.assert_called_once()
+        mock_fido_fetch.assert_called_once_with(
+            url='google.com',
+            method='POST',
+            headers={'header1': ['value1'], 'Content-Type': ['application/json']},
+            body=mock.ANY,
+        )
         mock_requests_post.assert_not_called()
 
         assert query.json() == 'ayy lmao'
@@ -88,7 +97,7 @@ class TestHypernovaQuery(object):
     def test_erroneous_send_asynchronous(self, mock_fido_fetch, mock_requests_post):
         mock_fido_fetch.return_value.wait.side_effect = NetworkError('ayy lmao')
 
-        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False)
+        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False, {})
         query.send()
 
         mock_fido_fetch.assert_called_once()
@@ -103,7 +112,7 @@ class TestHypernovaQuery(object):
         mock_fido_fetch.return_value.wait.return_value.body = b'<h1>504 Bad Gateway</h1>'
         mock_fido_fetch.return_value.wait.return_value.json.side_effect = AssertionError()
 
-        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False)
+        query = HypernovaQuery(TEST_JOB_GROUP, 'google.com', JSONEncoder(), False, {})
         query.send()
 
         mock_fido_fetch.assert_called_once()