Support AAD in soverign clouds (#39379)

Author: lzchen

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

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Support AAD for sovereign clouds
+  ([#39379](https://github.com/Azure/azure-sdk-for-python/pull/39379))
 - Support stable http semantic conventions for breeze exporter - REQUESTS
   ([#39208](https://github.com/Azure/azure-sdk-for-python/pull/39208))
 

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

@@ -7,6 +7,8 @@
 LIVE_ENDPOINT = "liveendpoint"
 INGESTION_ENDPOINT = "ingestionendpoint"
 INSTRUMENTATION_KEY = "instrumentationkey"
+# cspell:disable-next-line
+AAD_AUDIENCE = "aadaudience"
 
 # Validate UUID format
 # Specs taken from https://tools.ietf.org/html/rfc4122
@@ -26,6 +28,7 @@ def __init__(self, connection_string: typing.Optional[str] = None) -> None:
         self.endpoint = ""
         self.live_endpoint = ""
         self._connection_string = connection_string
+        self.aad_audience = ""
         self._initialize()
         self._validate_instrumentation_key()
 
@@ -42,17 +45,25 @@ def _initialize(self) -> None:
         # 3. Key from connection string in environment variable
         # 4. Key from instrumentation key in environment variable
         self.instrumentation_key = (
-            code_cs.get(INSTRUMENTATION_KEY) or code_ikey or env_cs.get(INSTRUMENTATION_KEY) or env_ikey  # type: ignore
+            code_cs.get(INSTRUMENTATION_KEY) or code_ikey or \
+                env_cs.get(INSTRUMENTATION_KEY) or env_ikey  # type: ignore
         )
         # The priority of the endpoints is as follows:
         # 1. The endpoint explicitly passed in connection string
         # 2. The endpoint from the connection string in environment variable
         # 3. The default breeze endpoint
         self.endpoint = (
-            code_cs.get(INGESTION_ENDPOINT) or env_cs.get(INGESTION_ENDPOINT) or "https://dc.services.visualstudio.com"
+            code_cs.get(INGESTION_ENDPOINT) or env_cs.get(INGESTION_ENDPOINT) or \
+                "https://dc.services.visualstudio.com"
         )
         self.live_endpoint = (
-            code_cs.get(LIVE_ENDPOINT) or env_cs.get(LIVE_ENDPOINT) or "https://rt.services.visualstudio.com"
+            code_cs.get(LIVE_ENDPOINT) or env_cs.get(LIVE_ENDPOINT) or \
+                "https://rt.services.visualstudio.com"
+        )
+        # The AUDIENCE is a url that identifies Azure Monitor in a specific cloud
+        # (For example: "https://monitor.azure.com/").
+        self.aad_audience = (
+            code_cs.get(AAD_AUDIENCE) or env_cs.get(AAD_AUDIENCE)  # type: ignore
         )
 
     def _validate_instrumentation_key(self) -> None:

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

@@ -204,6 +204,6 @@
 
 # AAD Auth
 
-_APPLICATION_INSIGHTS_RESOURCE_SCOPE = "https://monitor.azure.com//.default"
+_DEFAULT_AAD_SCOPE = "https://monitor.azure.com//.default"
 
 # cSpell:disable

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

@@ -22,7 +22,7 @@
 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,
+    _DEFAULT_AAD_SCOPE,
     _PYTHON_ENABLE_OPENTELEMETRY,
     _INSTRUMENTATIONS_BIT_MAP,
     _FUNCTIONS_WORKER_RUNTIME,
@@ -274,17 +274,25 @@ def _filter_custom_properties(properties: Attributes, filter=None) -> Dict[str,
     return truncated_properties
 
 
-def _get_auth_policy(credential, default_auth_policy):
+def _get_auth_policy(credential, default_auth_policy, aad_audience=None):
     if credential:
         if hasattr(credential, "get_token"):
             return BearerTokenCredentialPolicy(
                 credential,
-                _APPLICATION_INSIGHTS_RESOURCE_SCOPE,
+                _get_scope(aad_audience),
             )
         raise ValueError("Must pass in valid TokenCredential.")
     return default_auth_policy
 
 
+def _get_scope(aad_audience=None):
+    # The AUDIENCE is a url that identifies Azure Monitor in a specific cloud
+    # (For example: "https://monitor.azure.com/").
+    # The SCOPE is the audience + the permission
+    # (For example: "https://monitor.azure.com//.default").
+    return _DEFAULT_AAD_SCOPE if not aad_audience else "{}/.default".format(aad_audience)
+
+
 class Singleton(type):
     _instance = None
 

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

@@ -90,6 +90,7 @@ def __init__(self, **kwargs: Any) -> None:
         self._disable_offline_storage = kwargs.get("disable_offline_storage", False)
         self._endpoint = parsed_connection_string.endpoint
         self._instrumentation_key = parsed_connection_string.instrumentation_key
+        self._aad_audience = parsed_connection_string.aad_audience
         self._storage_maintenance_period = kwargs.get(
             "storage_maintenance_period", 60
         )  # Maintenance interval in seconds.
@@ -126,7 +127,7 @@ def __init__(self, **kwargs: Any) -> None:
             # Handle redirects in exporter, set new endpoint if redirected
             RedirectPolicy(permit_redirects=False),
             config.retry_policy,
-            _get_auth_policy(self._credential, config.authentication_policy),
+            _get_auth_policy(self._credential, config.authentication_policy, self._aad_audience),
             config.custom_hook_policy,
             config.logging_policy,
             # Explicitly disabling to avoid infinite loop of Span creation when data is exported

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/test_base_exporter.py

@@ -18,6 +18,7 @@
 )
 from azure.monitor.opentelemetry.exporter.statsbeat._state import _REQUESTS_MAP
 from azure.monitor.opentelemetry.exporter._constants import (
+    _DEFAULT_AAD_SCOPE,
     _REQ_DURATION_NAME,
     _REQ_EXCEPTION_NAME,
     _REQ_FAILURE_NAME,
@@ -941,16 +942,29 @@ def test_exporter_credential(self, mock_add_credential_policy):
         TEST_CREDENTIAL = "TEST_CREDENTIAL"
         base = BaseExporter(credential=TEST_CREDENTIAL, authentication_policy=TEST_AUTH_POLICY)
         self.assertEqual(base._credential, TEST_CREDENTIAL)
-        mock_add_credential_policy.assert_called_once_with(TEST_CREDENTIAL, TEST_AUTH_POLICY)
+        mock_add_credential_policy.assert_called_once_with(TEST_CREDENTIAL, TEST_AUTH_POLICY, None)
+
+    @mock.patch("azure.monitor.opentelemetry.exporter.export._base._get_auth_policy")
+    def test_exporter_credential_audience(self, mock_add_credential_policy):
+        test_cs = "AadAudience=test-aad"
+        TEST_CREDENTIAL = "TEST_CREDENTIAL"
+        base = BaseExporter(
+            connection_string=test_cs,
+            credential=TEST_CREDENTIAL,
+            authentication_policy=TEST_AUTH_POLICY,
+        )
+        self.assertEqual(base._credential, TEST_CREDENTIAL)
+        mock_add_credential_policy.assert_called_once_with(TEST_CREDENTIAL, TEST_AUTH_POLICY, "test-aad")
 
     def test_get_auth_policy(self):
         class TestCredential:
-            def get_token():
+            def get_token(self):
                 return "TEST_TOKEN"
 
         credential = TestCredential()
         result = _get_auth_policy(credential, TEST_AUTH_POLICY)
         self.assertEqual(result._credential, credential)
+        self.assertEqual(result._scopes, (_DEFAULT_AAD_SCOPE,))
 
     def test_get_auth_policy_no_credential(self):
         self.assertEqual(_get_auth_policy(credential=None, default_auth_policy=TEST_AUTH_POLICY), TEST_AUTH_POLICY)
@@ -964,6 +978,16 @@ def invalid_get_token():
             ValueError, _get_auth_policy, credential=InvalidTestCredential(), default_auth_policy=TEST_AUTH_POLICY
         )
 
+    def test_get_auth_policy_audience(self):
+        class TestCredential:
+            def get_token():
+                return "TEST_TOKEN"
+
+        credential = TestCredential()
+        result = _get_auth_policy(credential, TEST_AUTH_POLICY, aad_audience="test_audience")
+        self.assertEqual(result._credential, credential)
+        self.assertEqual(result._scopes, ("test_audience/.default",))
+
 
 def validate_telemetry_item(item1, item2):
     return (