Skip to content

Commit 4398b5b

Browse files
author
Srinath Narayanan
committed
added support for urllib3 connection retry
1 parent d25510b commit 4398b5b

File tree

3 files changed

+58
-1
lines changed

3 files changed

+58
-1
lines changed

‎sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,11 @@
2626
"""
2727
from typing import Dict, Any, Optional
2828
import six
29+
import requests
30+
from requests.adapters import HTTPAdapter
2931
from azure.core.paging import ItemPaged # type: ignore
3032
from azure.core import PipelineClient # type: ignore
33+
from azure.core.pipeline.transport import RequestsTransport
3134
from azure.core.pipeline.policies import ( # type: ignore
3235
ContentDecodePolicy,
3336
HeadersPolicy,
@@ -148,6 +151,16 @@ def __init__(
148151
self._useMultipleWriteLocations = False
149152
self._global_endpoint_manager = global_endpoint_manager._GlobalEndpointManager(self)
150153

154+
# creating a requests session used for connection pooling and re-used by all requests
155+
requests_session = requests.Session()
156+
157+
transport = None
158+
if self.connection_policy.ConnectionRetryConfiguration is not None:
159+
adapter = HTTPAdapter(max_retries=self.connection_policy.ConnectionRetryConfiguration)
160+
requests_session.mount('http://', adapter)
161+
requests_session.mount('https://', adapter)
162+
transport = RequestsTransport(session=requests_session)
163+
151164
proxies = kwargs.pop('proxies', {})
152165
if self.connection_policy.ProxyConfiguration and self.connection_policy.ProxyConfiguration.Host:
153166
host = self.connection_policy.ProxyConfiguration.Host
@@ -165,7 +178,7 @@ def __init__(
165178
NetworkTraceLoggingPolicy(**kwargs),
166179
]
167180

168-
self.pipeline_client = PipelineClient(url_connection, "empty-config", policies=policies)
181+
self.pipeline_client = PipelineClient(url_connection, "empty-config", transport=transport, policies=policies)
169182

170183
# Query compatibility mode.
171184
# Allows to specify compatibility mode used by client when making query requests. Should be removed when

‎sdk/cosmos/azure-cosmos/azure/cosmos/documents.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,8 @@ class ConnectionPolicy(object): # pylint: disable=too-many-instance-attributes
372372
:ivar boolean UseMultipleWriteLocations:
373373
Flag to enable writes on any locations (regions) for geo-replicated database accounts
374374
in the azure Cosmos service.
375+
:ivar (int or requests.packages.urllib3.util.retry) ConnectionRetryConfiguration:
376+
Retry Configuration to be used for urllib3 connection retries.
375377
"""
376378

377379
__defaultRequestTimeout = 60000 # milliseconds
@@ -391,6 +393,7 @@ def __init__(self):
391393
self.RetryOptions = _retry_options.RetryOptions()
392394
self.DisableSSLVerification = False
393395
self.UseMultipleWriteLocations = False
396+
self.ConnectionRetryConfiguration = None
394397

395398

396399
class _OperationType(object):

‎sdk/cosmos/azure-cosmos/test/crud_tests.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@
5353
from azure.cosmos.partition_key import PartitionKey
5454
import conftest
5555
from azure.cosmos import _retry_utility
56+
from requests.packages.urllib3.util.retry import Retry
57+
from requests.exceptions import ConnectionError
58+
5659

5760
pytestmark = pytest.mark.cosmosEmulator
5861

@@ -1959,6 +1962,44 @@ def test_client_request_timeout(self):
19591962
# client does a getDatabaseAccount on initialization, which will time out
19601963
cosmos_client.CosmosClient(CRUDTests.host, CRUDTests.masterKey, "Session", connection_policy=connection_policy)
19611964

1965+
def test_client_request_timeout_when_connection_retry_configuration_specified(self):
1966+
connection_policy = documents.ConnectionPolicy()
1967+
# making timeout 0 ms to make sure it will throw
1968+
connection_policy.RequestTimeout = 0
1969+
connection_policy.ConnectionRetryConfiguration = Retry(
1970+
total=3,
1971+
read=3,
1972+
connect=3,
1973+
backoff_factor=0.3,
1974+
status_forcelist=(500, 502, 504)
1975+
)
1976+
with self.assertRaises(Exception):
1977+
# client does a getDatabaseAccount on initialization, which will time out
1978+
cosmos_client.CosmosClient(CRUDTests.host, CRUDTests.masterKey, "Session", connection_policy=connection_policy)
1979+
1980+
def test_client_connection_retry_configuration(self):
1981+
total_time_for_two_retries = self.initialize_client_with_connection_retry_config(2)
1982+
total_time_for_three_retries = self.initialize_client_with_connection_retry_config(3)
1983+
self.assertGreater(total_time_for_three_retries, total_time_for_two_retries)
1984+
1985+
def initialize_client_with_connection_retry_config(self, retries):
1986+
from azure.core.exceptions import ServiceRequestError
1987+
connection_policy = documents.ConnectionPolicy()
1988+
connection_policy.ConnectionRetryConfiguration = Retry(
1989+
total=retries,
1990+
read=retries,
1991+
connect=retries,
1992+
backoff_factor=0.3,
1993+
status_forcelist=(500, 502, 504)
1994+
)
1995+
start_time = time.time()
1996+
try:
1997+
cosmos_client.CosmosClient("https://localhost:9999", CRUDTests.masterKey, "Session", connection_policy=connection_policy)
1998+
self.fail()
1999+
except ServiceRequestError as e:
2000+
end_time = time.time()
2001+
return end_time - start_time
2002+
19622003
def test_query_iterable_functionality(self):
19632004
def __create_resources(client):
19642005
"""Creates resources for this test.

0 commit comments

Comments
 (0)