feat: forward x-request-id headers (#270)

.gitignore

@@ -1,4 +1,5 @@
 *.py[cod]
+.python-version
 
 # C extensions
 *.so

CHANGELOG.rst

@@ -0,0 +1,22 @@
+Change Log
+==========
+
+..
+   All enhancements and patches to edx-rest-api-client will be documented
+   in this file.  It adheres to the structure of http://keepachangelog.com/ ,
+   but in reStructuredText instead of Markdown (for ease of incorporation into
+   Sphinx documentation and the PyPI description). Additionally, we no longer
+   track the date here since PyPi has its own history of dates based on when
+   the package is published.
+
+   This project adheres to Semantic Versioning (http://semver.org/).
+
+.. There should always be an "Unreleased" section for changes pending release.
+
+Unreleased
+----------
+* Nothing
+
+[5.5.1]
+--------
+feat: forward x-request-id headers if `crum` can find them

edx_rest_api_client/__version__.py

@@ -1 +1 @@
-__version__ = '5.5.0'
+__version__ = '5.5.1'

edx_rest_api_client/client.py

@@ -3,15 +3,15 @@
 import os
 import socket
 
+import crum
 import requests
 import requests.utils
 import slumber
-
 from edx_django_utils.cache import TieredCache
 from edx_django_utils.monitoring import set_custom_attribute
-from edx_rest_api_client.auth import BearerAuth, JwtAuth, SuppliedJwtAuth
-from edx_rest_api_client.__version__ import __version__
 
+from edx_rest_api_client.__version__ import __version__
+from edx_rest_api_client.auth import BearerAuth, JwtAuth, SuppliedJwtAuth
 
 # When caching tokens, use this value to err on expiring tokens a little early so they are
 # sure to be valid at the time they are used.
@@ -70,6 +70,14 @@ def _get_oauth_url(url):
     return stripped_url + '/oauth2/access_token'
 
 
+def get_request_id():
+    """
+    Helper to get the request id - usually set via an X-Request-ID header
+    """
+    request = crum.get_current_request()
+    return getattr(request, 'id', None)
+
+
 def get_oauth_access_token(url, client_id, client_secret, token_type='jwt', grant_type='client_credentials',
                            refresh_token=None,
                            timeout=(REQUEST_CONNECT_TIMEOUT, REQUEST_READ_TIMEOUT)):
@@ -272,17 +280,22 @@ def get_jwt_access_token(self):
         self._ensure_authentication()
         return self.auth.token
 
-    def request(self, method, url, **kwargs):  # pylint: disable=arguments-differ
+    def request(self, method, url, headers=None, **kwargs):  # pylint: disable=arguments-differ
         """
         Overrides Session.request to ensure that the session is authenticated.
 
         Note: Typically, users of the client won't call this directly, but will
         instead use Session.get or Session.post.
 
         """
+        request_id = get_request_id()
+        if headers is None:
+            headers = {}
+        if headers.get('X-Request-ID') is None and request_id is not None:
+            headers['X-Request-ID'] = request_id
         set_custom_attribute('api_client', 'OAuthAPIClient')
         self._ensure_authentication()
-        return super().request(method, url, **kwargs)
+        return super().request(method, url, headers=headers, **kwargs)
 
 
 class EdxRestApiClient(slumber.API):

edx_rest_api_client/tests/test_auth.py

@@ -1,5 +1,5 @@
 import datetime
-from unittest import mock, TestCase
+from unittest import TestCase, mock
 
 import jwt
 import requests

edx_rest_api_client/tests/test_client.py

@@ -1,24 +1,18 @@
 import datetime
 import json
 import os
-from unittest import mock, TestCase
+from unittest import TestCase, mock
 
 import ddt
 import requests
 import responses
-
 from edx_django_utils.cache import TieredCache
 from freezegun import freeze_time
 
 from edx_rest_api_client import __version__
 from edx_rest_api_client.auth import JwtAuth
-from edx_rest_api_client.client import (
-    EdxRestApiClient,
-    OAuthAPIClient,
-    get_and_cache_oauth_access_token,
-    get_oauth_access_token,
-    user_agent
-)
+from edx_rest_api_client.client import (EdxRestApiClient, OAuthAPIClient, get_and_cache_oauth_access_token,
+                                        get_oauth_access_token, user_agent)
 from edx_rest_api_client.tests.mixins import AuthenticationTestMixin
 
 URL = 'http://example.com/api/v2'
@@ -362,3 +356,19 @@ def test_get_jwt_access_token(self):
         client = OAuthAPIClient(self.base_url, self.client_id, self.client_secret)
         access_token = client.get_jwt_access_token()
         self.assertEqual(access_token, token)
+
+    @responses.activate
+    @mock.patch('edx_rest_api_client.client.get_request_id')
+    def test_request_id_forwarding(self, mock_get_request_id):
+        request_id = 'a-fake-request-id'
+        mock_get_request_id.return_value = request_id
+        token = 'abcd'
+        self._mock_auth_api(self.base_url + '/oauth2/access_token', 200, {'access_token': token, 'expires_in': 60})
+        client = OAuthAPIClient(self.base_url, self.client_id, self.client_secret)
+        post_url = self.base_url + '/oauth2/access_token'
+        responses.add(responses.POST,
+                      post_url,
+                      status=200,
+                      json={})
+        response = client.post(post_url, data={'test': 'ok'})
+        assert response.request.headers.get('X-Request-ID') == request_id