Set context and feature usage in transport

Author: samwelkanda

src/msgraph_core/graph_client_factory.py

@@ -8,11 +8,10 @@
 
 import httpx
 from kiota_http.kiota_client_factory import KiotaClientFactory
-from kiota_http.middleware import AsyncKiotaTransport
 from kiota_http.middleware.middleware import BaseMiddleware
 
 from ._enums import APIVersion, NationalClouds
-from .middleware import GraphTelemetryHandler
+from .middleware import AsyncGraphTransport, GraphTelemetryHandler
 
 
 class GraphClientFactory(KiotaClientFactory):
@@ -40,9 +39,10 @@ def create_with_default_middleware(
             middleware, current_transport
         )
 
-        client._transport = AsyncKiotaTransport(
+        client._transport = AsyncGraphTransport(
             transport=current_transport, pipeline=middleware_pipeline
         )
+        client._transport.pipeline
         return client
 
     @staticmethod
@@ -66,7 +66,7 @@ def create_with_custom_middleware(
             middleware, current_transport
         )
 
-        client._transport = AsyncKiotaTransport(
+        client._transport = AsyncGraphTransport(
             transport=current_transport, pipeline=middleware_pipeline
         )
         return client

src/msgraph_core/middleware/__init__.py

@@ -2,5 +2,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+from .async_graph_transport import AsyncGraphTransport
 from .request_context import GraphRequestContext
 from .telemetry import GraphTelemetryHandler

src/msgraph_core/middleware/async_graph_transport.py

@@ -0,0 +1,44 @@
+import json
+
+import httpx
+from kiota_http.middleware import MiddlewarePipeline, RedirectHandler, RetryHandler
+
+from .._enums import FeatureUsageFlag
+from .request_context import GraphRequestContext
+
+
+class AsyncGraphTransport(httpx.AsyncBaseTransport):
+    """A custom transport for requests to the Microsoft Graph API
+    """
+
+    def __init__(self, transport: httpx.AsyncBaseTransport, pipeline: MiddlewarePipeline) -> None:
+        self.transport = transport
+        self.pipeline = pipeline
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        if self.pipeline:
+            self.set_request_context_and_feature_usage(request)
+            response = await self.pipeline.send(request)
+            return response
+
+        response = await self.transport.handle_async_request(request)
+        return response
+
+    def set_request_context_and_feature_usage(self, request: httpx.Request) -> httpx.Request:
+
+        request_options = {}
+        options = request.headers.get('request_options', None)
+        if options:
+            request_options = json.loads(options)
+
+        context = GraphRequestContext(request_options, request.headers)
+        middleware = self.pipeline._first_middleware
+        while middleware:
+            if isinstance(middleware, RedirectHandler):
+                context.feature_usage = FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
+            if isinstance(middleware, RetryHandler):
+                context.feature_usage = FeatureUsageFlag.RETRY_HANDLER_ENABLED
+
+            middleware = middleware.next
+        request.context = context  #type: ignore
+        return request

src/msgraph_core/middleware/telemetry.py

@@ -3,11 +3,12 @@
 import platform
 
 import httpx
-from kiota_http.middleware import AsyncKiotaTransport, BaseMiddleware, RedirectHandler, RetryHandler
+from kiota_http.middleware import BaseMiddleware
 from urllib3.util import parse_url
 
 from .._constants import SDK_VERSION
-from .._enums import FeatureUsageFlag, NationalClouds
+from .._enums import NationalClouds
+from .async_graph_transport import AsyncGraphTransport
 from .request_context import GraphRequestContext
 
 
@@ -20,10 +21,9 @@ class GraphTelemetryHandler(BaseMiddleware):
     the SDK team improve the developer experience.
     """
 
-    async def send(self, request: GraphRequest, transport: AsyncKiotaTransport):
+    async def send(self, request: GraphRequest, transport: AsyncGraphTransport):
         """Adds telemetry headers and sends the http request.
         """
-        self.set_request_context_and_feature_usage(request, transport)
 
         if self.is_graph_url(request.url):
             self._add_client_request_id_header(request)
@@ -34,27 +34,6 @@ async def send(self, request: GraphRequest, transport: AsyncKiotaTransport):
         response = await super().send(request, transport)
         return response
 
-    def set_request_context_and_feature_usage(
-        self, request: GraphRequest, transport: AsyncKiotaTransport
-    ) -> GraphRequest:
-
-        request_options = {}
-        options = request.headers.pop('request_options', None)
-        if options:
-            request_options = json.loads(options)
-
-        request.context = GraphRequestContext(request_options, request.headers)
-        middleware = transport.pipeline._first_middleware
-        while middleware:
-            if isinstance(middleware, RedirectHandler):
-                request.context.feature_usage = FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
-            if isinstance(middleware, RetryHandler):
-                request.context.feature_usage = FeatureUsageFlag.RETRY_HANDLER_ENABLED
-
-            middleware = middleware.next
-
-        return request
-
     def is_graph_url(self, url):
         """Check if the request is made to a graph endpoint. We do not add telemetry headers to
         non-graph endpoints"""