1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
|
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
import unittest
import uuid
import pytest
from azure.core.exceptions import ServiceRequestError, ServiceResponseError
import test_config
from azure.cosmos import (CosmosClient, _retry_utility, DatabaseAccount, _global_endpoint_manager,
_location_cache)
from azure.cosmos._location_cache import RegionalRoutingContext
@pytest.mark.cosmosEmulator
class TestServiceRetryPolicies(unittest.TestCase):
host = test_config.TestConfig.host
masterKey = test_config.TestConfig.masterKey
connectionPolicy = test_config.TestConfig.connectionPolicy
TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID
TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID
REGION1 = "West US"
REGION2 = "East US"
REGION3 = "West US 2"
REGIONAL_ENDPOINT = RegionalRoutingContext(host, host)
@classmethod
def setUpClass(cls):
if (cls.masterKey == '[YOUR_KEY_HERE]' or
cls.host == '[YOUR_ENDPOINT_HERE]'):
raise Exception(
"You must specify your Azure Cosmos account values for "
"'masterKey' and 'host' at the top of this class to run the "
"tests.")
cls.client = CosmosClient(cls.host, cls.masterKey)
cls.created_database = cls.client.get_database_client(cls.TEST_DATABASE_ID)
cls.created_container = cls.created_database.get_container_client(cls.TEST_CONTAINER_ID)
def test_service_request_retry_policy(self):
mock_client = CosmosClient(self.host, self.masterKey)
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
# Save the original function
self.original_execute_function = _retry_utility.ExecuteFunction
# Change the location cache to have 3 preferred read regions and 3 available read endpoints by location
original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache
original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3]
original_location_cache.available_read_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT,
self.REGION2: self.REGIONAL_ENDPOINT,
self.REGION3: self.REGIONAL_ENDPOINT}
original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
self.REGIONAL_ENDPOINT]
try:
# Mock the function to return the ServiceRequestException we retry
mf = self.MockExecuteServiceRequestException()
_retry_utility.ExecuteFunction = mf
container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert mf.counter == 3
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
# Now we change the location cache to have only 1 preferred read region
original_location_cache.account_read_locations = [self.REGION1]
original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT]
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceRequestException()
_retry_utility.ExecuteFunction = mf
container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert mf.counter == 1
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
# Now we try it out with a write request
original_location_cache.account_write_locations = [self.REGION1, self.REGION2]
original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT,
self.REGION2: self.REGIONAL_ENDPOINT}
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceRequestException()
_retry_utility.ExecuteFunction = mf
container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceRequestError:
# Should retry twice in each region
assert mf.counter == 2
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
def test_service_response_retry_policy(self):
mock_client = CosmosClient(self.host, self.masterKey)
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
# Save the original function
self.original_execute_function = _retry_utility.ExecuteFunction
# Change the location cache to have 3 preferred read regions and 3 available read endpoints by location
original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache
original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3]
original_location_cache.available_read_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT,
self.REGION2: self.REGIONAL_ENDPOINT,
self.REGION3: self.REGIONAL_ENDPOINT}
original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
self.REGIONAL_ENDPOINT]
try:
# Mock the function to return the ServiceResponseException we retry
mf = self.MockExecuteServiceResponseException(Exception)
_retry_utility.ExecuteFunction = mf
container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 3
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
# Now we change the location cache to have only 1 preferred read region
original_location_cache.account_read_locations = [self.REGION1]
original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT]
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(Exception)
_retry_utility.ExecuteFunction = mf
container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 1
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
# Now we try it out with a write request
original_location_cache.account_write_locations = [self.REGION1, self.REGION2]
original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT,
self.REGION2: self.REGIONAL_ENDPOINT}
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(Exception)
_retry_utility.ExecuteFunction = mf
# Even though we have 2 preferred write endpoints,
# we will only run the exception once due to no retries on write requests
container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 1
finally:
_retry_utility.ExecuteFunction = self.original_execute_function
def test_service_request_connection_retry_policy(self):
# Mock the client retry policy to see the same-region retries that happen there
exception = ServiceRequestError("mock exception")
exception.exc_type = Exception
connection_retry_policy = test_config.MockConnectionRetryPolicy(resource_type="docs", error=exception)
mock_client = CosmosClient(self.host, self.masterKey, connection_retry_policy=connection_retry_policy)
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
# We retry ServiceRequestExceptions 3 times in the to the same endpoint before raising the exception
# regardless of operation type
try:
container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert connection_retry_policy.counter == 3
try:
container.read_item("some_id", "some_pk")
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert connection_retry_policy.counter == 3
def test_service_response_connection_retry_policy(self):
# Mock the client retry policy to see the same-region retries that happen there
exception = ServiceResponseError("mock exception")
exception.exc_type = Exception
connection_retry_policy = test_config.MockConnectionRetryPolicy(resource_type="docs", error=exception)
mock_client = CosmosClient(self.host, self.masterKey, connection_retry_policy=connection_retry_policy)
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
# We retry ServiceResponseExceptions 3 times in the to the same endpoint before raising the exception
# for read operations, but 0 times for write requests
try:
container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert connection_retry_policy.counter == 0
try:
container.read_item("some_id", "some_pk")
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert connection_retry_policy.counter == 3
def test_global_endpoint_manager_retry(self):
# For this test we mock both the ConnectionRetryPolicy and the GetDatabaseAccountStub
# - ConnectionRetryPolicy allows us to raise Service exceptions only for chosen requests and track endpoints used
# - GetDatabaseAccountStub allows us to receive any number of endpoints for that call independent of account used
exception = ServiceRequestError("mock exception")
exception.exc_type = Exception
self.original_get_database_account_stub = _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
connection_retry_policy = test_config.MockConnectionRetryPolicy(resource_type="docs", error=exception)
mock_client = CosmosClient(self.host, self.masterKey, connection_retry_policy=connection_retry_policy,
preferred_locations=[self.REGION1, self.REGION2])
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
try:
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert connection_retry_policy.counter == 3
# 4 total requests for each in-region (hub -> write locational endpoint)
assert len(connection_retry_policy.request_endpoints) == 8
finally:
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub
# Now we try with a read request - reset the policy to reset the counter
connection_retry_policy.request_endpoints = []
try:
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
container.read_item("some_id", "some_pk")
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert connection_retry_policy.counter == 3
# 4 total requests in each main region (preferred read region 1 -> preferred read region 2)
assert len(connection_retry_policy.request_endpoints) == 8
finally:
_global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub
class MockExecuteServiceRequestException(object):
def __init__(self):
self.counter = 0
def __call__(self, func, *args, **kwargs):
self.counter = self.counter + 1
exception = ServiceRequestError("mock exception")
exception.exc_type = Exception
raise exception
class MockExecuteServiceResponseException(object):
def __init__(self, err_type):
self.err_type = err_type
self.counter = 0
def __call__(self, func, *args, **kwargs):
self.counter = self.counter + 1
exception = ServiceResponseError("mock exception")
exception.exc_type = self.err_type
raise exception
def MockGetDatabaseAccountStub(self, endpoint):
read_regions = ["West US", "East US"]
read_locations = []
for loc in read_regions:
read_locations.append({'databaseAccountEndpoint': endpoint, 'name': loc})
write_regions = ["West US"]
write_locations = []
for loc in write_regions:
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(endpoint, loc)
write_locations.append({'databaseAccountEndpoint': locational_endpoint, 'name': loc})
multi_write = False
db_acc = DatabaseAccount()
db_acc.DatabasesLink = "/dbs/"
db_acc.MediaLink = "/media/"
db_acc._ReadableLocations = read_locations
db_acc._WritableLocations = write_locations
db_acc._EnableMultipleWritableLocations = multi_write
db_acc.ConsistencyPolicy = {"defaultConsistencyLevel": "Session"}
return db_acc
|
