diff --git a/packages/http/httpx/kiota_http/middleware/middleware.py b/packages/http/httpx/kiota_http/middleware/middleware.py index 97e713f7..17142649 100644 --- a/packages/http/httpx/kiota_http/middleware/middleware.py +++ b/packages/http/httpx/kiota_http/middleware/middleware.py @@ -19,43 +19,40 @@ class MiddlewarePipeline(): it here https://buffered.dev/middleware-python-requests/ """ - def __init__(self, transport: httpx.AsyncBaseTransport): + def __init__(self, transport: httpx.AsyncBaseTransport) -> None: super().__init__() - self._current_middleware = None - self._first_middleware = None + self._current_middleware: Optional[BaseMiddleware] = None + self._first_middleware: Optional[BaseMiddleware] = None self._transport = transport - def add_middleware(self, middleware): - if self._middleware_present(): - self._current_middleware.next = middleware - self._current_middleware = middleware - else: + def add_middleware(self, middleware: "BaseMiddleware") -> None: + if self._current_middleware is None: self._first_middleware = middleware - self._current_middleware = self._first_middleware + else: + self._current_middleware.next = middleware + self._current_middleware = middleware - async def send(self, request): + async def send(self, request: httpx.Request) -> httpx.Response: - if self._middleware_present(): + if self._first_middleware is not None: return await self._first_middleware.send(request, self._transport) - # No middleware in pipeline, delete request optoions from header and - # send the request - del request.headers['request_options'] + # No middleware in pipeline, remove the request options and send the request + request.extensions.pop(REQUEST_OPTIONS_KEY, None) return await self._transport.handle_async_request(request) - def _middleware_present(self): - return self._current_middleware - class BaseMiddleware(): """Base class for middleware. Handles moving a Request to the next middleware in the pipeline. If the current middleware is the last one in the pipeline, it makes a network request """ - def __init__(self): - self.next = None - self.parent_span = None + def __init__(self) -> None: + self.next: Optional[BaseMiddleware] = None + self.parent_span: Optional[trace.Span] = None - async def send(self, request, transport): + async def send( + self, request: httpx.Request, transport: httpx.AsyncBaseTransport + ) -> httpx.Response: if self.next is None: # Remove request options if there's no other middleware in the chain. if hasattr(request, "extensions") and isinstance(request.extensions, dict): @@ -65,19 +62,15 @@ async def send(self, request, transport): return response return await self.next.send(request, transport) - def _create_observability_span(self, request, span_name: str) -> trace.Span: + def _create_observability_span(self, request: httpx.Request, span_name: str) -> trace.Span: """Gets the parent_span from the request options and creates a new span. If no parent_span is found in the request, uses the parent_span in the object. If parent_span is None, current context will be used.""" - _span = None options = None if hasattr(request, "extensions") and isinstance(request.extensions, dict): options = request.extensions.get(REQUEST_OPTIONS_KEY) if options and (parent_span := options.get("parent_span", None)): self.parent_span = parent_span - _context = trace.set_span_in_context(parent_span) - _span = tracer.start_span(span_name, _context) - if _span is None: - _context = trace.set_span_in_context(self.parent_span) - _span = tracer.start_span(span_name, _context) - return _span + if self.parent_span is None: + return tracer.start_span(span_name) + return tracer.start_span(span_name, trace.set_span_in_context(self.parent_span)) diff --git a/packages/http/httpx/kiota_http/middleware/redirect_handler.py b/packages/http/httpx/kiota_http/middleware/redirect_handler.py index 62462b9d..a1e7d5fb 100644 --- a/packages/http/httpx/kiota_http/middleware/redirect_handler.py +++ b/packages/http/httpx/kiota_http/middleware/redirect_handler.py @@ -84,8 +84,7 @@ async def send( max_redirect -= 1 if not self.increment(response, max_redirect, history[:]): if max_redirect < 0: - response.history = history - exc = RedirectError(f"Too many redirects. {response.history}") + exc = RedirectError(f"Too many redirects. {history}") _redirect_span.record_exception(exc) raise exc break @@ -100,7 +99,6 @@ async def send( break finally: _redirect_span.end() - response.history = history return response diff --git a/packages/http/httpx/kiota_http/middleware/retry_handler.py b/packages/http/httpx/kiota_http/middleware/retry_handler.py index 43a1cdab..de1aa37c 100644 --- a/packages/http/httpx/kiota_http/middleware/retry_handler.py +++ b/packages/http/httpx/kiota_http/middleware/retry_handler.py @@ -3,6 +3,7 @@ import random import re from email.utils import parsedate_to_datetime +from typing import Optional from kiota_abstractions.request_option import RequestOption from opentelemetry.semconv.attributes.http_attributes import HTTP_RESPONSE_STATUS_CODE @@ -42,16 +43,18 @@ class RetryHandler(BaseMiddleware): ['HEAD', 'GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS', 'QUERY'] ) - def __init__(self, options: RetryHandlerOption = RetryHandlerOption()) -> None: + def __init__(self, options: Optional[RetryHandlerOption] = None) -> None: super().__init__() self.allowed_methods: frozenset[str] = self.DEFAULT_ALLOWED_METHODS self.backoff_factor: float = self.DEFAULT_BACKOFF_FACTOR self.backoff_max: int = self.MAXIMUM_BACKOFF - self.options = options - self.respect_retry_after_header: bool = self.options.DEFAULT_SHOULD_RETRY # type:ignore + self.options = RetryHandlerOption() if options is None else options + self.respect_retry_after_header: bool = self.options.DEFAULT_SHOULD_RETRY self.retry_on_status_codes: set[int] = self.DEFAULT_RETRY_STATUS_CODES - async def send(self, request: httpx.Request, transport: httpx.AsyncBaseTransport): + async def send( + self, request: httpx.Request, transport: httpx.AsyncBaseTransport + ) -> httpx.Response: """ Sends the http request object to the next middleware or retries the request if necessary. """ @@ -88,6 +91,7 @@ async def send(self, request: httpx.Request, transport: httpx.AsyncBaseTransport request.headers.update({RETRY_ATTEMPT: f'{retry_count}'}) _retry_span.set_attribute(HTTP_RESPONSE_STATUS_CODE, response.status_code) _retry_span.set_attribute('http.request.resend_count', retry_count) + _retry_span.end() continue _retry_span.end() break @@ -107,11 +111,15 @@ def _get_current_options(self, request: httpx.Request) -> RetryHandlerOption: """ request_options = request.extensions.get(REQUEST_OPTIONS_KEY) if request_options: - current_options = request_options.get(RetryHandlerOption.get_key(), self.options) + current_options: RetryHandlerOption = request_options.get( + RetryHandlerOption.get_key(), self.options + ) return current_options return self.options - def should_retry(self, request, options, response): + def should_retry( + self, request: httpx.Request, options: RetryHandlerOption, response: httpx.Response + ) -> bool: """ Determines whether the request should be retried Checks if the request method is in allowed methods @@ -121,10 +129,9 @@ def should_retry(self, request, options, response): return False if not self._is_request_payload_buffered(request): return False - val = options.max_retry and (response.status_code in self.retry_on_status_codes) - return val + return options.max_retry > 0 and response.status_code in self.retry_on_status_codes - def _is_method_retryable(self, request): + def _is_method_retryable(self, request: httpx.Request) -> bool: """ Checks if a given request should be retried upon, depending on whether the HTTP method is in the set of allowed methods @@ -133,7 +140,7 @@ def _is_method_retryable(self, request): return False return True - def _is_request_payload_buffered(self, request): + def _is_request_payload_buffered(self, request: httpx.Request) -> bool: """ Checks if the request payload is buffered/rewindable. Payloads with forward only streams will return false and have the responses @@ -145,7 +152,7 @@ def _is_request_payload_buffered(self, request): return False return True - def check_retry_valid(self, retry_count, options): + def check_retry_valid(self, retry_count: int, options: RetryHandlerOption) -> bool: """ Check that the max retries limit has not been hit """ @@ -153,7 +160,12 @@ def check_retry_valid(self, retry_count, options): return True return False - def get_delay_time(self, retry_count, response=None, delay=RetryHandlerOption.DEFAULT_DELAY): + def get_delay_time( + self, + retry_count: int, + response: Optional[httpx.Response] = None, + delay: float = RetryHandlerOption.DEFAULT_DELAY + ) -> float: """Get the time in seconds to delay before the next attempt. A ``Retry-After`` response header that parses to more than zero seconds is @@ -167,37 +179,41 @@ def get_delay_time(self, retry_count, response=None, delay=RetryHandlerOption.DE return retry_after return self._get_delay_time_exp_backoff(retry_count, delay) - def _get_delay_time_exp_backoff(self, retry_count, delay): + def _get_delay_time_exp_backoff(self, retry_count: int, delay: float) -> float: """ Get time in seconds to delay between retry attempts based on an exponential backoff value. """ - exp_backoff_value = self.backoff_factor * +(2**(retry_count - 1)) + exp_backoff_value: float = self.backoff_factor * +(2**(retry_count - 1)) backoff_value = exp_backoff_value + (random.randint(0, 1000) / 1000) + delay backoff = min(self.backoff_max, backoff_value) return backoff - def _get_retry_after(self, response): + def _get_retry_after(self, response: Optional[httpx.Response]) -> Optional[float]: """ Check if retry-after is specified in the response header and get the value """ + if response is None: + return None retry_after = response.headers.get("retry-after") if retry_after: return self._parse_retry_after(retry_after) return None - def _parse_retry_after(self, retry_after): + def _parse_retry_after(self, retry_after: str) -> Optional[float]: """ - Helper to parse Retry-After and get value in seconds. + Helper to parse Retry-After and get value in seconds. Returns None when the + value is neither a number of seconds nor an HTTP date. """ + if re.match(r"\d+,\d+$", retry_after): + retry_after = retry_after.split(",")[0] try: - retry_after = retry_after.split(",")[0] if re.match( - r"\d+,\d+$", retry_after - ) else retry_after - delay = int(retry_after) + return max(0, int(retry_after)) except ValueError: - # Not an integer? Try HTTP date + pass + try: retry_date = parsedate_to_datetime(retry_after) - delay = (retry_date - datetime.datetime.now(retry_date.tzinfo)).total_seconds() - return max(0, delay) + except ValueError: + return None + return max(0.0, (retry_date - datetime.datetime.now(retry_date.tzinfo)).total_seconds()) diff --git a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py index 2f9ed794..ec1983e4 100644 --- a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py +++ b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py @@ -1,7 +1,10 @@ """Test the BaseMiddleware class.""" +import httpx +import pytest from opentelemetry import trace -from kiota_http.middleware import BaseMiddleware +from kiota_http.middleware import REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline +from kiota_http.middleware import middleware as middleware_module def test_next_is_none(): @@ -27,4 +30,24 @@ def test_span_created(request_info, otel_tracer): span = middleware._create_observability_span(request_info, "test_span_created") span.end() after_state = outside_span.is_recording() - assert(before_state == after_state) \ No newline at end of file + assert(before_state == after_state) + + +def test_span_without_parent_span_uses_the_current_span(span_exporter): + middleware = BaseMiddleware() + with middleware_module.tracer.start_as_current_span("outer"): + middleware._create_observability_span(httpx.Request("GET", "https://example.com"), "inner").end() + inner, outer = span_exporter.get_finished_spans() + assert inner.parent.span_id == outer.context.span_id + + +@pytest.mark.asyncio +async def test_pipeline_without_middleware_sends_the_request(): + def request_handler(request: httpx.Request): + return httpx.Response(200) + + request = httpx.Request("GET", "https://example.com", extensions={REQUEST_OPTIONS_KEY: {}}) + async with httpx.MockTransport(request_handler) as transport: + response = await MiddlewarePipeline(transport).send(request) + assert response.status_code == 200 + assert REQUEST_OPTIONS_KEY not in request.extensions diff --git a/packages/http/httpx/tests/middleware_tests/test_redirect_handler.py b/packages/http/httpx/tests/middleware_tests/test_redirect_handler.py index 46f3fa48..46313fe3 100644 --- a/packages/http/httpx/tests/middleware_tests/test_redirect_handler.py +++ b/packages/http/httpx/tests/middleware_tests/test_redirect_handler.py @@ -37,7 +37,6 @@ def request_handler(request): attempts = redirects + 1 if should_redirect else 1 assert len(requests) == attempts assert response.status_code == (200 if should_redirect else 302) - assert len(response.history) == attempts - 1 spans = span_exporter.get_finished_spans() assert [span.name for span in spans] == ["RedirectHandler_send"] + [ f"RedirectHandler_send - redirect {index}" for index in range(attempts) @@ -124,6 +123,27 @@ def test_increment_redirects(): assert handler.increment(response, handler.options.max_redirect, history=history) +class SingleRedirectHandler(RedirectHandler): + + def increment(self, response, max_redirect, history): + return len(history) == 0 + + +@pytest.mark.asyncio +async def test_increment_override_stops_redirects(): + requests = [] + + def request_handler(request): + requests.append(request) + return httpx.Response(FOUND, headers={LOCATION_HEADER: f"/redirect/{len(requests)}"}) + + async with httpx.MockTransport(request_handler) as transport: + response = await SingleRedirectHandler().send(httpx.Request("GET", BASE_URL), transport) + + assert response.status_code == FOUND + assert [request.url.path for request in requests] == ["/", "/redirect/1"] + + def test_same_origin(mock_redirect_handler): origin1 = httpx.URL(BASE_URL) origin2 = httpx.URL(f"{BASE_URL}:443") diff --git a/packages/http/httpx/tests/middleware_tests/test_retry_handler.py b/packages/http/httpx/tests/middleware_tests/test_retry_handler.py index e6c49849..1cbe07ed 100644 --- a/packages/http/httpx/tests/middleware_tests/test_retry_handler.py +++ b/packages/http/httpx/tests/middleware_tests/test_retry_handler.py @@ -345,3 +345,34 @@ def request_handler(request: httpx.Request): assert responses[0].is_closed and responses[1].is_closed # the returned response is left open for the caller to read assert not responses[2].is_closed + + +@pytest.mark.asyncio +async def test_retry_spans_end(span_exporter): + def request_handler(request: httpx.Request): + if RETRY_ATTEMPT in request.headers: + return httpx.Response(200) + return httpx.Response(SERVICE_UNAVAILABLE) + + async with httpx.MockTransport(request_handler) as transport: + await RetryHandler(RetryHandlerOption(delay=0)).send(httpx.Request("GET", BASE_URL), transport) + assert [span.name for span in span_exporter.get_finished_spans()] == [ + "RetryHandler_send", "RetryHandler_send - attempt 0", "RetryHandler_send - attempt 1" + ] + + +def test_default_options_are_not_shared(): + first, second = RetryHandler(), RetryHandler() + first.options.max_retry = 7 + assert second.options.max_retry == RetryHandlerOption.DEFAULT_MAX_RETRIES + + +def test_get_delay_time_without_response(): + delay = RetryHandler().get_delay_time(1, delay=0) + assert 0.5 <= delay <= 1.5 + + +def test_get_retry_after_invalid_value(): + response = httpx.Response(SERVICE_UNAVAILABLE, headers={RETRY_AFTER: "soon"}) + assert RetryHandler()._get_retry_after(response) is None +