Refactor graph client to avoid preparing request

Author: samwelkanda

msgraph/core/_graph_client.py

@@ -2,6 +2,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import json
+
 from requests import Request, Session
 
 from ._client_factory import HTTPClientFactory
@@ -25,23 +27,24 @@
 ]
 
 
-def attach_context(func):
-    """Attaches a request context object to every graph request"""
+def collect_options(func):
+    """Collect middleware options into a middleware control dict and pass it as a header"""
     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
+        if 'headers' in kwargs.keys():
+            kwargs['headers'].update({'middleware_control': json.dumps(middleware_control)})
+        else:
+            kwargs['headers'] = dict()
+            kwargs['headers'].update({'middleware_control': json.dumps(middleware_control)})
 
-        return request
+        return func(*args, **kwargs)
 
     return wrapper
 
@@ -81,16 +84,34 @@ def __init__(self, **kwargs):
         """
         self.graph_session = self.get_graph_session(**kwargs)
 
+    @collect_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
         """
-        prepared_request = self.prepare_request('GET', self._graph_url(url), **kwargs)
-        return self.graph_session.send(prepared_request)
+        return self.graph_session.get(self._graph_url(url), **kwargs)
+
+    def options(self, url, **kwargs):
+        r"""Sends a OPTIONS 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.options(self._graph_url(url), **kwargs)
 
-    def post(self, url, **kwargs):
+    def head(self, url, **kwargs):
+        r"""Sends a HEAD 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.head(self._graph_url(url), **kwargs)
+
+    def post(self, url, data=None, json=None, **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
@@ -99,8 +120,7 @@ def post(self, url, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        prepared_request = self.prepare_request('POST', self._graph_url(url), **kwargs)
-        return self.graph_session.send(prepared_request)
+        return self.graph_session.post(self._graph_url(url), data=data, json=json, **kwargs)
 
     def put(self, url, data=None, **kwargs):
         r"""Sends a PUT request. Returns :class:`Response` object.
@@ -110,8 +130,8 @@ def put(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        prepared_request = self.prepare_request('PUT', self._graph_url(url), **kwargs)
-        return self.graph_session.send(prepared_request)
+
+        return self.graph_session.put(self._graph_url(url), data=data, **kwargs)
 
     def patch(self, url, data=None, **kwargs):
         r"""Sends a PATCH request. Returns :class:`Response` object.
@@ -121,17 +141,15 @@ def patch(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        prepared_request = self.prepare_request('PATCH', self._graph_url(url), **kwargs)
-        return self.graph_session.send(prepared_request)
+        return self.graph_session.patch(self._graph_url(url), data=data, **kwargs)
 
     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
         """
-        prepared_request = self.prepare_request('DELETE', self._graph_url(url), **kwargs)
-        return self.graph_session.send(prepared_request)
+        return self.graph_session.delete(self._graph_url(url), **kwargs)
 
     def _graph_url(self, url: str) -> str:
         """Appends BASE_URL to user provided path
@@ -140,23 +158,6 @@ 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,
-            headers=kwargs.get('headers'),
-            files=kwargs.get('files'),
-            data=kwargs.get('data'),
-            json=kwargs.get('json'),
-            params=kwargs.get('params'),
-            auth=kwargs.get('auth'),
-            cookies=kwargs.get('cookies'),
-            hooks=kwargs.get('hooks'),
-        )
-        prepared = Session().prepare_request(req)
-        return prepared
-
     @staticmethod
     def get_graph_session(**kwargs):
         """Method to always return a single instance of a HTTP Client"""

msgraph/core/middleware/middleware.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import json
 import ssl
 
 from requests.adapters import HTTPAdapter
@@ -31,9 +32,14 @@ def add_middleware(self, middleware):
 
     def send(self, request, **kwargs):
 
-        if not hasattr(request, 'context'):
-            headers = request.headers
-            request.context = RequestContext(dict(), headers)
+        middleware_control_json = request.headers.pop('middleware_control', None)
+        if middleware_control_json:
+            middleware_control = json.loads(middleware_control_json)
+        else:
+            middleware_control = dict()
+
+        headers = request.headers
+        request.context = RequestContext(middleware_control, headers)
 
         if self._middleware_present():
             return self._first_middleware.send(request, **kwargs)

tests/integration/test_graphclient.py

@@ -3,26 +3,20 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import pytest
+from azure.identity import EnvironmentCredential
 from requests import Session
 
 from msgraph.core import APIVersion, GraphClient
 from msgraph.core.middleware.authorization import AuthorizationHandler
 
 
-class _CustomTokenCredential:
-    def get_token(self, scopes):
-        return ['{token:https://graph.microsoft.com/}']
-
-
 def test_graph_client_with_default_middleware():
     """
     Test that a graph client uses default middleware if none are provided
     """
-    credential = _CustomTokenCredential()
-    client = GraphClient(credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me'
-    )
+    # credential = _CustomTokenCredential()
+    client = GraphClient(credential=EnvironmentCredential())
+    response = client.get('https://graph.microsoft.com/v1.0/users')
     assert response.status_code == 200
 
 
@@ -32,52 +26,42 @@ def test_graph_client_with_user_provided_session():
     """
 
     session = Session()
-    credential = _CustomTokenCredential()
-    client = GraphClient(session=session, credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me'
-    )
+    client = GraphClient(session=session, credential=EnvironmentCredential())
+    response = client.get('https://graph.microsoft.com/v1.0/users', )
     assert response.status_code == 200
 
 
 def test_graph_client_with_custom_settings():
     """
     Test that the graph client works with user provided configuration
     """
-    credential = _CustomTokenCredential()
+    credential = EnvironmentCredential()
     client = GraphClient(api_version=APIVersion.beta, credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me'
-    )
+    response = client.get('https://graph.microsoft.com/v1.0/users', )
     assert response.status_code == 200
 
 
 def test_graph_client_with_custom_middleware():
     """
     Test client factory works with user provided middleware
     """
-    credential = _CustomTokenCredential()
+    credential = EnvironmentCredential()
     middleware = [
         AuthorizationHandler(credential),
     ]
     client = GraphClient(middleware=middleware)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me'
-    )
+    response = client.get('https://graph.microsoft.com/v1.0/users', )
     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']
+    credential = EnvironmentCredential()
+    scopes = ['https://graph.microsoft.com/.default']
     client = GraphClient(credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me',
-        scopes=scopes
-    )
+    response = client.get('https://graph.microsoft.com/v1.0/users', scopes=scopes)
     assert response.status_code == 200
     assert hasattr(response.request, 'context')
 
@@ -86,13 +70,10 @@ 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']
+    credential = EnvironmentCredential()
+    scopes = ['https://graph.microsoft.com/.default']
     client = GraphClient(credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me',
-        scopes=scopes
-    )
+    response = client.get('https://graph.microsoft.com/v1.0/users', scopes=scopes)
     assert response.status_code == 200
     assert 'scopes' in response.request.context.middleware_control.keys()
     assert response.request.context.middleware_control['scopes'] == scopes
@@ -103,12 +84,8 @@ def test_graph_client_allows_passing_optional_kwargs():
     Test the graph client allows passing optional kwargs native to the requests library
     such as stream, proxy and cert.
     """
-    credential = _CustomTokenCredential()
-    scopes = ['User.Read.All']
+    credential = EnvironmentCredential()
+    scopes = ['https://graph.microsoft.com/.default']
     client = GraphClient(credential=credential)
-    response = client.get(
-        'https://proxy.apisandbox.msdn.microsoft.com/svc?url=https://graph.microsoft.com/v1.0/me',
-        scopes=scopes,
-        stream=True
-    )
+    response = client.get('https://graph.microsoft.com/v1.0/users', scopes=scopes, stream=True)
     assert response.status_code == 200