Merge pull request #1129 from rksharma-owg/codex/graph-kiota-request-options

fix: honor Kiota request extensions in Graph transport

Author: baywet

pyproject.toml

@@ -11,9 +11,9 @@ version = "1.5.1"
 authors = [{name = "Microsoft", email = "graphtooling+python@microsoft.com"}]
 description = "Core component of the Microsoft Graph Python SDK"
 dependencies = [
-    "microsoft-kiota-abstractions >=1.11.6,<2.0.0",
-    "microsoft-kiota-authentication-azure >=1.11.6,<2.0.0",
-    "microsoft-kiota-http >=1.11.6,<2.0.0",
+    "microsoft-kiota-abstractions >=1.13.0,<2.0.0",
+    "microsoft-kiota-authentication-azure >=1.13.0,<2.0.0",
+    "microsoft-kiota-http >=1.13.0,<2.0.0",
     "httpx[http2] >=0.23.0",
 ]
 requires-python = ">=3.10"

requirements-dev.txt

@@ -144,13 +144,13 @@ httpx[http2]==0.28.1
 
 hyperframe==6.1.0 ; python_full_version >= '3.6.1'
 
-microsoft-kiota-abstractions==1.11.6
+microsoft-kiota-abstractions==1.13.0
 
-microsoft-kiota-authentication-azure==1.11.6
+microsoft-kiota-authentication-azure==1.13.0
 
-microsoft-kiota-http==1.11.6
+microsoft-kiota-http==1.13.0
 
-microsoft-kiota-serialization-json==1.11.6
+microsoft-kiota-serialization-json==1.13.0
 
 multidict==6.7.1 ; python_version >= '3.7'
 

src/msgraph_core/middleware/async_graph_transport.py

@@ -6,6 +6,9 @@
 from .._enums import FeatureUsageFlag
 from .request_context import GraphRequestContext
 
