Misc fixes for Cosmos SDK

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

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):

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`.

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.