Handle Errors and Tables in Query (#20658)

* initial commit

* batch set

* batch

* tests + changes

* more tests

* lint

* changelog

* Apply suggestions from code review

* comment

* lint

* querry

* comments

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

@@ -4,8 +4,14 @@
 
 ### Features Added
 
+- Added `QueryPartialErrorException` and `LogsQueryError` to handle errors.
+- Added `partial_error` and `is_error` attributes to `LogsQueryResult`.
+- Added an option `allow_partial_errors` that defaults to False, which can be set to not throw if there are any partial errors.
+
 ### Breaking Changes
 
+- `LogsQueryResult` now iterates over the tables directly as a convinience.
+
 ### Bugs Fixed
 
 ### Other Changes

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

@@ -7,6 +7,11 @@
 from ._logs_query_client import LogsQueryClient
 from ._metrics_query_client import MetricsQueryClient
 
+from ._exceptions import (
+    LogsQueryError,
+    QueryPartialErrorException
+)
+
 from ._models import (
     MetricAggregationType,
     LogsQueryResult,
@@ -30,6 +35,8 @@
     "MetricAggregationType",
     "LogsQueryClient",
     "LogsQueryResult",
+    "LogsQueryError",
+    "QueryPartialErrorException",
     "LogsTable",
     "LogsBatchQuery",
     "MetricsQueryClient",

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

@@ -0,0 +1,79 @@
+#
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+from azure.core.exceptions import HttpResponseError
+
+class LogsQueryError(object):
+    """The code and message for an error.
+
+    All required parameters must be populated in order to send to Azure.
+
+    :ivar code: A machine readable error code.
+    :vartype code: str
+    :ivar message: A human readable error message.
+    :vartype message: str
+    :ivar details: error details.
+    :vartype details: list[~monitor_query_client.models.ErrorDetail]
+    :ivar innererror: Inner error details if they exist.
+    :vartype innererror: ~azure.monitor.query.LogsQueryError
+    :ivar additional_properties: Additional properties that can be provided on the error info
+     object.
+    :vartype additional_properties: object
+    :ivar bool is_error: Boolean check for error item when iterating over list of
+        results. Always True for an instance of a LogsQueryError.
+    """
+    def __init__(
+        self,
+        **kwargs
+    ):
+        self.code = kwargs.get('code', None)
+        self.message = kwargs.get('message', None)
+        self.details = kwargs.get('details', None)
+        self.innererror = kwargs.get('innererror', None)
+        self.additional_properties = kwargs.get('additional_properties', None)
+        self.is_error = True
+
+    @classmethod
+    def _from_generated(cls, generated):
+        if not generated:
+            return None
+        details = None
+        if generated.details is not None:
+            details = [d.serialize() for d in generated.details]
+        return cls(
+            code=generated.code,
+            message=generated.message,
+            innererror=cls._from_generated(generated.innererror) if generated.innererror else None,
+            additional_properties=generated.additional_properties,
+            details=details,
+        )
+
+class QueryPartialErrorException(HttpResponseError):
+    """There is a partial failure in query operation. This is thrown for a single query operation
+     when allow_partial_errors is set to False.
+
+    :ivar code: A machine readable error code.
+    :vartype code: str
+    :ivar message: A human readable error message.
+    :vartype message: str
+    :ivar details: error details.
+    :vartype details: list[~monitor_query_client.models.ErrorDetail]
+    :ivar innererror: Inner error details if they exist.
+    :vartype innererror: ~azure.monitor.query.LogsQueryError
+    :ivar additional_properties: Additional properties that can be provided on the error info
+     object.
+    :vartype additional_properties: object
+    """
+
+    def __init__(self, **kwargs):
+        error = kwargs.pop('error', None)
+        if error:
+            self.code = error.code
+            self.message = error.message
+            self.details = [d.serialize() for d in error.details] if error.details else None
+            self.innererror = LogsQueryError._from_generated(error.innererror) if error.innererror else None
+            self.additional_properties = error.additional_properties
+        super(QueryPartialErrorException, self).__init__(message=self.message)

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

@@ -41,13 +41,21 @@ def get_metrics_authentication_policy(
 
     raise TypeError("Unsupported credential")
 
-def process_error(exception):
-    raise_error = HttpResponseError
-    raise raise_error(message=exception.message, response=exception.response)
-
-def order_results(request_order, mapping, obj):
+def order_results(request_order, mapping, obj, err, allow_partial_errors=False):
     ordered = [mapping[id] for id in request_order]
-    return [obj._from_generated(rsp) for rsp in ordered] # pylint: disable=protected-access
+    results = []
+    for item in ordered:
+        if not item.body.error:
+            results.append(obj._from_generated(item.body)) # pylint: disable=protected-access
+        else:
+            error = item.body.error
+            if allow_partial_errors and error.code == 'PartialError':
+                res = obj._from_generated(item.body) # pylint: disable=protected-access
+                res.partial_error = err._from_generated(error) # pylint: disable=protected-access
+                results.append(res)
+            else:
+                results.append(err._from_generated(error)) # pylint: disable=protected-access
+    return results
 
 def construct_iso8601(timespan=None):
     if not timespan:
@@ -90,3 +98,23 @@ def native_col_type(col_type, value):
 
 def process_row(col_types, row):
     return [native_col_type(col_types[ind], val) for ind, val in enumerate(row)]
+
+def process_error(error, model):
+    try:
+        model = model._from_generated(error.model.error) # pylint: disable=protected-access
+    except AttributeError: # model can be none
+        pass
+    raise HttpResponseError(
+        message=error.message,
+        response=error.response,
+        model=model)
+
+def process_prefer(server_timeout, include_statistics, include_visualization):
+    prefer = ""
+    if server_timeout:
+        prefer += "wait=" + str(server_timeout) + ","
+    if include_statistics:
+        prefer += "include-statistics=true,"
+    if include_visualization:
+        prefer += "include-render=true"
+    return prefer.rstrip(",")

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

@@ -12,8 +12,9 @@
 from ._generated._monitor_query_client import MonitorQueryClient
 
 from ._generated.models import BatchRequest, QueryBody as LogsQueryBody
-from ._helpers import get_authentication_policy, process_error, construct_iso8601, order_results
+from ._helpers import get_authentication_policy, construct_iso8601, order_results, process_error, process_prefer
 from ._models import LogsBatchQuery, LogsQueryResult
+from ._exceptions import  LogsQueryError, QueryPartialErrorException
 
 if TYPE_CHECKING:
     from azure.core.credentials import TokenCredential
@@ -76,6 +77,8 @@ def query(self, workspace_id, query, **kwargs):
         :keyword additional_workspaces: A list of workspaces that are included in the query.
          These can be qualified workspace names, workspace Ids, or Azure resource Ids.
         :paramtype additional_workspaces: list[str]
+        :keyword allow_partial_errors: Defaults to False. If set to true, partial errors are not thrown.
+        :paramtype allow_partial_errors: bool
         :return: LogsQueryResult, or the result of cls(response)
         :rtype: ~azure.monitor.query.LogsQueryResult
         :raises: ~azure.core.exceptions.HttpResponseError
@@ -89,6 +92,7 @@ def query(self, workspace_id, query, **kwargs):
             :dedent: 0
             :caption: Get a response for a single Log Query
         """
+        allow_partial_errors = kwargs.pop('allow_partial_errors', False)
         if 'timespan' not in kwargs:
             raise TypeError("query() missing 1 required keyword-only argument: 'timespan'")
         timespan = construct_iso8601(kwargs.pop('timespan'))
@@ -97,17 +101,7 @@ def query(self, workspace_id, query, **kwargs):
         server_timeout = kwargs.pop("server_timeout", None)
         workspaces = kwargs.pop("additional_workspaces", None)
 
-        prefer = ""
-        if server_timeout:
-            prefer += "wait=" + str(server_timeout)
-        if include_statistics:
-            if len(prefer) > 0:
-                prefer += ","
-            prefer += "include-statistics=true"
-        if include_visualization:
-            if len(prefer) > 0:
-                prefer += ","
-            prefer += "include-render=true"
+        prefer = process_prefer(server_timeout, include_statistics, include_visualization)
 
         body = LogsQueryBody(
             query=query,
@@ -117,14 +111,23 @@ def query(self, workspace_id, query, **kwargs):
         )
 
         try:
-            return LogsQueryResult._from_generated(self._query_op.execute( # pylint: disable=protected-access
+            generated_response = self._query_op.execute( # pylint: disable=protected-access
                 workspace_id=workspace_id,
                 body=body,
                 prefer=prefer,
                 **kwargs
-            ))
-        except HttpResponseError as e:
-            process_error(e)
+            )
+        except HttpResponseError as err:
+            process_error(err, LogsQueryError)
+        response = LogsQueryResult._from_generated(generated_response) # pylint: disable=protected-access
+        if not generated_response.error:
+            return response
+        if not allow_partial_errors:
+            raise QueryPartialErrorException(error=generated_response.error)
+        response.partial_error = LogsQueryError._from_generated( # pylint: disable=protected-access
+            generated_response.error
+            )
+        return response
 
     @distributed_trace
     def query_batch(self, queries, **kwargs):
@@ -136,6 +139,9 @@ def query_batch(self, queries, **kwargs):
 
         :param queries: The list of Kusto queries to execute.
         :type queries: list[dict] or list[~azure.monitor.query.LogsBatchQuery]
+        :keyword bool allow_partial_errors: If set to True, a `LogsQueryResult` object is returned
+         when a partial error occurs. The error can be accessed using the `partial_error`
+         attribute in the object.
         :return: List of LogsQueryResult, or the result of cls(response)
         :rtype: list[~azure.monitor.query.LogsQueryResult]
         :raises: ~azure.core.exceptions.HttpResponseError
@@ -149,6 +155,7 @@ def query_batch(self, queries, **kwargs):
             :dedent: 0
             :caption: Get a response for multiple Log Queries.
         """
+        allow_partial_errors = kwargs.pop('allow_partial_errors', False)
         try:
             queries = [LogsBatchQuery(**q) for q in queries]
         except (KeyError, TypeError):
@@ -161,7 +168,12 @@ def query_batch(self, queries, **kwargs):
         batch = BatchRequest(requests=queries)
         generated = self._query_op.batch(batch, **kwargs)
         mapping = {item.id: item for item in generated.responses}
-        return order_results(request_order, mapping, LogsQueryResult)
+        return order_results(
+            request_order,
+            mapping,
+            LogsQueryResult,
+            LogsQueryError,
+            allow_partial_errors)
 
     def close(self):
         # type: () -> None

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

@@ -165,21 +165,27 @@ class LogsQueryResult(object):
     :ivar visualization: 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.
     :vartype visualization: object
-    :ivar error: Any error info.
-    :vartype error: ~azure.core.exceptions.HttpResponseError
+    :ivar partial_error: Any error info. This is none except in the case where `allow_partial_errors`
+     is explicitly set to True.
+    :vartype partial_error: ~azure.core.exceptions.HttpResponseError
+    :ivar bool is_error: Boolean check for error item when iterating over list of
+        results. Always False for an instance of a LogsQueryResult.
     """
     def __init__(
         self,
         **kwargs
     ):
         self.tables = kwargs.get('tables', None)
-        self.error = kwargs.get('error', None)
+        self.partial_error = None
         self.statistics = kwargs.get('statistics', None)
         self.visualization = kwargs.get('visualization', None)
+        self.is_error = False
+
+    def __iter__(self):
+        return iter(self.tables)
 
     @classmethod
     def _from_generated(cls, generated):
-
         if not generated:
             return cls()
         tables = None
@@ -195,7 +201,6 @@ def _from_generated(cls, generated):
             tables=tables,
             statistics=generated.statistics,
             visualization=generated.render,
-            error=generated.error
         )