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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
|
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
import unittest
import uuid
import pytest
from aiohttp.client_exceptions import (ClientConnectionError, ClientConnectionResetError,
ClientOSError, ServerConnectionError)
from azure.core.exceptions import ServiceRequestError, ServiceResponseError
import test_config
from azure.cosmos import DatabaseAccount, _location_cache
from azure.cosmos._location_cache import RegionalRoutingContext
from azure.cosmos.aio import CosmosClient, _retry_utility_async, _global_endpoint_manager_async
from azure.cosmos.exceptions import CosmosHttpResponseError
@pytest.mark.cosmosEmulator
class TestServiceRetryPoliciesAsync(unittest.IsolatedAsyncioTestCase):
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.")
async def asyncSetUp(self):
self.client = CosmosClient(self.host, self.masterKey)
self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID)
self.created_container = self.created_database.get_container_client(self.TEST_CONTAINER_ID)
async def asyncTearDown(self):
self.connectionPolicy.ConnectionRetryConfiguration = None
await self.client.close()
async def test_service_request_retry_policy_async(self):
# ServiceRequestErrors will always retry, and will retry once per preferred region
async with CosmosClient(self.host, self.masterKey) as mock_client:
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
# Save the original function
self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync
# 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_async.ExecuteFunctionAsync = mf
await container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert mf.counter == 3
finally:
_retry_utility_async.ExecuteFunctionAsync = 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_async.ExecuteFunctionAsync = mf
await container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert mf.counter == 1
finally:
_retry_utility_async.ExecuteFunctionAsync = 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_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert mf.counter == 2
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
async def test_service_response_retry_policy_async(self):
# For ServiceResponseErrors, we only do cross region retries on read requests or on ClientConnectionErrors
# We also only do retries within the ConnectionRetryPolicy for the cases above
async with CosmosClient(self.host, self.masterKey) as mock_client:
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
# Save the original function
self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync
# 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 ClientConnectionError we retry
mf = self.MockExecuteServiceResponseException(AttributeError, None)
_retry_utility_async.ExecuteFunctionAsync = mf
await container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 3
finally:
_retry_utility_async.ExecuteFunctionAsync = 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(AttributeError, None)
_retry_utility_async.ExecuteFunctionAsync = mf
await container.read_item(created_item['id'], created_item['pk'])
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 1
finally:
_retry_utility_async.ExecuteFunctionAsync = 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(AttributeError, None)
_retry_utility_async.ExecuteFunctionAsync = mf
# Even though we have 2 preferred write endpoints,
# we will only run the exception once due to no retries on write requests
await 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_async.ExecuteFunctionAsync = self.original_execute_function
# If we do a write request with a ClientConnectionError,
# we will do cross-region retries like with read requests
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(ClientConnectionError, ClientConnectionError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 2
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
async def test_service_request_connection_retry_policy_async(self):
# Mock the client retry policy to see the same-region retries that happen there
exception = ServiceRequestError("mock exception")
exception.exc_type = Exception
connection_policy = self.connectionPolicy
connection_retry_policy = test_config.MockConnectionRetryPolicyAsync(resource_type="docs", error=exception)
connection_policy.ConnectionRetryConfiguration = connection_retry_policy
async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy) as mock_client:
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:
await 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:
await container.read_item("some_id", "some_pk")
pytest.fail("Exception was not raised.")
except ServiceRequestError:
assert connection_retry_policy.counter == 3
async def test_service_response_connection_retry_policy_async(self):
# Mock the client retry policy to see the same-region retries that happen there
exception = ServiceResponseError("mock exception")
exception.exc_type = Exception
connection_policy = self.connectionPolicy
connection_retry_policy = test_config.MockConnectionRetryPolicyAsync(resource_type="docs", error=exception)
connection_policy.ConnectionRetryConfiguration = connection_retry_policy
async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy) as mock_client:
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:
await 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:
await container.read_item("some_id", "some_pk")
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert connection_retry_policy.counter == 3
async def test_service_response_errors_async(self):
# Test for errors that are subclasses of ClientConnectionError for write requests
# Save the original ExecuteAsyncFunction function
self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync
async with CosmosClient(self.host, self.masterKey) as mock_client:
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
await container.read()
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]
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:
# Start with a normal ServiceResponseException with no special casing
mf = self.MockExecuteServiceResponseException(AttributeError, AttributeError())
_retry_utility_async.ExecuteFunctionAsync = mf
# Even though we have 2 preferred write endpoints,
# we will only run the exception once due to no retries on write requests
await 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_async.ExecuteFunctionAsync = self.original_execute_function
# Now we test the base ClientConnectionError to see in-region retry
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 2
assert len(original_location_cache.location_unavailability_info_by_endpoint) == 1
host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host)
assert host_unavailable is not None
assert len(host_unavailable.get('operationType')) == 1
assert 'Write' in host_unavailable.get('operationType')
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
await container.read()
# We send another request to the same error - since we marked one of the two in-region endpoints unavailable we don't retry
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 1
host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host)
assert host_unavailable is not None
assert len(host_unavailable.get('operationType')) == 1
assert 'Write' in host_unavailable.get('operationType')
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
# Reset the location cache's unavailable endpoints in order to try the same with other exceptions
original_location_cache.location_unavailability_info_by_endpoint = {}
original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
# Now we test ClientConnectionResetError, the subclass of ClientConnectionError
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(ClientConnectionResetError, ClientConnectionResetError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 2
assert len(original_location_cache.location_unavailability_info_by_endpoint) == 1
host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host)
assert host_unavailable is not None
assert len(host_unavailable.get('operationType')) == 1
assert 'Write' in host_unavailable.get('operationType')
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
# Reset the location cache's unavailable endpoints in order to try the same with other exceptions
original_location_cache.location_unavailability_info_by_endpoint = {}
original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
# Now we test ServerConnectionError, the subclass of ClientConnectionError
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(ServerConnectionError, ServerConnectionError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 2
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
# Reset the location cache's unavailable endpoints in order to try the same with other exceptions
original_location_cache.location_unavailability_info_by_endpoint = {}
original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
# Now we test ClientOSError, the subclass of ClientConnectionError
try:
# Reset the function to reset the counter
mf = self.MockExecuteServiceResponseException(ClientOSError, ClientOSError())
_retry_utility_async.ExecuteFunctionAsync = mf
await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
pytest.fail("Exception was not raised.")
except ServiceResponseError:
assert mf.counter == 2
finally:
_retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
async def test_global_endpoint_manager_retry_async(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_async._GlobalEndpointManager._GetDatabaseAccountStub
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
connection_policy = self.connectionPolicy
connection_retry_policy = test_config.MockConnectionRetryPolicyAsync(resource_type="docs", error=exception)
connection_policy.ConnectionRetryConfiguration = connection_retry_policy
async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy,
preferred_locations=[self.REGION1, self.REGION2]) as mock_client:
db = mock_client.get_database_client(self.TEST_DATABASE_ID)
container = db.get_container_client(self.TEST_CONTAINER_ID)
try:
await 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
except CosmosHttpResponseError as e:
print(e)
finally:
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub
# Now we try with a read request - reset the policy to reset the counter
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
connection_retry_policy.request_endpoints = []
try:
await 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_async._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, inner_exception):
self.err_type = err_type
self.inner_exception = inner_exception
self.counter = 0
def __call__(self, func, *args, **kwargs):
self.counter = self.counter + 1
exception = ServiceResponseError("mock exception")
exception.exc_type = self.err_type
exception.inner_exception = self.inner_exception
raise exception
async 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 = self.host.replace("localhost", "127.0.0.1")
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
|
