[Monitor Query] Typing updates (#35180)

Regenerated with latest autorest and updated some samples. Changed the model types for status to use Literal with the corresponding enum values to improve typing. Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>
Azure · Apr 19, 2024 · d5c3578 · d5c3578
1 parent dccef72
commit d5c3578
Show file tree

Hide file tree

Showing 29 changed files with 184 additions and 142 deletions.
diff --git a/sdk/monitor/azure-monitor-query/README.md b/sdk/monitor/azure-monitor-query/README.md
@@ -147,12 +147,14 @@ query = """AzureActivity | take 5"""
 
 try:
     response = client.query_resource(os.environ['LOGS_RESOURCE_ID'], query, timespan=timedelta(days=1))
-    if response.status == LogsQueryStatus.PARTIAL:
+    if response.status == LogsQueryStatus.SUCCESS:
+        data = response.tables
+    else:
+        # LogsQueryPartialResult
         error = response.partial_error
         data = response.partial_data
         print(error)
-    elif response.status == LogsQueryStatus.SUCCESS:
-        data = response.tables
+
     for table in data:
         df = pd.DataFrame(data=table.rows, columns=table.columns)
         print(df)
@@ -175,7 +177,7 @@ For example:
 import os
 import pandas as pd
 from datetime import datetime, timezone
-from azure.monitor.query import LogsQueryClient, LogsQueryStatus
+from azure.monitor.query import LogsQueryClient, LogsQueryResult
 from azure.identity import DefaultAzureCredential
 from azure.core.exceptions import HttpResponseError
 
@@ -193,12 +195,14 @@ try:
         query=query,
         timespan=(start_time, end_time)
         )
-    if response.status == LogsQueryStatus.PARTIAL:
+    if response.status == LogsQueryStatus.SUCCESS:
+        data = response.tables
+    else:
+        # LogsQueryPartialResult
         error = response.partial_error
         data = response.partial_data
         print(error)
-    elif response.status == LogsQueryStatus.SUCCESS:
-        data = response.tables
+
     for table in data:
         df = pd.DataFrame(data=table.rows, columns=table.columns)
         print(df)
