Support redirect in live metrics (#35910)

Author: lzchen

sdk/monitor/azure-monitor-opentelemetry-exporter/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- Implement redirect for live metrics
+    ([#35910](https://github.com/Azure/azure-sdk-for-python/pull/35910))
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_quickpulse/_constants.py

@@ -52,4 +52,11 @@ class _DocumentIngressDocumentType(Enum):
     Event = "Event"
     Trace = "Trace"
 
+# Response Headers
+
+_QUICKPULSE_ETAG_HEADER_NAME = "x-ms-qps-configuration-etag"
+_QUICKPULSE_POLLING_HEADER_NAME = "x-ms-qps-service-polling-interval-hint"
+_QUICKPULSE_REDIRECT_HEADER_NAME = "x-ms-qps-service-endpoint-redirect-v2"
+_QUICKPULSE_SUBSCRIBED_HEADER_NAME = "x-ms-qps-subscribed"
+
 # cSpell:enable

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_quickpulse/_exporter.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 import logging
 from typing import Any, Optional
+import weakref
 
 from opentelemetry.context import (
     _SUPPRESS_INSTRUMENTATION_KEY,
@@ -28,10 +29,12 @@
     _LONG_PING_INTERVAL_SECONDS,
     _POST_CANCEL_INTERVAL_SECONDS,
     _POST_INTERVAL_SECONDS,
+    _QUICKPULSE_SUBSCRIBED_HEADER_NAME,
 )
 from azure.monitor.opentelemetry.exporter._quickpulse._generated._configuration import QuickpulseClientConfiguration
 from azure.monitor.opentelemetry.exporter._quickpulse._generated._client import QuickpulseClient
 from azure.monitor.opentelemetry.exporter._quickpulse._generated.models import MonitoringDataPoint
+from azure.monitor.opentelemetry.exporter._quickpulse._policy import _QuickpulseRedirectPolicy
 from azure.monitor.opentelemetry.exporter._quickpulse._state import (
     _get_global_quickpulse_state,
     _is_ping_state,
@@ -85,9 +88,10 @@ def __init__(self, connection_string: Optional[str]) -> None:
         # TODO: Support AADaudience (scope)/credentials
         # Pass `None` for now until swagger definition is fixed
         config = QuickpulseClientConfiguration(credential=None)  # type: ignore
+        qp_redirect_policy = _QuickpulseRedirectPolicy(permit_redirects=False)
         policies = [
-            # TODO: Support redirect
-            config.redirect_policy,
+            # Custom redirect policy for QP
+            qp_redirect_policy,
             # Needed for serialization
             ContentDecodePolicy(),
             # Logging for client calls
@@ -102,6 +106,9 @@ def __init__(self, connection_string: Optional[str]) -> None:
             endpoint=self._live_endpoint,
             policies=policies
         )
+        # Create a weakref of the client to the redirect policy so the endpoint can be
+        # dynamically modified if redirect does occur
+        qp_redirect_policy._qp_client_ref = weakref.ref(self._client)
 
         MetricExporter.__init__(
             self,
@@ -146,7 +153,7 @@ def export(
                 # If no response, assume unsuccessful
                 result = MetricExportResult.FAILURE
             else:
-                header = post_response._response_headers.get("x-ms-qps-subscribed")  # pylint: disable=protected-access
+                header = post_response._response_headers.get(_QUICKPULSE_SUBSCRIBED_HEADER_NAME)  # pylint: disable=protected-access
                 if header != "true":
                     # User leaving the live metrics page will be treated as an unsuccessful
                     result = MetricExportResult.FAILURE
@@ -240,7 +247,7 @@ def _ticker(self) -> None:
                     self._base_monitoring_data_point,
                 )
                 if ping_response:
-                    header = ping_response._response_headers.get("x-ms-qps-subscribed")  # pylint: disable=protected-access
+                    header = ping_response._response_headers.get(_QUICKPULSE_SUBSCRIBED_HEADER_NAME)  # pylint: disable=protected-access
                     if header and header == "true":
                         # Switch state to post if subscribed
                         _set_global_quickpulse_state(_QuickpulseState.POST_SHORT)

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_quickpulse/_policy.py

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+from typing import Any, Optional
+from urllib.parse import urlparse
+from weakref import ReferenceType
+
+from azure.core.pipeline import PipelineResponse, policies
+
+from azure.monitor.opentelemetry.exporter._quickpulse._constants import _QUICKPULSE_REDIRECT_HEADER_NAME
+from azure.monitor.opentelemetry.exporter._quickpulse._generated import QuickpulseClient
+
+
+# Quickpulse endpoint handles redirects via header instead of status codes
+# We use a custom RedirectPolicy to handle this use case
+# pylint: disable=protected-access
+class _QuickpulseRedirectPolicy(policies.RedirectPolicy):
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Weakref to QuickPulseClient instance
+        self._qp_client_ref: Optional[ReferenceType[QuickpulseClient]] = None
+        super().__init__(**kwargs)
+
+    # Gets the redirect location from header
+    def get_redirect_location(self, response: PipelineResponse) -> Optional[str] :
+        redirect_location = response.http_response.headers.get(_QUICKPULSE_REDIRECT_HEADER_NAME)
+        qp_client = None
+        if redirect_location:
+            redirected_url = urlparse(redirect_location)
+            if redirected_url.scheme and redirected_url.netloc:
+                if self._qp_client_ref:
+                    qp_client = self._qp_client_ref()
+                if qp_client and qp_client._client:
+                    # Set new endpoint to redirect location
+                    qp_client._client._base_url = f"{redirected_url.scheme}://{redirected_url.netloc}"
+        return redirect_location  # type: ignore

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/quickpulse/test_policy.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+import weakref
+import unittest
+from unittest import mock
+
+from azure.monitor.opentelemetry.exporter._quickpulse._policy import _QuickpulseRedirectPolicy
+
+
+class TestQuickpulseRedirectPolicy(unittest.TestCase):
+
+    def test_get_redirect_location(self):
+        policy = _QuickpulseRedirectPolicy()
+        pipeline_resp_mock = mock.Mock()
+        http_resp_mock = mock.Mock()
+        headers = {
+            "x-ms-qps-service-endpoint-redirect-v2": "https://eastus.livediagnostics.monitor.azure.com/QuickPulseService.svc"
+        }
+        http_resp_mock.headers = headers
+        pipeline_resp_mock.http_response = http_resp_mock
+        policy = _QuickpulseRedirectPolicy()
+        qp_client_mock = mock.Mock()
+        client_mock = mock.Mock()
+        client_mock._base_url = "previous_url"
+        qp_client_mock._client = client_mock
+        qp_client_ref = weakref.ref(qp_client_mock)
+        policy._qp_client_ref = qp_client_ref
+        redirect_location = policy.get_redirect_location(pipeline_resp_mock)
+        self.assertEqual(
+            redirect_location,
+            "https://eastus.livediagnostics.monitor.azure.com/QuickPulseService.svc",
+        )
+        self.assertEqual(
+            client_mock._base_url,
+            "https://eastus.livediagnostics.monitor.azure.com",
+        )
+
+    def test_get_redirect_location_no_header(self):
+        policy = _QuickpulseRedirectPolicy()
+        pipeline_resp_mock = mock.Mock()
+        http_resp_mock = mock.Mock()
+        headers = {}
+        http_resp_mock.headers = headers
+        pipeline_resp_mock.http_response = http_resp_mock
+        policy = _QuickpulseRedirectPolicy()
+        self.assertIsNone(policy.get_redirect_location(pipeline_resp_mock))
+        redirect_location = policy.get_redirect_location(pipeline_resp_mock)
+ 
+    def test_get_redirect_location_invalid_url(self):
+        policy = _QuickpulseRedirectPolicy()
+        pipeline_resp_mock = mock.Mock()
+        http_resp_mock = mock.Mock()
+        headers = {
+            "x-ms-qps-service-endpoint-redirect-v2": "invalid_url"
+        }
+        http_resp_mock.headers = headers
+        pipeline_resp_mock.http_response = http_resp_mock
+        policy = _QuickpulseRedirectPolicy()
+        qp_client_mock = mock.Mock()
+        client_mock = mock.Mock()
+        client_mock._base_url = "previous_url"
+        qp_client_mock._client = client_mock
+        qp_client_ref = weakref.ref(qp_client_mock)
+        policy._qp_client_ref = qp_client_ref
+        redirect_location = policy.get_redirect_location(pipeline_resp_mock)
+        self.assertEqual(redirect_location, "invalid_url")
+        self.assertEqual(client_mock._base_url,"previous_url")
+
+    def test_get_redirect_location_no_client(self):
+        policy = _QuickpulseRedirectPolicy()
+        pipeline_resp_mock = mock.Mock()
+        http_resp_mock = mock.Mock()
+        headers = {
+            "x-ms-qps-service-endpoint-redirect-v2": "https://eastus.livediagnostics.monitor.azure.com/QuickPulseService.svc"
+        }
+        http_resp_mock.headers = headers
+        pipeline_resp_mock.http_response = http_resp_mock
+        policy = _QuickpulseRedirectPolicy()
+        redirect_location = policy.get_redirect_location(pipeline_resp_mock)
+        self.assertEqual(redirect_location, "https://eastus.livediagnostics.monitor.azure.com/QuickPulseService.svc")
+        self.assertIsNone(policy._qp_client_ref)