+# Older supported Kiota releases do not export this request extension key.
+REQUEST_OPTIONS_KEY = 'kiota_request_options'
+
 
 class AsyncGraphTransport(httpx.AsyncBaseTransport):
     """A custom transport for requests to the Microsoft Graph API
@@ -16,7 +19,9 @@ def __init__(self, transport: httpx.AsyncBaseTransport, pipeline: MiddlewarePipe
         self.pipeline = pipeline
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
-        if self.pipeline and hasattr(request, 'options'):
+        if self.pipeline and (
+            REQUEST_OPTIONS_KEY in request.extensions or hasattr(request, 'options')
+        ):
             self.set_request_context_and_feature_usage(request)
             response = await self.pipeline.send(request)
             return response
@@ -26,7 +31,11 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
 
     def set_request_context_and_feature_usage(self, request: httpx.Request) -> httpx.Request:
 
-        request_options = request.options  # type:ignore
+        if REQUEST_OPTIONS_KEY in request.extensions:
+            request_options = request.extensions[REQUEST_OPTIONS_KEY]
+        else:
+            request_options = request.options  # type:ignore
+            request.extensions[REQUEST_OPTIONS_KEY] = request_options
 
         context = GraphRequestContext(request_options, request.headers)
         middleware = self.pipeline._first_middleware

tests/middleware/test_async_graph_transport.py

@@ -1,8 +1,18 @@
+import asyncio
+
+import httpx
 import pytest
+from kiota_abstractions.authentication import AnonymousAuthenticationProvider
+from kiota_abstractions.method import Method
+from kiota_abstractions.request_information import RequestInformation
+from kiota_http.httpx_request_adapter import HttpxRequestAdapter
 from kiota_http.kiota_client_factory import KiotaClientFactory
+from kiota_http.middleware.options import RedirectHandlerOption
 
 from msgraph_core._enums import FeatureUsageFlag
+from msgraph_core.graph_client_factory import GraphClientFactory
 from msgraph_core.middleware import AsyncGraphTransport, GraphRequestContext
+from msgraph_core.middleware.async_graph_transport import REQUEST_OPTIONS_KEY
 
 
 def test_set_request_context_and_feature_usage(mock_request, mock_transport):
@@ -16,3 +26,139 @@ def test_set_request_context_and_feature_usage(mock_request, mock_transport):
     assert mock_request.context.feature_usage == hex(
         FeatureUsageFlag.RETRY_HANDLER_ENABLED | FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
     )
+
+
+@pytest.mark.parametrize(
+    'content_type', [
+        'application/octet-stream',
+        'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+    ]
+)
+def test_binary_download_follows_redirect_with_kiota_request_extensions(content_type):
+    calls = []
+    contexts = []
+
+    def handle_request(request):
+        calls.append(str(request.url))
+        contexts.append(request.context)
+        if request.url.host == 'graph.example':
+            return httpx.Response(302, headers={'Location': 'https://download.example/file'})
+        return httpx.Response(
+            200, content=b'binary content', headers={'Content-Type': content_type}
+        )
+
+    async def download():
+        client = GraphClientFactory.create_with_default_middleware(
+            client=httpx.AsyncClient(transport=httpx.MockTransport(handle_request))
+        )
+        try:
+            adapter = HttpxRequestAdapter(AnonymousAuthenticationProvider(), http_client=client)
+            request_info = RequestInformation()
+            request_info.http_method = Method.GET
+            request_info.url = 'https://graph.example/drive/item/content'
+            return await adapter.send_primitive_async(request_info, 'bytes', {})
+        finally:
+            await client.aclose()
+
+    assert asyncio.run(download()) == b'binary content'
+    assert calls == ['https://graph.example/drive/item/content', 'https://download.example/file']
+    assert all(isinstance(context, GraphRequestContext) for context in contexts)
+
+
+def test_extension_options_take_precedence_over_legacy_attribute(mock_transport):
+    middleware = KiotaClientFactory.get_default_middleware(None)
+    pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, mock_transport)
+    transport = AsyncGraphTransport(mock_transport, pipeline)
+    request = httpx.Request('GET', 'https://example.org', extensions={REQUEST_OPTIONS_KEY: {}})
+    request.options = {'legacy': True}
+
+    transport.set_request_context_and_feature_usage(request)
+
+    assert request.context.middleware_control == {}
+
+
+def test_request_without_options_bypasses_graph_pipeline():
+    calls = []
+
+    def handle_request(request):
+        calls.append(request)
+        return httpx.Response(200, content=b'body')
+
+    async def send():
+        underlying_transport = httpx.MockTransport(handle_request)
+        middleware = KiotaClientFactory.get_default_middleware(None)
+        pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, underlying_transport)
+        transport = AsyncGraphTransport(underlying_transport, pipeline)
+        return await transport.handle_async_request(httpx.Request('GET', 'https://example.org'))
+
+    assert asyncio.run(send()).status_code == 200
+    assert len(calls) == 1
+    assert not hasattr(calls[0], 'context')
+
+
+def test_extension_only_request_uses_graph_pipeline():
+    requests = []
+
+    def handle_request(request):
+        requests.append(request)
+        return httpx.Response(200, content=b'body')
+
+    async def send():
+        underlying_transport = httpx.MockTransport(handle_request)
+        middleware = KiotaClientFactory.get_default_middleware(None)
+        pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, underlying_transport)
+        transport = AsyncGraphTransport(underlying_transport, pipeline)
+        request = httpx.Request('GET', 'https://example.org', extensions={REQUEST_OPTIONS_KEY: {}})
+        assert not hasattr(request, 'options')
+        return await transport.handle_async_request(request)
+
+    assert asyncio.run(send()).content == b'body'
+    assert len(requests) == 1
+    assert isinstance(requests[0].context, GraphRequestContext)
+    assert requests[0].context.middleware_control == {}
+    assert requests[0].context.feature_usage == hex(
+        FeatureUsageFlag.RETRY_HANDLER_ENABLED | FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
+    )
+
+
+@pytest.mark.parametrize(
+    'legacy_redirect, extension_options, expected_status, expected_calls', [
+        (False, None, 302, 1),
+        (None, None, 200, 2),
+        (True, False, 302, 1),
+        (False, True, 200, 2),
+        (False, {}, 200, 2),
+    ]
+)
+def test_request_redirect_options_are_honored(
+    legacy_redirect, extension_options, expected_status, expected_calls
+):
+    requests = []
+
+    def handle_request(request):
+        requests.append(request)
+        if request.url.path == '/start':
+            return httpx.Response(302, headers={'Location': 'https://example.org/end'})
+        return httpx.Response(200, content=b'body')
+
+    async def send():
+        underlying_transport = httpx.MockTransport(handle_request)
+        middleware = KiotaClientFactory.get_default_middleware(None)
+        pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, underlying_transport)
+        transport = AsyncGraphTransport(underlying_transport, pipeline)
+        request = httpx.Request('GET', 'https://example.org/start')
+        request.options = {}
+        if legacy_redirect is not None:
+            option = RedirectHandlerOption(should_redirect=legacy_redirect)
+            request.options[option.get_key()] = option
+        if extension_options is not None:
+            options = extension_options
+            if isinstance(options, bool):
+                option = RedirectHandlerOption(should_redirect=options)
+                options = {option.get_key(): option}
+            request.extensions[REQUEST_OPTIONS_KEY] = options
+        return await transport.handle_async_request(request)
+
+    assert asyncio.run(send()).status_code == expected_status
+    assert len(requests) == expected_calls
+    assert isinstance(requests[0].context, GraphRequestContext)