Feat/request context

Pipfile

@@ -5,15 +5,15 @@ name = "pypi"
 
 [packages] # Packages required to run the application
 requests = "==2.23.0"
-pre-commit = "==2.10.1"
+azure-identity = "==1.6.0"
 
 [dev-packages] # Packages required to develop the application
 coverage = "==5.0.3"
 responses = "==0.10.12"
 flit = "==2.2.0"
-azure-identity = "==1.5.0"
 isort = "==5.7.0"
 yapf = "==0.30.0"
 mypy = "==0.800"
 pylint = "==2.7.4"
-pytest = "*"
+pytest = "==6.2.4"
+pre-commit = "==2.13.0"

msgraphcore/__init__.py

@@ -1,3 +1,5 @@
-from .constants import SDK_VERSION
+from msgraphcore.client_factory import HTTPClientFactory
+from msgraphcore.constants import SDK_VERSION
+from msgraphcore.graph_client import GraphClient
 
 __version__ = SDK_VERSION

msgraphcore/enums.py

@@ -9,6 +9,17 @@ class APIVersion(str, Enum):
     v1 = 'v1.0'
 
 
+class FeatureUsageFlag(int, Enum):
+    """Enumerated list of values used to flag usage of specific middleware"""
+
+    NONE = 0
+    REDIRECT_HANDLER_ENABLED = 1
+    RETRY_HANDLER_ENABLED = 2
+    AUTH_HANDLER_ENABLED = 4
+    DEFAULT_HTTP_PROVIDER_ENABLED = 8
+    LOGGING_HANDLER_ENABLED = 16
+
+
 class NationalClouds(str, Enum):
     """Enumerated list of supported sovereign clouds"""
 

msgraphcore/graph_client.py

@@ -1,11 +1,33 @@
 from typing import List, Optional
 
-from requests import Session
+from requests import Request, Session
 
 from msgraphcore.client_factory import HTTPClientFactory
 from msgraphcore.middleware.abc_token_credential import TokenCredential
 from msgraphcore.middleware.middleware import BaseMiddleware
-from msgraphcore.middleware.options.middleware_control import middleware_control
+from msgraphcore.middleware.request_context import RequestContext
+
+supported_options = ['scopes', 'custom_option']
+
+
+def attach_context(func):
+    def wrapper(*args, **kwargs):
+        middleware_control = dict()
+
+        for option in supported_options:
+            value = kwargs.pop(option, None)
+            if value:
+                middleware_control.update({option: value})
+
+        headers = kwargs.get('headers', {})
+        request_context = RequestContext(middleware_control, headers)
+
+        request = func(*args, **kwargs)
+        request.context = request_context
+
+        return request
+
+    return wrapper
 
 
 class GraphClient:
@@ -43,17 +65,16 @@ def __init__(self, **kwargs):
         """
         self.graph_session = self.get_graph_session(**kwargs)
 
-    @middleware_control.get_middleware_options
     def get(self, url: str, **kwargs):
         r"""Sends a GET request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.get(self._graph_url(url), **kwargs)
+        prepared_request = self.prepare_request('GET', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepared_request)
 
-    @middleware_control.get_middleware_options
-    def post(self, url, data=None, json=None, **kwargs):
+    def post(self, url, **kwargs):
         r"""Sends a POST request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
         :param data: (optional) Dictionary, list of tuples, bytes, or file-like
@@ -62,9 +83,9 @@ def post(self, url, data=None, json=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.post(self._graph_url(url), data, json, **kwargs)
+        prepared_request = self.prepare_request('POST', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepared_request)
 
-    @middleware_control.get_middleware_options
     def put(self, url, data=None, **kwargs):
         r"""Sends a PUT request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
