diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index 1fa3d3db9934..cb6cb7e7b0d3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -26,8 +26,11 @@ """ from typing import Dict, Any, Optional import six +import requests +from requests.adapters import HTTPAdapter from azure.core.paging import ItemPaged # type: ignore from azure.core import PipelineClient # type: ignore +from azure.core.pipeline.transport import RequestsTransport from azure.core.pipeline.policies import ( # type: ignore ContentDecodePolicy, HeadersPolicy, @@ -148,6 +151,16 @@ def __init__( self._useMultipleWriteLocations = False self._global_endpoint_manager = global_endpoint_manager._GlobalEndpointManager(self) + # creating a requests session used for connection pooling and re-used by all requests + requests_session = requests.Session() + + transport = None + if self.connection_policy.ConnectionRetryConfiguration is not None: + adapter = HTTPAdapter(max_retries=self.connection_policy.ConnectionRetryConfiguration) + requests_session.mount('http://', adapter) + requests_session.mount('https://', adapter) + transport = RequestsTransport(session=requests_session) + proxies = kwargs.pop('proxies', {}) if self.connection_policy.ProxyConfiguration and self.connection_policy.ProxyConfiguration.Host: host = self.connection_policy.ProxyConfiguration.Host @@ -165,7 +178,7 @@ def __init__( NetworkTraceLoggingPolicy(**kwargs), ] - self.pipeline_client = PipelineClient(url_connection, "empty-config", policies=policies) + self.pipeline_client = PipelineClient(url_connection, "empty-config", transport=transport, policies=policies) # Query compatibility mode. # Allows to specify compatibility mode used by client when making query requests. Should be removed when diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/documents.py b/sdk/cosmos/azure-cosmos/azure/cosmos/documents.py index a2661a71bb20..02b80331b281 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/documents.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/documents.py @@ -372,6 +372,8 @@ class ConnectionPolicy(object): # pylint: disable=too-many-instance-attributes :ivar boolean UseMultipleWriteLocations: Flag to enable writes on any locations (regions) for geo-replicated database accounts in the azure Cosmos service. + :ivar (int or requests.packages.urllib3.util.retry) ConnectionRetryConfiguration: + Retry Configuration to be used for urllib3 connection retries. """ __defaultRequestTimeout = 60000 # milliseconds @@ -391,6 +393,7 @@ def __init__(self): self.RetryOptions = _retry_options.RetryOptions() self.DisableSSLVerification = False self.UseMultipleWriteLocations = False + self.ConnectionRetryConfiguration = None class _OperationType(object): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py index c4087542f003..77b523c35679 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/offer.py @@ -24,7 +24,7 @@ from typing import Dict, Any -class Offer(dict): +class Offer(object): """ Represents a offer in an Azure Cosmos DB SQL API container. To read and update offers use the associated methods on the :class:`Container`. diff --git a/sdk/cosmos/azure-cosmos/test/crud_tests.py b/sdk/cosmos/azure-cosmos/test/crud_tests.py index 14529aede445..1854b23fae78 100644 --- a/sdk/cosmos/azure-cosmos/test/crud_tests.py +++ b/sdk/cosmos/azure-cosmos/test/crud_tests.py @@ -53,6 +53,9 @@ from azure.cosmos.partition_key import PartitionKey import conftest from azure.cosmos import _retry_utility +from requests.packages.urllib3.util.retry import Retry +from requests.exceptions import ConnectionError + pytestmark = pytest.mark.cosmosEmulator @@ -1959,6 +1962,44 @@ def test_client_request_timeout(self): # client does a getDatabaseAccount on initialization, which will time out cosmos_client.CosmosClient(CRUDTests.host, CRUDTests.masterKey, "Session", connection_policy=connection_policy) + def test_client_request_timeout_when_connection_retry_configuration_specified(self): + connection_policy = documents.ConnectionPolicy() + # making timeout 0 ms to make sure it will throw + connection_policy.RequestTimeout = 0 + connection_policy.ConnectionRetryConfiguration = Retry( + total=3, + read=3, + connect=3, + backoff_factor=0.3, + status_forcelist=(500, 502, 504) + ) + with self.assertRaises(Exception): + # client does a getDatabaseAccount on initialization, which will time out + cosmos_client.CosmosClient(CRUDTests.host, CRUDTests.masterKey, "Session", connection_policy=connection_policy) + + def test_client_connection_retry_configuration(self): + total_time_for_two_retries = self.initialize_client_with_connection_retry_config(2) + total_time_for_three_retries = self.initialize_client_with_connection_retry_config(3) + self.assertGreater(total_time_for_three_retries, total_time_for_two_retries) + + def initialize_client_with_connection_retry_config(self, retries): + from azure.core.exceptions import ServiceRequestError + connection_policy = documents.ConnectionPolicy() + connection_policy.ConnectionRetryConfiguration = Retry( + total=retries, + read=retries, + connect=retries, + backoff_factor=0.3, + status_forcelist=(500, 502, 504) + ) + start_time = time.time() + try: + cosmos_client.CosmosClient("https://localhost:9999", CRUDTests.masterKey, "Session", connection_policy=connection_policy) + self.fail() + except ServiceRequestError as e: + end_time = time.time() + return end_time - start_time + def test_query_iterable_functionality(self): def __create_resources(client): """Creates resources for this test.