# The MIT License (MIT)
# Copyright (c) 2014 Microsoft Corporation
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""Internal methods for executing functions in the Azure Cosmos database service.
"""
import json
import time
import logging
from typing import Optional
from azure.core.exceptions import AzureError, ClientAuthenticationError, ServiceRequestError, ServiceResponseError
from azure.core.pipeline import PipelineRequest
from azure.core.pipeline.policies import RetryPolicy
from . import _container_recreate_retry_policy, _health_check_retry_policy
from . import _default_retry_policy
from . import _endpoint_discovery_retry_policy
from . import _gone_retry_policy
from . import _resource_throttle_retry_policy
from . import _service_request_retry_policy, _service_response_retry_policy
from . import _session_retry_policy
from . import _timeout_failover_retry_policy
from . import exceptions
from ._constants import _Constants
from .documents import _OperationType
from .exceptions import CosmosHttpResponseError
from .http_constants import HttpHeaders, StatusCodes, SubStatusCodes, ResourceType
from ._cosmos_http_logging_policy import _log_diagnostics_error
# pylint: disable=protected-access, disable=too-many-lines, disable=too-many-statements, disable=too-many-branches
# args [0] is the request object
# args [1] is the connection policy
# args [2] is the pipeline client
# args [3] is the http request
def Execute(client, global_endpoint_manager, function, *args, **kwargs): # pylint: disable=too-many-locals
"""Executes the function with passed parameters applying all retry policies
:param object client:
Document client instance
:param object global_endpoint_manager:
Instance of _GlobalEndpointManager class
:param function function:
Function to be called wrapped with retries
:param list args:
:returns: the result of running the passed in function as a (result, headers) tuple
:rtype: tuple of (dict, dict)
"""
# Capture the client timeout and start time at the beginning
timeout = kwargs.get('timeout')
operation_start_time = kwargs.get(_Constants.OperationStartTime, time.time())
# Track the last error for chaining
last_error = None
pk_range_wrapper = None
if args and global_endpoint_manager.is_circuit_breaker_applicable(args[0]):
pk_range_wrapper = global_endpoint_manager.create_pk_range_wrapper(args[0])
# instantiate all retry policies here to be applied for each request execution
endpointDiscovery_retry_policy = _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy(
client.connection_policy, global_endpoint_manager, *args
)
health_check_retry_policy = _health_check_retry_policy.HealthCheckRetryPolicy(
client.connection_policy, *args
)
resourceThrottle_retry_policy = _resource_throttle_retry_policy.ResourceThrottleRetryPolicy(
client.connection_policy.RetryOptions.MaxRetryAttemptCount,
client.connection_policy.RetryOptions.FixedRetryIntervalInMilliseconds,
client.connection_policy.RetryOptions.MaxWaitTimeInSeconds,
)
defaultRetry_policy = _default_retry_policy.DefaultRetryPolicy(*args)
sessionRetry_policy = _session_retry_policy._SessionRetryPolicy(
client.connection_policy.EnableEndpointDiscovery, global_endpoint_manager, pk_range_wrapper, *args
)
partition_key_range_gone_retry_policy = _gone_retry_policy.PartitionKeyRangeGoneRetryPolicy(client, *args)
timeout_failover_retry_policy = _timeout_failover_retry_policy._TimeoutFailoverRetryPolicy(
client.connection_policy, global_endpoint_manager, pk_range_wrapper, *args
)
service_response_retry_policy = _service_response_retry_policy.ServiceResponseRetryPolicy(
client.connection_policy, global_endpoint_manager, pk_range_wrapper, *args,
)
service_request_retry_policy = _service_request_retry_policy.ServiceRequestRetryPolicy(
client.connection_policy, global_endpoint_manager, pk_range_wrapper, *args,
)
# Get logger
logger = kwargs.get("logger", logging.getLogger("azure.cosmos._retry_utility"))
# HttpRequest we would need to modify for Container Recreate Retry Policy
request = None
if args and len(args) > 3:
# Reference HttpRequest instance in args
request = args[3]
container_recreate_retry_policy = _container_recreate_retry_policy.ContainerRecreateRetryPolicy(
client, client._container_properties_cache, request, *args)
else:
container_recreate_retry_policy = _container_recreate_retry_policy.ContainerRecreateRetryPolicy(
client, client._container_properties_cache, None, *args)
while True:
start_time = time.time()
# Check timeout before executing function
if timeout:
elapsed = time.time() - operation_start_time
if elapsed >= timeout:
raise exceptions.CosmosClientTimeoutError(error=last_error)
try:
if args:
result = ExecuteFunction(function, global_endpoint_manager, *args, **kwargs)
global_endpoint_manager.record_success(args[0])
else:
result = ExecuteFunction(function, *args, **kwargs)
# Check timeout after successful execution
if timeout:
elapsed = time.time() - operation_start_time
if elapsed >= timeout:
raise exceptions.CosmosClientTimeoutError(error=last_error)
if not client.last_response_headers:
client.last_response_headers = {}
# setting the throttle related response headers before returning the result
client.last_response_headers[
HttpHeaders.ThrottleRetryCount
] = resourceThrottle_retry_policy.current_retry_attempt_count
client.last_response_headers[
HttpHeaders.ThrottleRetryWaitTimeInMs
] = resourceThrottle_retry_policy.cumulative_wait_time_in_milliseconds
# TODO: It is better to raise Exceptions manually in the method related to the request,
# a rework of retry would be needed to be able to retry exceptions raised that way.
# for now raising a manual exception here should allow it to be retried.
# If container does not have throughput, results will return empty list.
# We manually raise a 404. We raise it here, so we can handle it in retry utilities.
if result and isinstance(result[0], dict) and 'Offers' in result[0] and \
not result[0]['Offers'] and request.method == 'POST':
# Grab the link used for getting throughput properties to add to message.
link = json.loads(request.body)["parameters"][0]["value"]
response = exceptions._InternalCosmosException(status_code=StatusCodes.NOT_FOUND,
headers={HttpHeaders.SubStatus:
SubStatusCodes.THROUGHPUT_OFFER_NOT_FOUND})
e_offer = exceptions.CosmosResourceNotFoundError(
status_code=StatusCodes.NOT_FOUND,
message="Could not find ThroughputProperties for container " + link,
response=response)
response_headers = result[1] if len(result) > 1 else {}
logger_attributes = {
"duration": time.time() - start_time,
"verb": request.method,
"status_code": e_offer.status_code,
"sub_status_code": e_offer.sub_status,
}
_log_diagnostics_error(client._enable_diagnostics_logging, request, response_headers, e_offer,
logger_attributes, global_endpoint_manager, logger=logger)
raise e_offer
return result
except exceptions.CosmosHttpResponseError as e:
last_error = e
if request:
# update session token for relevant operations
client._UpdateSessionIfRequired(request.headers, {}, e.headers)
if request and _has_database_account_header(request.headers):
retry_policy = health_check_retry_policy
# Re-assign retry policy based on error code
elif e.status_code == StatusCodes.FORBIDDEN and e.sub_status in\
[SubStatusCodes.DATABASE_ACCOUNT_NOT_FOUND, SubStatusCodes.WRITE_FORBIDDEN]:
retry_policy = endpointDiscovery_retry_policy
elif e.status_code == StatusCodes.TOO_MANY_REQUESTS:
retry_policy = resourceThrottle_retry_policy
elif (
e.status_code == StatusCodes.NOT_FOUND
and e.sub_status
and e.sub_status == SubStatusCodes.READ_SESSION_NOTAVAILABLE
):
retry_policy = sessionRetry_policy
elif exceptions._partition_range_is_gone(e):
retry_policy = partition_key_range_gone_retry_policy
elif exceptions._container_recreate_exception(e):
retry_policy = container_recreate_retry_policy
# Before we retry if retry policy is container recreate, we need refresh the cache of the
# container properties and pass in the new RID in the headers.
client._refresh_container_properties_cache(retry_policy.container_link)
if e.sub_status != SubStatusCodes.COLLECTION_RID_MISMATCH and retry_policy.check_if_rid_different(
retry_policy.container_link, client._container_properties_cache, retry_policy.container_rid):
retry_policy.refresh_container_properties_cache = False
else:
cached_container = client._container_properties_cache[retry_policy.container_link]
# If partition key value was previously extracted from the document definition
# reattempt to extract partition key with updated partition key definition
if retry_policy.should_extract_partition_key(cached_container):
new_partition_key = retry_policy._extract_partition_key(
client, container_cache=cached_container, body=request.body
)
request.headers[HttpHeaders.PartitionKey] = new_partition_key
# If getting throughput, we have to replace the container link received from stale cache
# with refreshed cache
if retry_policy.should_update_throughput_link(request.body, cached_container):
new_body = retry_policy._update_throughput_link(request.body)
request.body = new_body
retry_policy.container_rid = cached_container["_rid"]
request.headers[retry_policy._intended_headers] = retry_policy.container_rid
elif e.status_code == StatusCodes.REQUEST_TIMEOUT or e.status_code >= StatusCodes.INTERNAL_SERVER_ERROR:
if args:
# record the failure for circuit breaker tracking
global_endpoint_manager.record_failure(args[0])
retry_policy = timeout_failover_retry_policy
else:
retry_policy = defaultRetry_policy
# If none of the retry policies applies or there is no retry needed, set the
# throttle related response headers and re-throw the exception back arg[0]
# is the request. It needs to be modified for write forbidden exception
if not retry_policy.ShouldRetry(e):
if not client.last_response_headers:
client.last_response_headers = {}
client.last_response_headers[
HttpHeaders.ThrottleRetryCount
] = resourceThrottle_retry_policy.current_retry_attempt_count
client.last_response_headers[
HttpHeaders.ThrottleRetryWaitTimeInMs
] = resourceThrottle_retry_policy.cumulative_wait_time_in_milliseconds
if args and args[0].should_clear_session_token_on_session_read_failure:
client.session.clear_session_token(client.last_response_headers)
raise
# Now check timeout before retrying
if timeout:
elapsed = time.time() - operation_start_time
if elapsed >= timeout:
raise exceptions.CosmosClientTimeoutError(error=last_error)
# Wait for retry_after_in_milliseconds time before the next retry
time.sleep(retry_policy.retry_after_in_milliseconds / 1000.0)
except ServiceRequestError as e:
if request and _has_database_account_header(request.headers):
if not health_check_retry_policy.ShouldRetry(e):
raise e
else:
_handle_service_request_retries(client, service_request_retry_policy, e, *args)
except ServiceResponseError as e:
if request and _has_database_account_header(request.headers):
if not health_check_retry_policy.ShouldRetry(e):
raise e
else:
if args:
global_endpoint_manager.record_failure(args[0])
_handle_service_response_retries(request, client, service_response_retry_policy, e, *args)
def ExecuteFunction(function, *args, **kwargs):
"""Stub method so that it can be used for mocking purposes as well.
:param Callable function: the function to execute.
:param list args: the explicit arguments for the function.
:returns: the result of executing the function with the passed in arguments
:rtype: tuple(dict, dict)
"""
return function(*args, **kwargs)
def _has_read_retryable_headers(request_headers):
if _OperationType.IsReadOnlyOperation(request_headers.get(HttpHeaders.ThinClientProxyOperationType)):
return True
return False
def _has_database_account_header(request_headers):
if request_headers.get(HttpHeaders.ThinClientProxyResourceType) == ResourceType.DatabaseAccount:
return True
return False
def _handle_service_request_retries(
client,
request_retry_policy,
exception,
*args
):
# we resolve the request endpoint to the next preferred region
# once we are out of preferred regions we stop retrying
retry_policy = request_retry_policy
if not retry_policy.ShouldRetry():
if args and args[0].should_clear_session_token_on_session_read_failure and client.session:
client.session.clear_session_token(client.last_response_headers)
raise exception
def _handle_service_response_retries(request, client, response_retry_policy, exception, *args):
if request and (_has_read_retryable_headers(request.headers) or (args and is_write_retryable(args[0], client))):
# we resolve the request endpoint to the next preferred region
# once we are out of preferred regions we stop retrying
retry_policy = response_retry_policy
if not retry_policy.ShouldRetry():
if args and args[0].should_clear_session_token_on_session_read_failure and client.session:
client.session.clear_session_token(client.last_response_headers)
raise exception
else:
raise exception
def is_write_retryable(request_params, client):
return (request_params.retry_write > 0 or
(client.connection_policy.RetryNonIdempotentWrites > 0 and
not request_params.operation_type == _OperationType.Patch))
def _configure_timeout(request: PipelineRequest, absolute: Optional[int], per_request: int) -> None:
if absolute is not None:
if absolute <= 0:
raise exceptions.CosmosClientTimeoutError()
if per_request:
# Both socket timeout and client timeout have been provided - use the shortest value.
request.context.options['connection_timeout'] = min(per_request, absolute)
else:
# Only client timeout provided.
request.context.options['connection_timeout'] = absolute
elif per_request:
# Only socket timeout provided.
request.context.options['connection_timeout'] = per_request
[docs]
class ConnectionRetryPolicy(RetryPolicy):
def __init__(self, **kwargs):
clean_kwargs = {k: v for k, v in kwargs.items() if v is not None}
super(ConnectionRetryPolicy, self).__init__(**clean_kwargs)
[docs]
def send(self, request):
"""Sends the PipelineRequest object to the next policy. Uses retry settings if necessary.
Also enforces an absolute client-side timeout that spans multiple retry attempts.
:param request: The PipelineRequest object
:type request: ~azure.core.pipeline.PipelineRequest
:return: Returns the PipelineResponse or raises error if maximum retries exceeded.
:rtype: ~azure.core.pipeline.PipelineResponse
:raises ~azure.core.exceptions.AzureError: Maximum retries exceeded.
:raises ~azure.cosmos.exceptions.CosmosClientTimeoutError: Specified timeout exceeded.
:raises ~azure.core.exceptions.ClientAuthenticationError: Authentication failed.
"""
absolute_timeout = request.context.options.pop('timeout', None)
per_request_timeout = request.context.options.pop('connection_timeout', 0)
request_params = request.context.options.pop('request_params', None)
global_endpoint_manager = request.context.options.pop('global_endpoint_manager', None)
retry_error = None
retry_active = True
response = None
retry_settings = self.configure_retries(request.context.options)
while retry_active:
start_time = time.time()
try:
_configure_timeout(request, absolute_timeout, per_request_timeout)
response = self.next.send(request)
break
except ClientAuthenticationError: # pylint:disable=try-except-raise
# the authentication policy failed such that the client's request can't
# succeed--we'll never have a response to it, so propagate the exception
raise
except exceptions.CosmosClientTimeoutError as timeout_error:
timeout_error.inner_exception = retry_error
timeout_error.response = response
timeout_error.history = retry_settings['history']
raise
except ServiceRequestError as err:
retry_error = err
# the request ran into a socket timeout or failed to establish a new connection
# since request wasn't sent, raise exception immediately to be dealt with in client retry policies
# This logic is based on the _retry.py file from azure-core
if (not _has_database_account_header(request.http_request.headers)
and not request_params.healthy_tentative_location):
if retry_settings['connect'] > 0:
retry_active = self.increment(retry_settings, response=request, error=err)
if retry_active:
self.sleep(retry_settings, request.context.transport)
continue
raise err
except ServiceResponseError as err:
retry_error = err
# Only read operations can be safely retried with ServiceResponseError
if (not _has_read_retryable_headers(request.http_request.headers) or
_has_database_account_header(request.http_request.headers) or
request_params.healthy_tentative_location):
raise err
# This logic is based on the _retry.py file from azure-core
if retry_settings['read'] > 0:
# record the failure for circuit breaker tracking for retries in connection retry policy
# retries in the execute function will mark those failures
global_endpoint_manager.record_failure(request_params)
retry_active = self.increment(retry_settings, response=request, error=err)
if retry_active:
self.sleep(retry_settings, request.context.transport)
continue
raise err
except CosmosHttpResponseError as err:
raise err
except AzureError as err:
retry_error = err
if (_has_database_account_header(request.http_request.headers) or
request_params.healthy_tentative_location):
raise err
if _has_read_retryable_headers(request.http_request.headers) and retry_settings['read'] > 0:
global_endpoint_manager.record_failure(request_params)
retry_active = self.increment(retry_settings, response=request, error=err)
if retry_active:
self.sleep(retry_settings, request.context.transport)
continue
raise err
finally:
end_time = time.time()
if absolute_timeout:
absolute_timeout -= (end_time - start_time)
self.update_context(response.context, retry_settings)
return response