diff --git a/sdk/confidentialledger/azure-confidentialledger/CHANGELOG.md b/sdk/confidentialledger/azure-confidentialledger/CHANGELOG.md index 22366f06a123..b1d4b8316f28 100644 --- a/sdk/confidentialledger/azure-confidentialledger/CHANGELOG.md +++ b/sdk/confidentialledger/azure-confidentialledger/CHANGELOG.md @@ -1,5 +1,15 @@ # Release History +## 2.0.0b3 (Unreleased) + +### Features Added + +- Added redirect URL caching for write operations (POST/PUT/PATCH/DELETE). After the first redirect from the load balancer, subsequent writes go directly to the primary node, significantly reducing latency. + +### Bugs Fixed + +- Fixed missing `disable_redirect_cleanup=True` in the async client, which could cause authentication failures on redirects. + ## 2.0.0b2 (2026-01-29) ### Bugs Fixed diff --git a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_client.py b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_client.py index 92370f080501..fa11e8b73f11 100644 --- a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_client.py +++ b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_client.py @@ -16,6 +16,7 @@ from ._configuration import ConfidentialLedgerClientConfiguration from ._operations import _ConfidentialLedgerClientOperationsMixin +from ._redirect_caching_policy import RedirectCachingPolicy from ._utils.serialization import Deserializer, Serializer @@ -46,7 +47,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self._config.user_agent_policy, self._config.proxy_policy, policies.ContentDecodePolicy(**kwargs), - self._config.redirect_policy, + kwargs.get("redirect_policy") or RedirectCachingPolicy(**kwargs), self._config.retry_policy, self._config.authentication_policy, self._config.custom_hook_policy, diff --git a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_redirect_caching_policy.py b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_redirect_caching_policy.py new file mode 100644 index 000000000000..fe1a2214d723 --- /dev/null +++ b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_redirect_caching_policy.py @@ -0,0 +1,261 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ + +"""Custom redirect policies that cache the redirect target URL for write operations. + +Write requests (POST / PUT / PATCH / DELETE) that receive a 307 redirect from +the load-balancer are followed, and the target URL's base (scheme + host) is +cached so that subsequent writes skip the load-balancer entirely. + +Read requests (GET, HEAD, OPTIONS, …) **never** consult or populate the cache +and always go through the load-balancer. + +The cache is invalidated on 5xx responses or transport errors so that a +failover is respected on the next write. + +Thread-safety +~~~~~~~~~~~~~ +* Reads of the cached value are lock-free (CPython GIL guarantees atomic + reference reads — *Volatile.Read* semantics). +* Writes are protected by a :class:`threading.Lock` (*Volatile.Write* + semantics) so that at most one thread mutates the cached reference at a + time. +""" + +import logging +import threading +from typing import FrozenSet, Optional +from urllib.parse import urlparse, urlunparse + +from azure.core.pipeline import PipelineRequest, PipelineResponse +from azure.core.pipeline.policies import AsyncHTTPPolicy, HTTPPolicy + +_LOGGER = logging.getLogger(__name__) + +_WRITE_METHODS: FrozenSet[str] = frozenset({"POST", "PUT", "PATCH", "DELETE"}) + +_REDIRECT_STATUS_CODES = frozenset({301, 302, 307, 308}) + + +class _RedirectUrlCache: + """Thread-safe cache for a redirect target base URL.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._cached_base_url: Optional[str] = None + + # -- Volatile.Read -------------------------------------------------- + def get(self) -> Optional[str]: + """Return the cached base URL or *None* (lock-free).""" + return self._cached_base_url + + # -- Volatile.Write ------------------------------------------------- + def set(self, url: str) -> None: # pylint: disable=redefined-builtin + """Extract and cache the base URL (scheme + host) from *url*. + + :param str url: The full redirect target URL. + """ + parsed = urlparse(url) + base_url = f"{parsed.scheme}://{parsed.netloc}" + with self._lock: + self._cached_base_url = base_url + + def invalidate(self) -> None: + """Clear the cached value.""" + with self._lock: + self._cached_base_url = None + + +def _rewrite_url(original_url: str, cached_base_url: str) -> str: + """Replace the scheme + host of *original_url* with *cached_base_url*. + + :param str original_url: The original request URL. + :param str cached_base_url: The cached base URL (scheme + host) to use. + :return: The rewritten URL. + :rtype: str + """ + original = urlparse(original_url) + cached = urlparse(cached_base_url) + return urlunparse( + ( + cached.scheme, + cached.netloc, + original.path, + original.params, + original.query, + original.fragment, + ) + ) + + +def _is_redirect(status_code: int) -> bool: + return status_code in _REDIRECT_STATUS_CODES + + +class RedirectCachingPolicy(HTTPPolicy): + """Synchronous redirect policy with write-URL caching. + + Replaces the default :class:`~azure.core.pipeline.policies.RedirectPolicy` + in the pipeline. See module docstring for caching semantics. + + :keyword bool permit_redirects: Whether redirects are followed at all. + Defaults to ``True``. + :keyword int redirect_max: Maximum number of redirects to follow per + request. Defaults to ``5``. + """ + + def __init__( + self, + *, + permit_redirects: bool = True, + redirect_max: int = 5, + **kwargs, # pylint: disable=unused-argument + ) -> None: + super().__init__() + self._permit_redirects = permit_redirects + self._max_redirects = redirect_max + self._cache = _RedirectUrlCache() + + def send(self, request: PipelineRequest) -> PipelineResponse: + method = request.http_request.method.upper() + is_write = method in _WRITE_METHODS + + # For writes, rewrite the URL to the cached primary (if warm). + if is_write: + cached = self._cache.get() + if cached: + request.http_request.url = _rewrite_url( + request.http_request.url, cached + ) + _LOGGER.debug( + "Using cached redirect URL for %s: %s", + method, + request.http_request.url, + ) + + # Send the request downstream. + try: + response = self.next.send(request) + except Exception: + if is_write: + self._cache.invalidate() + _LOGGER.debug("Transport error on write; invalidated redirect cache") + raise + + if not self._permit_redirects: + return response + + # Follow redirect chain. + redirects_remaining = self._max_redirects + while ( + _is_redirect(response.http_response.status_code) + and redirects_remaining > 0 + ): + redirect_url = response.http_response.headers.get("Location") + if not redirect_url: + break + + # Only cache for write methods. + if is_write: + self._cache.set(redirect_url) + _LOGGER.debug("Cached redirect target for writes: %s", redirect_url) + + request.http_request.url = redirect_url + redirects_remaining -= 1 + + try: + response = self.next.send(request) + except Exception: + if is_write: + self._cache.invalidate() + _LOGGER.debug( + "Transport error following redirect; invalidated cache" + ) + raise + + # Invalidate cache on server errors for writes. + if is_write and response.http_response.status_code >= 500: + self._cache.invalidate() + _LOGGER.debug("5xx on write; invalidated redirect cache") + + return response + + +class AsyncRedirectCachingPolicy(AsyncHTTPPolicy): + """Asynchronous redirect policy with write-URL caching. + + Async counterpart of :class:`RedirectCachingPolicy`. + """ + + def __init__( + self, + *, + permit_redirects: bool = True, + redirect_max: int = 5, + **kwargs, # pylint: disable=unused-argument + ) -> None: + super().__init__() + self._permit_redirects = permit_redirects + self._max_redirects = redirect_max + self._cache = _RedirectUrlCache() + + async def send(self, request: PipelineRequest) -> PipelineResponse: + method = request.http_request.method.upper() + is_write = method in _WRITE_METHODS + + if is_write: + cached = self._cache.get() + if cached: + request.http_request.url = _rewrite_url( + request.http_request.url, cached + ) + _LOGGER.debug( + "Using cached redirect URL for %s: %s", + method, + request.http_request.url, + ) + + try: + response = await self.next.send(request) + except Exception: + if is_write: + self._cache.invalidate() + _LOGGER.debug("Transport error on write; invalidated redirect cache") + raise + + if not self._permit_redirects: + return response + + redirects_remaining = self._max_redirects + while ( + _is_redirect(response.http_response.status_code) + and redirects_remaining > 0 + ): + redirect_url = response.http_response.headers.get("Location") + if not redirect_url: + break + + if is_write: + self._cache.set(redirect_url) + _LOGGER.debug("Cached redirect target for writes: %s", redirect_url) + + request.http_request.url = redirect_url + redirects_remaining -= 1 + + try: + response = await self.next.send(request) + except Exception: + if is_write: + self._cache.invalidate() + _LOGGER.debug( + "Transport error following redirect; invalidated cache" + ) + raise + + if is_write and response.http_response.status_code >= 500: + self._cache.invalidate() + _LOGGER.debug("5xx on write; invalidated redirect cache") + + return response diff --git a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_utils/model_base.py b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_utils/model_base.py index 12926fa98dcf..3cd9cb9ceabf 100644 --- a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_utils/model_base.py +++ b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_utils/model_base.py @@ -817,16 +817,16 @@ def _get_deserialize_callable_from_annotation( # pylint: disable=too-many-retur # is it optional? try: - if any(a for a in annotation.__args__ if a == type(None)): # pyright: ignore + if any(a for a in annotation.__args__ if a == type(None)): # pyright: ignore # pylint: disable=unidiomatic-typecheck if len(annotation.__args__) <= 2: # pyright: ignore if_obj_deserializer = _get_deserialize_callable_from_annotation( - next(a for a in annotation.__args__ if a != type(None)), module, rf # pyright: ignore + next(a for a in annotation.__args__ if a != type(None)), module, rf # pyright: ignore # pylint: disable=unidiomatic-typecheck ) return functools.partial(_deserialize_with_optional, if_obj_deserializer) # the type is Optional[Union[...]], we need to remove the None type from the Union annotation_copy = copy.copy(annotation) - annotation_copy.__args__ = [a for a in annotation_copy.__args__ if a != type(None)] # pyright: ignore + annotation_copy.__args__ = [a for a in annotation_copy.__args__ if a != type(None)] # pyright: ignore # pylint: disable=unidiomatic-typecheck return _get_deserialize_callable_from_annotation(annotation_copy, module, rf) except AttributeError: pass diff --git a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_version.py b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_version.py index 8eb37199ee54..2e7efc3e2e20 100644 --- a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_version.py +++ b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_version.py @@ -6,4 +6,4 @@ # Changes may cause incorrect behavior and will be lost if the code is regenerated. # -------------------------------------------------------------------------- -VERSION = "2.0.0b2" +VERSION = "2.0.0b3" diff --git a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/aio/_client.py b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/aio/_client.py index 7f49efc68499..7e465fb3869c 100644 --- a/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/aio/_client.py +++ b/sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/aio/_client.py @@ -14,6 +14,7 @@ from azure.core.pipeline import policies from azure.core.rest import AsyncHttpResponse, HttpRequest +from .._redirect_caching_policy import AsyncRedirectCachingPolicy from .._utils.serialization import Deserializer, Serializer from ._configuration import ConfidentialLedgerClientConfiguration from ._operations import _ConfidentialLedgerClientOperationsMixin @@ -46,7 +47,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self._config.user_agent_policy, self._config.proxy_policy, policies.ContentDecodePolicy(**kwargs), - self._config.redirect_policy, + kwargs.get("redirect_policy") or AsyncRedirectCachingPolicy(**kwargs), self._config.retry_policy, self._config.authentication_policy, self._config.custom_hook_policy, diff --git a/sdk/confidentialledger/azure-confidentialledger/tests/test_redirect_caching_policy.py b/sdk/confidentialledger/azure-confidentialledger/tests/test_redirect_caching_policy.py new file mode 100644 index 000000000000..4ab2293452c1 --- /dev/null +++ b/sdk/confidentialledger/azure-confidentialledger/tests/test_redirect_caching_policy.py @@ -0,0 +1,256 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ + +"""Unit tests for RedirectCachingPolicy and AsyncRedirectCachingPolicy.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from azure.confidentialledger._redirect_caching_policy import ( + AsyncRedirectCachingPolicy, + RedirectCachingPolicy, + _RedirectUrlCache, + _rewrite_url, +) + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _make_request(method: str = "POST", url: str = "https://lb.example.com/app/transactions"): + """Return a fake PipelineRequest.""" + http_request = MagicMock() + http_request.method = method + http_request.url = url + pipeline_request = MagicMock() + pipeline_request.http_request = http_request + return pipeline_request + + +def _make_response(status_code: int, headers=None): + """Return a fake PipelineResponse whose http_response has the given status and headers.""" + http_response = MagicMock() + http_response.status_code = status_code + http_response.headers = headers or {} + pipeline_response = MagicMock() + pipeline_response.http_response = http_response + return pipeline_response + + +# ── _RedirectUrlCache ──────────────────────────────────────────────────────── + + +class TestRedirectUrlCache: + def test_initial_value_is_none(self): + cache = _RedirectUrlCache() + assert cache.get() is None + + def test_set_and_get(self): + cache = _RedirectUrlCache() + cache.set("https://primary.example.com:443/some/path?q=1") + assert cache.get() == "https://primary.example.com:443" + + def test_invalidate_clears(self): + cache = _RedirectUrlCache() + cache.set("https://primary.example.com/path") + cache.invalidate() + assert cache.get() is None + + +# ── _rewrite_url ───────────────────────────────────────────────────────────── + + +class TestRewriteUrl: + def test_preserves_path_and_query(self): + result = _rewrite_url( + "https://lb.example.com/app/tx?api=v1", + "https://primary.example.com", + ) + assert result == "https://primary.example.com/app/tx?api=v1" + + def test_replaces_scheme_and_host(self): + result = _rewrite_url( + "http://old-host:8080/path", + "https://new-host:443", + ) + assert result == "https://new-host:443/path" + + +# ── RedirectCachingPolicy (sync) ───────────────────────────────────────────── + + +class TestRedirectCachingPolicy: + def _make_policy(self, responses): + """Create a policy with a mocked next node that returns *responses* in order.""" + policy = RedirectCachingPolicy() + policy.next = MagicMock() + policy.next.send = MagicMock(side_effect=responses) + return policy + + def test_cache_populated_on_first_redirect(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp]) + + request = _make_request("POST") + result = policy.send(request) + + assert result.http_response.status_code == 200 + assert policy._cache.get() == "https://primary.example.com" + + def test_subsequent_write_uses_cached_url(self): + # First request: redirect populates cache + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + # Second request: should go directly + final_resp2 = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp1, final_resp2]) + + policy.send(_make_request("POST")) + req2 = _make_request("POST", "https://lb.example.com/app/transactions") + policy.send(req2) + + # The request URL should have been rewritten to the cached primary + assert "primary.example.com" in req2.http_request.url + + def test_get_never_uses_cache(self): + # Warm the cache via a POST redirect + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + # GET request + final_resp2 = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp1, final_resp2]) + + policy.send(_make_request("POST")) + get_req = _make_request("GET", "https://lb.example.com/app/transactions") + policy.send(get_req) + + # GET should NOT have been rewritten + assert get_req.http_request.url == "https://lb.example.com/app/transactions" + + def test_5xx_invalidates_cache(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + error_resp = _make_response(503) + policy = self._make_policy([redirect_resp, final_resp1, error_resp]) + + policy.send(_make_request("POST")) + assert policy._cache.get() is not None + + policy.send(_make_request("POST")) + assert policy._cache.get() is None + + def test_transport_error_invalidates_cache(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp, ConnectionError("connection lost")]) + + policy.send(_make_request("POST")) + assert policy._cache.get() is not None + + with pytest.raises(ConnectionError): + policy.send(_make_request("POST")) + assert policy._cache.get() is None + + def test_delete_is_a_write_method(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp]) + + policy.send(_make_request("DELETE")) + assert policy._cache.get() == "https://primary.example.com" + + def test_permit_redirects_false_skips_redirect(self): + policy = RedirectCachingPolicy(permit_redirects=False) + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + policy.next = MagicMock() + policy.next.send = MagicMock(return_value=redirect_resp) + + result = policy.send(_make_request("POST")) + # Should return the 307 without following + assert result.http_response.status_code == 307 + assert policy._cache.get() is None + + +# ── AsyncRedirectCachingPolicy ─────────────────────────────────────────────── + + +class TestAsyncRedirectCachingPolicy: + def _make_policy(self, responses): + """Create an async policy with a mocked next node.""" + policy = AsyncRedirectCachingPolicy() + policy.next = MagicMock() + side_effects = [] + for r in responses: + if isinstance(r, Exception): + side_effects.append(r) + else: + side_effects.append(r) + policy.next.send = AsyncMock(side_effect=side_effects) + return policy + + @pytest.mark.asyncio + async def test_cache_populated_on_first_redirect(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp]) + + result = await policy.send(_make_request("POST")) + assert result.http_response.status_code == 200 + assert policy._cache.get() == "https://primary.example.com" + + @pytest.mark.asyncio + async def test_subsequent_write_uses_cached_url(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + final_resp2 = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp1, final_resp2]) + + await policy.send(_make_request("POST")) + req2 = _make_request("POST", "https://lb.example.com/app/transactions") + await policy.send(req2) + + assert "primary.example.com" in req2.http_request.url + + @pytest.mark.asyncio + async def test_get_never_uses_cache(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + final_resp2 = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp1, final_resp2]) + + await policy.send(_make_request("POST")) + get_req = _make_request("GET", "https://lb.example.com/app/transactions") + await policy.send(get_req) + + assert get_req.http_request.url == "https://lb.example.com/app/transactions" + + @pytest.mark.asyncio + async def test_5xx_invalidates_cache(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp1 = _make_response(200) + error_resp = _make_response(503) + policy = self._make_policy([redirect_resp, final_resp1, error_resp]) + + await policy.send(_make_request("POST")) + assert policy._cache.get() is not None + + await policy.send(_make_request("POST")) + assert policy._cache.get() is None + + @pytest.mark.asyncio + async def test_transport_error_invalidates_cache(self): + redirect_resp = _make_response(307, {"Location": "https://primary.example.com/app/transactions"}) + final_resp = _make_response(200) + policy = self._make_policy([redirect_resp, final_resp, ConnectionError("connection lost")]) + + await policy.send(_make_request("POST")) + assert policy._cache.get() is not None + + with pytest.raises(ConnectionError): + await policy.send(_make_request("POST")) + assert policy._cache.get() is None