@@ -73,9 +94,9 @@ def put(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.put(self._graph_url(url), data, **kwargs)
+        prepared_request = self.prepare_request('PUT', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepared_request)
 
-    @middleware_control.get_middleware_options
     def patch(self, url, data=None, **kwargs):
         r"""Sends a PATCH request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
@@ -84,16 +105,17 @@ def patch(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.patch(self._graph_url(url), data, **kwargs)
+        prepared_request = self.prepare_request('PATCH', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepared_request)
 
-    @middleware_control.get_middleware_options
     def delete(self, url, **kwargs):
         r"""Sends a DELETE request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.delete(self._graph_url(url), **kwargs)
+        prepared_request = self.prepare_request('DELETE', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepared_request)
 
     def _graph_url(self, url: str) -> str:
         """Appends BASE_URL to user provided path
@@ -102,12 +124,18 @@ def _graph_url(self, url: str) -> str:
         """
         return self.graph_session.base_url + url if (url[0] == '/') else url
 
+    @attach_context
+    def prepare_request(self, method, url, **kwargs):
+        req = Request(method, url, **kwargs)
+        prepared = Session().prepare_request(req)
+        return prepared
+
     @staticmethod
     def get_graph_session(**kwargs):
         """Method to always return a single instance of a HTTP Client"""
 
-        credential = kwargs.get('credential')
-        middleware = kwargs.get('middleware')
+        credential = kwargs.pop('credential', None)
+        middleware = kwargs.pop('middleware', None)
 
         if credential and middleware:
             raise ValueError(
@@ -117,5 +145,5 @@ def get_graph_session(**kwargs):
             raise ValueError("Invalid parameters!. Missing TokenCredential or middleware")
 
         if credential:
-            return HTTPClientFactory(**kwargs).create_with_default_middleware(credential)
+            return HTTPClientFactory(**kwargs).create_with_default_middleware(credential, **kwargs)
         return HTTPClientFactory(**kwargs).create_with_custom_middleware(middleware)

msgraphcore/middleware/authorization.py

@@ -1,7 +1,6 @@
-from msgraphcore.constants import AUTH_MIDDLEWARE_OPTIONS
+from msgraphcore.enums import FeatureUsageFlag
 from msgraphcore.middleware.abc_token_credential import TokenCredential
 from msgraphcore.middleware.middleware import BaseMiddleware
-from msgraphcore.middleware.options.middleware_control import middleware_control
 
 
 class AuthorizationHandler(BaseMiddleware):
@@ -12,7 +11,11 @@ def __init__(self, credential: TokenCredential, **kwargs):
         self.retry_count = 0
 
     def send(self, request, **kwargs):
-        request.headers.update({'Authorization': 'Bearer {}'.format(self._get_access_token())})
+        context = request.context
+        request.headers.update(
+            {'Authorization': 'Bearer {}'.format(self._get_access_token(context))}
+        )
+        context.set_feature_usage = FeatureUsageFlag.AUTH_HANDLER_ENABLED
         response = super().send(request, **kwargs)
 
         # Token might have expired just before transmission, retry the request one more time
@@ -21,13 +24,9 @@ def send(self, request, **kwargs):
             return self.send(request, **kwargs)
         return response
 
-    def _get_access_token(self):
-        return self.credential.get_token(*self.get_scopes())[0]
+    def _get_access_token(self, context):
+        return self.credential.get_token(*self.get_scopes(context))[0]
 
-    def get_scopes(self):
+    def get_scopes(self, context):
         # Checks if there are any options for this middleware
-        auth_options_present = middleware_control.get(AUTH_MIDDLEWARE_OPTIONS)
-        # If there is, get the scopes from the options
-        if auth_options_present:
-            return auth_options_present.scopes
-        return self.scopes
+        return context.middleware_control.get('scopes', self.scopes)

msgraphcore/middleware/middleware.py

@@ -1,8 +1,11 @@
 import ssl
+import uuid
 
 from requests.adapters import HTTPAdapter
 from urllib3 import PoolManager
 
+from msgraphcore.middleware.request_context import RequestContext
+
 
 class MiddlewarePipeline(HTTPAdapter):
     """MiddlewarePipeline, entry point of middleware
@@ -22,6 +25,11 @@ def add_middleware(self, middleware):
             self._middleware = middleware
 
     def send(self, request, **kwargs):
+
+        if not hasattr(request, 'context'):
+            headers = request.headers
+            request.context = RequestContext(dict(), headers)
+
         if self._middleware_present():
             return self._middleware.send(request, **kwargs)
         # No middleware in pipeline, call superclass' send

msgraphcore/middleware/request_context.py

@@ -0,0 +1,18 @@
+import uuid
+
+from msgraphcore.enums import FeatureUsageFlag
+
+
+class RequestContext:
+    def __init__(self, middleware_control, headers):
+        self.middleware_control = middleware_control
+        self.client_request_id = headers.get('client-request-id', str(uuid.uuid4()))
+        self._feature_usage = FeatureUsageFlag.NONE
+
+    @property
+    def feature_usage(self):
+        return hex(self._feature_usage)
+
+    @feature_usage.setter
+    def set_feature_usage(self, flag: FeatureUsageFlag):
+        self._feature_usage = self._feature_usage | flag

samples/samples.py

@@ -3,11 +3,11 @@
 
 from azure.identity import InteractiveBrowserCredential
 
-from msgraphcore import GraphSession
+from msgraphcore.graph_client import GraphClient
 
 scopes = ['user.read']
-browser_credential = InteractiveBrowserCredential(client_id='ENTER_YOUR_CLIENT_ID')
-graph_session = GraphSession(browser_credential, scopes)
+browser_credential = InteractiveBrowserCredential(client_id='YOUR_CLIENT_ID')
+graph_session = GraphClient(credential=browser_credential)
 
 
 def post_sample():

tests/integration/test_graph_client.py → tests/integration/test_graphclient.py

@@ -49,7 +49,7 @@ def test_graph_client_with_custom_settings():
     assert response.status_code == 200
 
 
-def test_client_factory_with_custom_middleware():
+def test_graph_client_with_custom_middleware():
     """
     Test client factory works with user provided middleware
     """
@@ -62,3 +62,34 @@ def test_client_factory_with_custom_middleware():
         'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me'
     )
     assert response.status_code == 200
+
+
+def test_graph_client_adds_context_to_request():
+    """
+    Test the graph client adds a context object to a request
+    """
+    credential = _CustomTokenCredential()
+    scopes = ['User.Read.All']
+    client = GraphClient(credential=credential)
+    response = client.get(
+        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me',
+        scopes=scopes
+    )
+    assert response.status_code == 200
+    assert hasattr(response.request, 'context')
+
+
+def test_graph_client_picks_options_from_kwargs():
+    """
+    Test the graph client picks middleware options from kwargs and sets them in the context
+    """
+    credential = _CustomTokenCredential()
+    scopes = ['User.Read.All']
+    client = GraphClient(credential=credential)
+    response = client.get(
+        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me',
+        scopes=scopes
+    )
+    assert response.status_code == 200
+    assert 'scopes' in response.request.context.middleware_control.keys()
+    assert response.request.context.middleware_control['scopes'] == scopes