diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py index fd37c277cf17..b341a9bf082c 100644 --- a/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py +++ b/sdk/monitor/azure-monitor-query/azure/monitor/query/_exceptions.py @@ -42,7 +42,7 @@ def _from_generated(cls, generated): return None details = None if generated.details is not None: - details=[d.serialize() for d in generated.details] + details = [d.serialize() for d in generated.details] return cls( code=generated.code, message=generated.message, diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_helpers.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_helpers.py index d19d240bb27a..4124b6e3f21d 100644 --- a/sdk/monitor/azure-monitor-query/azure/monitor/query/_helpers.py +++ b/sdk/monitor/azure-monitor-query/azure/monitor/query/_helpers.py @@ -103,3 +103,17 @@ def process_error(error): if not error: return None raise HttpResponseError(message=error.message, response=error.response) + +def process_prefer(server_timeout, include_statistics, include_visualization): + 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" + return prefer diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/_logs_query_client.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/_logs_query_client.py index b2c8e4aa43cb..0e8af480096a 100644 --- a/sdk/monitor/azure-monitor-query/azure/monitor/query/_logs_query_client.py +++ b/sdk/monitor/azure-monitor-query/azure/monitor/query/_logs_query_client.py @@ -12,7 +12,7 @@ from ._generated._monitor_query_client import MonitorQueryClient from ._generated.models import BatchRequest, QueryBody as LogsQueryBody -from ._helpers import get_authentication_policy, construct_iso8601, order_results, process_error +from ._helpers import get_authentication_policy, construct_iso8601, order_results, process_error, process_prefer from ._models import LogsBatchQuery, LogsQueryResult from ._exceptions import LogsQueryError, QueryPartialErrorException @@ -112,6 +112,7 @@ def query(self, workspace_id, query, **kwargs): if len(prefer) > 0: prefer += "," prefer += "include-render=true" + prefer = process_prefer(server_timeout, include_statistics, include_visualization) body = LogsQueryBody( query=query, @@ -129,14 +130,15 @@ def query(self, workspace_id, query, **kwargs): ) except HttpResponseError as err: process_error(err) - response = LogsQueryResult._from_generated(generated_response) + response = LogsQueryResult._from_generated(generated_response) # pylint: disable=protected-access if not generated_response.error: return response - else: - if not allow_partial_errors: - raise QueryPartialErrorException(error=generated_response.error) - response.partial_error = LogsQueryError._from_generated(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): diff --git a/sdk/monitor/azure-monitor-query/azure/monitor/query/aio/_logs_query_client_async.py b/sdk/monitor/azure-monitor-query/azure/monitor/query/aio/_logs_query_client_async.py index b0ffbddb476d..8e1ac58dde26 100644 --- a/sdk/monitor/azure-monitor-query/azure/monitor/query/aio/_logs_query_client_async.py +++ b/sdk/monitor/azure-monitor-query/azure/monitor/query/aio/_logs_query_client_async.py @@ -13,7 +13,7 @@ from .._generated.aio._monitor_query_client import MonitorQueryClient from .._generated.models import BatchRequest, QueryBody as LogsQueryBody -from .._helpers import construct_iso8601, order_results, process_error +from .._helpers import construct_iso8601, order_results, process_error, process_prefer from .._models import LogsQueryResult, LogsBatchQuery from ._helpers_asyc import get_authentication_policy from .._exceptions import LogsQueryError, QueryPartialErrorException @@ -83,17 +83,7 @@ async def query( server_timeout = kwargs.pop("server_timeout", None) additional_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, @@ -111,14 +101,15 @@ async def query( ) except HttpResponseError as err: process_error(err) - response = LogsQueryResult._from_generated(generated_response) + response = LogsQueryResult._from_generated(generated_response) # pylint: disable=protected-access if not generated_response.error: return response - else: - if not allow_partial_errors: - raise QueryPartialErrorException(error=generated_response.error) - response.partial_error = LogsQueryError._from_generated(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_async async def query_batch(