Support AAD auth for live metrics (#37258)

Author: lzchen

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

@@ -6,6 +6,8 @@
 
 - Allow passing in of custom `TracerProvider` for `AzureMonitorTraceExporter`
     ([#36363](https://github.com/Azure/azure-sdk-for-python/pull/36363))
+- Support AAD Auth for live metrics
+    ([#37258](https://github.com/Azure/azure-sdk-for-python/pull/37258))
 
 ### Breaking Changes
 

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

@@ -200,4 +200,8 @@
 
 _SAMPLE_RATE_KEY = "_MS.sampleRate"
 
+# AAD Auth
+
+_APPLICATION_INSIGHTS_RESOURCE_SCOPE = "https://monitor.azure.com//.default"
+
 # cSpell:disable

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

@@ -46,7 +46,11 @@
     _metric_to_quick_pulse_data_points,
 )
 from azure.monitor.opentelemetry.exporter._connection_string_parser import ConnectionStringParser
-from azure.monitor.opentelemetry.exporter._utils import _ticks_since_dot_net_epoch, PeriodicTask
+from azure.monitor.opentelemetry.exporter._utils import (
+    _get_auth_policy,
+    _ticks_since_dot_net_epoch,
+    PeriodicTask,
+)
 
 
 _logger = logging.getLogger(__name__)
@@ -75,19 +79,20 @@ class _UnsuccessfulQuickPulsePostError(Exception):
 
 class _QuickpulseExporter(MetricExporter):
 
-    def __init__(self, connection_string: Optional[str]) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         """Metric exporter for Quickpulse.
 
         :param str connection_string: The connection string used for your Application Insights resource.
+        :keyword TokenCredential credential: Token credential, such as ManagedIdentityCredential or
+            ClientSecretCredential, used for Azure Active Directory (AAD) authentication. Defaults to None.
         :rtype: None
         """
-        parsed_connection_string = ConnectionStringParser(connection_string)
+        parsed_connection_string = ConnectionStringParser(kwargs.get('connection_string'))
 
         self._live_endpoint = parsed_connection_string.live_endpoint
         self._instrumentation_key = parsed_connection_string.instrumentation_key
-        # TODO: Support AADaudience (scope)/credentials
-        # Pass `None` for now until swagger definition is fixed
-        config = QuickpulseClientConfiguration(credential=None)  # type: ignore
+        self._credential = kwargs.get('credential')
+        config = QuickpulseClientConfiguration(credential=self._credential)  # type: ignore
         qp_redirect_policy = _QuickpulseRedirectPolicy(permit_redirects=False)
         policies = [
             # Custom redirect policy for QP
@@ -96,13 +101,13 @@ def __init__(self, connection_string: Optional[str]) -> None:
             ContentDecodePolicy(),
             # Logging for client calls
             config.http_logging_policy,
-            # TODO: Support AADaudience (scope)/credentials
+            _get_auth_policy(self._credential, config.authentication_policy),
             config.authentication_policy,
             # Explicitly disabling to avoid tracing live metrics calls
             # DistributedTracingPolicy(),
         ]
         self._client = QuickpulseClient(
-            credential=None, # type: ignore
+            credential=self._credential, # type: ignore
             endpoint=self._live_endpoint,
             policies=policies
         )

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

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 # cSpell:disable
 from datetime import datetime
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable
 
 import platform
 import psutil
@@ -70,21 +70,25 @@ def enable_live_metrics(**kwargs: Any) -> None:  # pylint: disable=C4758
 
     :keyword str connection_string: The connection string used for your Application Insights resource.
     :keyword Resource resource: The OpenTelemetry Resource used for this Python application.
+    :keyword TokenCredential credential: Token credential, such as ManagedIdentityCredential or
+        ClientSecretCredential, used for Azure Active Directory (AAD) authentication. Defaults to None.
     :rtype: None
     """
-    _QuickpulseManager(kwargs.get('connection_string'), kwargs.get('resource'))
+    _QuickpulseManager(**kwargs)
     set_statsbeat_live_metrics_feature_set()
 
 
 # pylint: disable=protected-access,too-many-instance-attributes
 class _QuickpulseManager(metaclass=Singleton):
 
-    def __init__(self, connection_string: Optional[str], resource: Optional[Resource]) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         _set_global_quickpulse_state(_QuickpulseState.PING_SHORT)
-        self._exporter = _QuickpulseExporter(connection_string)
+        self._exporter = _QuickpulseExporter(**kwargs)
         part_a_fields = {}
-        if resource:
-            part_a_fields = _populate_part_a_fields(resource)
+        resource = kwargs.get('resource')
+        if not resource:
+            resource = Resource.create({})
+        part_a_fields = _populate_part_a_fields(resource)
         id_generator = RandomIdGenerator()
         self._base_monitoring_data_point = MonitoringDataPoint(
             version=_get_sdk_version(),
@@ -97,7 +101,10 @@ def __init__(self, connection_string: Optional[str], resource: Optional[Resource
             performance_collection_supported=True,
         )
         self._reader = _QuickpulseMetricReader(self._exporter, self._base_monitoring_data_point)
-        self._meter_provider = MeterProvider([self._reader])
+        self._meter_provider = MeterProvider(
+            metric_readers=[self._reader],
+            resource=resource,
+        )
         self._meter = self._meter_provider.get_meter("azure_monitor_live_metrics")
 
         self._request_duration = self._meter.create_histogram(

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_utils.py

@@ -17,13 +17,15 @@
 from opentelemetry.sdk.util import ns_to_iso_str
 from opentelemetry.util.types import Attributes
 
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.monitor.opentelemetry.exporter._generated.models import ContextTagKeys, TelemetryItem
 from azure.monitor.opentelemetry.exporter._version import VERSION as ext_version
 from azure.monitor.opentelemetry.exporter._constants import (
+    _AKS_ARM_NAMESPACE_ID,
+    _APPLICATION_INSIGHTS_RESOURCE_SCOPE,
     _INSTRUMENTATIONS_BIT_MAP,
-    _WEBSITE_SITE_NAME,
     _FUNCTIONS_WORKER_RUNTIME,
-    _AKS_ARM_NAMESPACE_ID,
+    _WEBSITE_SITE_NAME,
 )
 
 
@@ -260,6 +262,19 @@ def _filter_custom_properties(properties: Attributes, filter=None) -> Dict[str,
     return truncated_properties
 
 
+def _get_auth_policy(credential, default_auth_policy):
+    if credential:
+        if hasattr(credential, 'get_token'):
+            return BearerTokenCredentialPolicy(
+                credential,
+                _APPLICATION_INSIGHTS_RESOURCE_SCOPE,
+            )
+        raise ValueError(
+            'Must pass in valid TokenCredential.'
+        )
+    return default_auth_policy
+
+
 class Singleton(type):
     _instance = None
     def __call__(cls, *args, **kwargs):

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/_base.py

@@ -10,7 +10,6 @@
 
 from azure.core.exceptions import HttpResponseError, ServiceRequestError
 from azure.core.pipeline.policies import (
-    BearerTokenCredentialPolicy,
     ContentDecodePolicy,
     HttpLoggingPolicy,
     RedirectPolicy,
@@ -44,6 +43,7 @@
 )
 from azure.monitor.opentelemetry.exporter._connection_string_parser import ConnectionStringParser
 from azure.monitor.opentelemetry.exporter._storage import LocalFileStorage
+from azure.monitor.opentelemetry.exporter._utils import _get_auth_policy
 from azure.monitor.opentelemetry.exporter.statsbeat._state import (
     get_statsbeat_initial_success,
     get_statsbeat_shutdown,
@@ -58,7 +58,6 @@
 _AZURE_TEMPDIR_PREFIX = "Microsoft/AzureMonitor"
 _TEMPDIR_PREFIX = "opentelemetry-python-"
 _SERVICE_API_LATEST = "2020-09-15_Preview"
-_APPLICATION_INSIGHTS_RESOURCE_SCOPE = "https://monitor.azure.com//.default"
 
 class ExportResult(Enum):
     SUCCESS = 0
@@ -346,19 +345,6 @@ def _is_stats_exporter(self):
         return self.__class__.__name__ == "_StatsBeatExporter"
 
 
-def _get_auth_policy(credential, default_auth_policy):
-    if credential:
-        if hasattr(credential, 'get_token'):
-            return BearerTokenCredentialPolicy(
-                credential,
-                _APPLICATION_INSIGHTS_RESOURCE_SCOPE,
-            )
-        raise ValueError(
-            'Must pass in valid TokenCredential.'
-        )
-    return default_auth_policy
-
-
 def _is_invalid_code(response_code: Optional[int]) -> bool:
     """Determine if response is a invalid response.
 

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

@@ -91,7 +91,7 @@ def setUpClass(cls):
             performance_collection_supported=True,
         )
         cls._exporter = _QuickpulseExporter(
-            "InstrumentationKey=4321abcd-5678-4efa-8abc-1234567890ac;LiveEndpoint=https://eastus.livediagnostics.monitor.azure.com/"
+            connection_string="InstrumentationKey=4321abcd-5678-4efa-8abc-1234567890ac;LiveEndpoint=https://eastus.livediagnostics.monitor.azure.com/"
         )
         cls._reader = _QuickpulseMetricReader(
             cls._exporter,

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

@@ -64,7 +64,10 @@ def test_enable_live_metrics(self, manager_mock):
             connection_string="test_cs",
             resource=mock_resource,
         )
-        manager_mock.assert_called_with("test_cs", mock_resource)
+        manager_mock.assert_called_with(
+            connection_string="test_cs",
+            resource=mock_resource
+        )
 
 
 class TestQuickpulseManager(unittest.TestCase):
@@ -118,6 +121,7 @@ def test_init(self, generator_mock):
         self.assertEqual(qpm._reader._base_monitoring_data_point, qpm._base_monitoring_data_point)
         self.assertTrue(isinstance(qpm._meter_provider, MeterProvider))
         self.assertEqual(qpm._meter_provider._sdk_config.metric_readers, [qpm._reader])
+        self.assertEqual(qpm._meter_provider._sdk_config.resource, resource)
         self.assertTrue(isinstance(qpm._meter, Meter))
         self.assertEqual(qpm._meter.name, "azure_monitor_live_metrics")
         self.assertTrue(isinstance(qpm._request_duration, Histogram))