[Monitor Query] MetricsClient sovereign cloud support (#35502)

This adds some documentation and samples regarding using MetricsClient
with sovereign clouds. This adds a feature that extrapolates the correct
audience based on the passed in endpoint. This saves the user from
having to pass in extra information via the "audience" kwarg.


Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>

sdk/monitor/azure-monitor-query/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- An `audience` keyword argument can now be passed to the `MetricsClient` constructor to specify the audience for the authentication token. This is useful when querying metrics in sovereign clouds. ([#35502](https://github.com/Azure/azure-sdk-for-python/pull/35502))
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/monitor/azure-monitor-query/README.md

@@ -69,17 +69,20 @@ async_metrics_client = MetricsClient("https://<regional endpoint>", credential)
 
 #### Configure client for Azure sovereign cloud
 
-By default, `LogsQueryClient` and `MetricsQueryClient` are configured to use the Azure Public Cloud. To use a sovereign cloud instead, provide the correct `endpoint` argument. For example:
+By default, all clients are configured to use the Azure public cloud. To use a sovereign cloud, provide the correct `endpoint` argument when using `LogsQueryClient` or `MetricsQueryClient`. For `MetricsClient`, provide the correct `audience` argument instead. For example:
 
 ```python
 from azure.identity import AzureAuthorityHosts, DefaultAzureCredential
-from azure.monitor.query import LogsQueryClient, MetricsQueryClient
+from azure.monitor.query import LogsQueryClient, MetricsQueryClient, MetricsClient
 
 # Authority can also be set via the AZURE_AUTHORITY_HOST environment variable.
 credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_GOVERNMENT)
 
 logs_query_client = LogsQueryClient(credential, endpoint="https://api.loganalytics.us/v1")
 metrics_query_client = MetricsQueryClient(credential, endpoint="https://management.usgovcloudapi.net")
+metrics_client = MetricsClient(
+    "https://usgovvirginia.metrics.monitor.azure.us", credential, audience="https://metrics.monitor.azure.us"
+)
 ```
 
 **Note**: Currently, `MetricsQueryClient` uses the Azure Resource Manager (ARM) endpoint for querying metrics. You need the corresponding management endpoint for your cloud when using this client. This detail is subject to change in the future.

sdk/monitor/azure-monitor-query/azure/monitor/query/_metrics_client.py

@@ -28,6 +28,26 @@ class MetricsClient:  # pylint: disable=client-accepts-api-version-keyword
         resources. For global resources, the region should be 'global'. Required.
     :param credential: The credential to authenticate the client.
     :type credential: ~azure.core.credentials.TokenCredential
+    :keyword str audience: The audience to use when requesting tokens for Microsoft Entra ID. Defaults to the public
+        cloud audience (https://metrics.monitor.azure.com).
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/sample_authentication.py
+            :start-after: [START create_metrics_client]
+            :end-before: [END create_metrics_client]
+            :language: python
+            :dedent: 4
+            :caption: Creating the MetricsClient with a TokenCredential.
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/sample_authentication.py
+            :start-after: [START create_metrics_client_sovereign_cloud]
+            :end-before: [END create_metrics_client_sovereign_cloud]
+            :language: python
+            :dedent: 4
+            :caption: Creating the MetricsClient for use with a sovereign cloud (i.e. non-public cloud).
     """
 
     def __init__(self, endpoint: str, credential: TokenCredential, **kwargs: Any) -> None:
@@ -59,7 +79,7 @@ def query_resources(
         order_by: Optional[str] = None,
         filter: Optional[str] = None,
         roll_up_by: Optional[str] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> List[MetricsQueryResult]:
         """Lists the metric values for multiple resources.
 
@@ -154,7 +174,7 @@ def query_resources(
             orderby=order_by,
             filter=filter,
             rollupby=roll_up_by,  # cspell:ignore rollupby
-            **kwargs
+            **kwargs,
         )
 
         # In rare cases, the generated value is a JSON string instead of a dict. This potentially stems from a bug in

sdk/monitor/azure-monitor-query/azure/monitor/query/aio/_metrics_client_async.py

@@ -29,6 +29,26 @@ class MetricsClient:  # pylint: disable=client-accepts-api-version-keyword
         resources. For global resources, the region should be 'global'. Required.
     :param credential: The credential to authenticate the client.
     :type credential: ~azure.core.credentials_async.AsyncTokenCredential
+    :keyword str audience: The audience to use when requesting tokens for Microsoft Entra ID. Defaults to the public
+        cloud audience (https://metrics.monitor.azure.com).
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/async_samples/sample_authentication_async.py
+            :start-after: [START create_metrics_client_async]
+            :end-before: [END create_metrics_client_async]
+            :language: python
+            :dedent: 4
+            :caption: Creating the asynchronous MetricsClient with a TokenCredential.
+
+    .. admonition:: Example:
+
+        .. literalinclude:: ../samples/async_samples/sample_authentication_async.py
+            :start-after: [START create_metrics_client_sovereign_cloud_async]
+            :end-before: [END create_metrics_client_sovereign_cloud_async]
+            :language: python
+            :dedent: 4
+            :caption: Creating the MetricsClient for use with a sovereign cloud (i.e. non-public cloud).
     """
 
     def __init__(self, endpoint: str, credential: AsyncTokenCredential, **kwargs: Any) -> None:
@@ -60,7 +80,7 @@ async def query_resources(
         order_by: Optional[str] = None,
         filter: Optional[str] = None,
         roll_up_by: Optional[str] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> List[MetricsQueryResult]:
         """Lists the metric values for multiple resources.
 
@@ -155,7 +175,7 @@ async def query_resources(
             orderby=order_by,
             filter=filter,
             rollupby=roll_up_by,  # cspell:ignore rollupby
-            **kwargs
+            **kwargs,
         )
 
         # In rare cases, the generated value is a JSON string instead of a dict. This potentially stems from a bug in

sdk/monitor/azure-monitor-query/samples/async_samples/sample_authentication_async.py

@@ -56,11 +56,38 @@ async def create_metrics_query_client_sovereign_cloud_async():
     # [END create_metrics_query_client_sovereign_cloud_async]
 
 
+async def create_metrics_client_async():
+    # [START create_metrics_client_async]
+    from azure.identity.aio import DefaultAzureCredential
+    from azure.monitor.query.aio import MetricsClient
+
+    credential = DefaultAzureCredential()
+    client = MetricsClient("https://eastus.metrics.monitor.azure.com", credential)
+    # [END create_metrics_client_async]
+
+
+async def create_metrics_client_sovereign_cloud_async():
+    # [START create_metrics_client_sovereign_cloud_async]
+    from azure.identity import AzureAuthorityHosts
+    from azure.identity.aio import DefaultAzureCredential
+    from azure.monitor.query.aio import MetricsClient
+
+    credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_GOVERNMENT)
+    client = MetricsClient(
+        "https://usgovvirginia.metrics.monitor.azure.us",
+        credential,
+        audience="https://metrics.monitor.azure.us",
+    )
+    # [END create_metrics_client_sovereign_cloud_async]
+
+
 async def main():
     await create_logs_query_client_async()
     await create_logs_query_client_sovereign_cloud_async()
     await create_metrics_query_client_async()
     await create_metrics_query_client_sovereign_cloud_async()
+    await create_metrics_client_async()
+    await create_metrics_client_sovereign_cloud_async()
 
 
 if __name__ == '__main__':

sdk/monitor/azure-monitor-query/samples/sample_authentication.py

@@ -53,8 +53,34 @@ def create_metrics_query_client_sovereign_cloud():
     # [END create_metrics_query_client_sovereign_cloud]
 
 
+def create_metrics_client():
+    # [START create_metrics_client]
+    from azure.identity import DefaultAzureCredential
+    from azure.monitor.query import MetricsClient
+
+    credential = DefaultAzureCredential()
+    client = MetricsClient("https://eastus.metrics.monitor.azure.com", credential)
+    # [END create_metrics_client]
+
+
+def create_metrics_client_sovereign_cloud():
+    # [START create_metrics_client_sovereign_cloud]
+    from azure.identity import AzureAuthorityHosts, DefaultAzureCredential
+    from azure.monitor.query import MetricsClient
+
+    credential = DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_GOVERNMENT)
+    client = MetricsClient(
+        "https://usgovvirginia.metrics.monitor.azure.us",
+        credential,
+        audience="https://metrics.monitor.azure.us",
+    )
+    # [END create_metrics_client_sovereign_cloud]
+
+
 if __name__ == '__main__':
     create_logs_query_client()
     create_logs_query_client_sovereign_cloud()
     create_metrics_query_client()
     create_metrics_query_client_sovereign_cloud()
+    create_metrics_client()
+    create_metrics_client_sovereign_cloud()

sdk/monitor/azure-monitor-query/tests/base_testcase.py

@@ -21,7 +21,7 @@
 METRICS_CLIENT_ENVIRONMENT_AUDIENCE_MAP = {
     "AzureCloud": "https://metrics.monitor.azure.com",
     "AzureChinaCloud": "https://metrics.monitor.azure.cn",
-    "AzureUSGovernment": "https:/metrics.monitor.azure.us"
+    "AzureUSGovernment": "https://metrics.monitor.azure.us"
 }
 
 TLD_MAP = {

sdk/monitor/azure-monitor-query/tests/test_metrics_client.py

@@ -44,3 +44,12 @@ def test_batch_metrics_granularity(self, recorded_test, monitor_info):
             assert metric.timeseries
             for t in metric.timeseries:
                 assert t.metadata_values is not None
+
+    def test_client_different_endpoint(self):
+        credential = self.get_credential(MetricsClient)
+        endpoint = "https://usgovvirginia.metrics.monitor.azure.us"
+        audience = "https://metrics.monitor.azure.us"
+        client = MetricsClient(endpoint, credential, audience=audience)
+
+        assert client._endpoint == endpoint
+        assert f"{audience}/.default" in client._client._config.authentication_policy._scopes

sdk/monitor/azure-monitor-query/tests/test_metrics_client_async.py

@@ -52,3 +52,13 @@ async def test_batch_metrics_granularity(self, recorded_test, monitor_info):
                 assert metric.timeseries
                 for t in metric.timeseries:
                     assert t.metadata_values is not None
+
+    @pytest.mark.asyncio
+    async def test_client_different_endpoint(self):
+        credential = self.get_credential(MetricsClient, is_async=True)
+        endpoint = "https://usgovvirginia.metrics.monitor.azure.us"
+        audience = "https://metrics.monitor.azure.us"
+        client = MetricsClient(endpoint, credential, audience=audience)
+
+        assert client._endpoint == endpoint
+        assert f"{audience}/.default" in client._client._config.authentication_policy._scopes