@@ -294,10 +298,7 @@ requests = [
 results = client.query_batch(requests)
 
 for res in results:
-    if res.status == LogsQueryStatus.FAILURE:
-        # this will be a LogsQueryError
-        print(res.message)
-    elif res.status == LogsQueryStatus.PARTIAL:
+    if res.status == LogsQueryStatus.PARTIAL:
         ## this will be a LogsQueryPartialResult
         print(res.partial_error)
         for table in res.partial_data:
@@ -308,6 +309,9 @@ for res in results:
         table = res.tables[0]
         df = pd.DataFrame(table.rows, columns=table.columns)
         print(df)
+    else:
+        # this will be a LogsQueryError
+        print(res.message)
 
 ```
 

diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py
@@ -5,7 +5,7 @@
 # license information.
 # --------------------------------------------------------------------------
 import sys
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Literal
 
 from ._enums import LogsQueryStatus
 
@@ -27,7 +27,7 @@ class LogsQueryError:
     """A human readable error message."""
     details: Optional[List[JSON]] = None
     """A list of additional details about the error."""
-    status: LogsQueryStatus
+    status: Literal[LogsQueryStatus.FAILURE]
     """Status for error item when iterating over list of results. Always "Failure" for an instance of a
     LogsQueryError."""
 

diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_generated/aio/operations/_operations.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_generated/aio/operations/_operations.py
@@ -9,7 +9,7 @@
 import datetime
 from io import IOBase
 import sys
-from typing import Any, Callable, Dict, IO, Optional, TypeVar, Union, cast, overload
+from typing import Any, Callable, Dict, IO, Optional, Type, TypeVar, Union, cast, overload
 
 from azure.core.exceptions import (
     ClientAuthenticationError,
@@ -136,7 +136,7 @@ async def get(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -450,7 +450,7 @@ async def execute(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -581,7 +581,7 @@ async def resource_get(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -892,7 +892,7 @@ async def resource_execute(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1286,7 +1286,7 @@ async def batch(self, body: Union[JSON, IO[bytes]], **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1415,7 +1415,7 @@ async def resource_get_xms(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1726,7 +1726,7 @@ async def resource_execute_xms(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -2159,7 +2159,7 @@ async def get(self, workspace_id: str, **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -2561,7 +2561,7 @@ async def post(self, workspace_id: str, **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/.../azure-monitor-query/azure/monitor/query/_generated/metrics/aio/operations/_operations.py b/.../azure-monitor-query/azure/monitor/query/_generated/metrics/aio/operations/_operations.py
@@ -8,7 +8,7 @@
 # --------------------------------------------------------------------------
 import datetime
 import sys
-from typing import Any, AsyncIterable, Callable, Dict, Optional, TypeVar, cast
+from typing import Any, AsyncIterable, Callable, Dict, Optional, Type, TypeVar, cast
 import urllib.parse
 
 from azure.core.async_paging import AsyncItemPaged, AsyncList
@@ -126,7 +126,7 @@ def list(self, resource_uri: str, *, metricnamespace: Optional[str] = None, **kw
 
         cls: ClsType[JSON] = kwargs.pop("cls", None)
 
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -367,7 +367,7 @@ async def list(
                       for metrics.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -474,7 +474,7 @@ def list(self, resource_uri: str, *, start_time: Optional[str] = None, **kwargs:
 
         cls: ClsType[JSON] = kwargs.pop("cls", None)
 
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/...-monitor-query/azure/monitor/query/_generated/metrics/batch/aio/operations/_operations.py b/...-monitor-query/azure/monitor/query/_generated/metrics/batch/aio/operations/_operations.py
@@ -9,7 +9,7 @@
 import datetime
 from io import IOBase
 import sys
-from typing import Any, Callable, Dict, IO, List, Optional, TypeVar, Union, cast, overload
+from typing import Any, Callable, Dict, IO, List, Optional, Type, TypeVar, Union, cast, overload
 
 from azure.core.exceptions import (
     ClientAuthenticationError,
@@ -566,7 +566,7 @@ async def batch(
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/...zure-monitor-query/azure/monitor/query/_generated/metrics/batch/operations/_operations.py b/...zure-monitor-query/azure/monitor/query/_generated/metrics/batch/operations/_operations.py
@@ -9,7 +9,7 @@
 import datetime
 from io import IOBase
 import sys
-from typing import Any, Callable, Dict, IO, List, Optional, TypeVar, Union, cast, overload
+from typing import Any, Callable, Dict, IO, List, Optional, Type, TypeVar, Union, cast, overload
 
 from azure.core.exceptions import (
     ClientAuthenticationError,
@@ -628,7 +628,7 @@ def batch(
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/...itor/azure-monitor-query/azure/monitor/query/_generated/metrics/operations/_operations.py b/...itor/azure-monitor-query/azure/monitor/query/_generated/metrics/operations/_operations.py
@@ -8,7 +8,7 @@
 # --------------------------------------------------------------------------
 import datetime
 import sys
-from typing import Any, Callable, Dict, Iterable, Optional, TypeVar, cast
+from typing import Any, Callable, Dict, Iterable, Optional, Type, TypeVar, cast
 import urllib.parse
 
 from azure.core.exceptions import (
@@ -244,7 +244,7 @@ def list(self, resource_uri: str, *, metricnamespace: Optional[str] = None, **kw
 
         cls: ClsType[JSON] = kwargs.pop("cls", None)
 
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -485,7 +485,7 @@ def list(
                       for metrics.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -592,7 +592,7 @@ def list(self, resource_uri: str, *, start_time: Optional[str] = None, **kwargs:
 
         cls: ClsType[JSON] = kwargs.pop("cls", None)
 
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_generated/operations/_operations.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_generated/operations/_operations.py
@@ -9,7 +9,7 @@
 import datetime
 from io import IOBase
 import sys
-from typing import Any, Callable, Dict, IO, Optional, TypeVar, Union, cast, overload
+from typing import Any, Callable, Dict, IO, Optional, Type, TypeVar, Union, cast, overload
 
 from azure.core.exceptions import (
     ClientAuthenticationError,
@@ -341,7 +341,7 @@ def get(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -655,7 +655,7 @@ def execute(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -786,7 +786,7 @@ def resource_get(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1097,7 +1097,7 @@ def resource_execute(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1491,7 +1491,7 @@ def batch(self, body: Union[JSON, IO[bytes]], **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1620,7 +1620,7 @@ def resource_get_xms(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -1931,7 +1931,7 @@ def resource_execute_xms(
                     "statistics": {}  # Optional. Statistics represented in JSON format.
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -2364,7 +2364,7 @@ def get(self, workspace_id: str, **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,
@@ -2766,7 +2766,7 @@ def post(self, workspace_id: str, **kwargs: Any) -> JSON:
                     ]
                 }
         """
-        error_map = {
+        error_map: MutableMapping[int, Type[HttpResponseError]] = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
             409: ResourceExistsError,

diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_models.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_models.py
@@ -8,7 +8,7 @@
 import uuid
 from datetime import datetime, timedelta
 import sys
-from typing import Any, Optional, List, Union, Tuple, Dict, Iterator
+from typing import Any, Optional, List, Union, Tuple, Dict, Iterator, Literal
 
 from ._enums import LogsQueryStatus, MetricAggregationType, MetricClass, MetricNamespaceClassification, MetricUnit
 from ._exceptions import LogsQueryError
@@ -376,7 +376,7 @@ class LogsQueryResult:
     visualization: Optional[JSON] = None
     """This will include a visualization property in the response that specifies the type of visualization selected
     by the query and any properties for that visualization."""
-    status: LogsQueryStatus
+    status: Literal[LogsQueryStatus.SUCCESS]
     """The status of the result. Always 'Success' for an instance of a LogsQueryResult."""
 
     def __init__(self, **kwargs: Any) -> None:
@@ -554,7 +554,7 @@ class LogsQueryPartialResult:
     selected by the query and any properties for that visualization."""
     partial_error: Optional[LogsQueryError] = None
     """The partial error info."""
-    status: LogsQueryStatus
+    status: Literal[LogsQueryStatus.PARTIAL]
     """The status of the result. Always 'PartialError' for an instance of a LogsQueryPartialResult."""
 
     def __init__(self, **kwargs: Any) -> None: