def test_default_retry_policy_for_create(self):
connection_policy = documents.ConnectionPolicy()
client = document_client.DocumentClient(
Test_retry_policy_tests.host,
{'masterKey': Test_retry_policy_tests.masterKey},
connection_policy)
document_definition = {
'id': 'doc',
'name': 'sample document',
'key': 'value'
}
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunctionConnectionReset
created_document = {}
try:
created_document = client.CreateDocument(
self.created_collection['_self'], document_definition)
except Exception as err:
self.assertEqual(err.status_code, 10054)
self.assertDictEqual(created_document, {})
self.assertEqual(self.counter, 7)
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
def test_resource_throttle_retry_policy_max_wait_time(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.RetryOptions = retry_options.RetryOptions(5, 2000, 3)
client = document_client.DocumentClient(
Test_retry_policy_tests.host,
{'masterKey': Test_retry_policy_tests.masterKey},
connection_policy)
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunction
document_definition = {
'id': 'doc',
'name': 'sample document',
'key': 'value'
}
try:
client.CreateDocument(self.created_collection['_self'],
document_definition)
except errors.HTTPFailure as e:
self.assertEqual(e.status_code, 429)
self.assertGreaterEqual(
client.last_response_headers[
http_constants.HttpHeaders.ThrottleRetryWaitTimeInMs],
connection_policy.RetryOptions.MaxWaitTimeInSeconds * 1000)
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
def test_globaldb_endpoint_discovery(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = False
read_location_client = document_client.DocumentClient(Test_globaldb_tests.read_location_host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
document_definition = { 'id': 'doc',
'name': 'sample document',
'key': 'value'}
# Create Document will fail for the read location client since it has EnableEndpointDiscovery set to false, and hence the request will directly go to
# the endpoint that was used to create the client instance(which happens to be a read endpoint)
self.__AssertHTTPFailureWithStatus(
403,
3,
read_location_client.CreateDocument,
self.test_coll['_self'],
document_definition)
# Query databases will pass for the read location client as it's a GET operation
list(read_location_client.QueryDatabases({
'query': 'SELECT * FROM root r WHERE r.id=@id',
'parameters': [
{ 'name':'@id', 'value': self.test_db['id'] }
]
}))
connection_policy.EnableEndpointDiscovery = True
read_location_client = document_client.DocumentClient(Test_globaldb_tests.read_location_host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
# CreateDocument call will go to the WriteEndpoint as EnableEndpointDiscovery is set to True and client will resolve the right endpoint based on the operation
created_document = read_location_client.CreateDocument(self.test_coll['_self'], document_definition)
self.assertEqual(created_document['id'], document_definition['id'])
def test_default_retry_policy_for_read(self):
connection_policy = documents.ConnectionPolicy()
client = document_client.DocumentClient(
Test_retry_policy_tests.host,
{'masterKey': Test_retry_policy_tests.masterKey},
connection_policy)
document_definition = {
'id': 'doc',
'name': 'sample document',
'key': 'value'
}
created_document = client.CreateDocument(
self.created_collection['_self'], document_definition)
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunctionConnectionReset
doc = client.ReadDocument(created_document['_self'], {})
self.assertEqual(doc['id'], 'doc')
self.assertEqual(self.counter, 3)
self.counter = 0
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
client.DeleteDocument(doc['_self'])
def test_resource_throttle_retry_policy_fixed_retry_after(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.RetryOptions = retry_options.RetryOptions(5, 2000)
client = document_client.DocumentClient(
Test_retry_policy_tests.host,
{'masterKey': Test_retry_policy_tests.masterKey},
connection_policy)
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunction
document_definition = {
'id': 'doc',
'name': 'sample document',
'key': 'value'
}
try:
client.CreateDocument(self.created_collection['_self'],
document_definition)
except errors.HTTPFailure as e:
self.assertEqual(e.status_code, StatusCodes.TOO_MANY_REQUESTS)
self.assertEqual(
connection_policy.RetryOptions.MaxRetryAttemptCount,
client.last_response_headers[HttpHeaders.ThrottleRetryCount])
self.assertGreaterEqual(
client.last_response_headers[
HttpHeaders.ThrottleRetryWaitTimeInMs],
connection_policy.RetryOptions.MaxRetryAttemptCount *
connection_policy.RetryOptions.FixedRetryIntervalInMilliseconds
)
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
def test_globaldb_preferred_locations(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = True
client = document_client.DocumentClient(
Test_globaldb_tests.host,
{'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
document_definition = {
'id': 'doc',
'name': 'sample document',
'key': 'value'
}
created_document = client.CreateDocument(self.test_coll['_self'],
document_definition)
self.assertEqual(created_document['id'], document_definition['id'])
# Delay to get these resources replicated to read location due to Eventual consistency
time.sleep(5)
client.ReadDocument(created_document['_self'])
content_location = str(client.last_response_headers[
http_constants.HttpHeaders.ContentLocation])
content_location_url = urlparse(content_location)
write_location_url = urlparse(Test_globaldb_tests.write_location_host)
# If no preferred locations is set, we return the write endpoint as ReadEndpoint for better latency performance
self.assertEqual(str(content_location_url.hostname),
str(write_location_url.hostname))
self.assertEqual(client.ReadEndpoint,
Test_globaldb_tests.write_location_host)
connection_policy.PreferredLocations = [
Test_globaldb_tests.read_location2
]
client = document_client.DocumentClient(
Test_globaldb_tests.host,
{'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
document_definition['id'] = 'doc2'
created_document = client.CreateDocument(self.test_coll['_self'],
document_definition)
# Delay to get these resources replicated to read location due to Eventual consistency
time.sleep(5)
client.ReadDocument(created_document['_self'])
content_location = str(client.last_response_headers[
http_constants.HttpHeaders.ContentLocation])
content_location_url = urlparse(content_location)
read_location2_url = urlparse(Test_globaldb_tests.read_location2_host)
# Test that the preferred location is set as ReadEndpoint instead of default write endpoint when no preference is set
self.assertEqual(str(content_location_url.hostname),
str(read_location2_url.hostname))
self.assertEqual(client.ReadEndpoint,
Test_globaldb_tests.read_location2_host)
def test_resource_throttle_retry_policy_query(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.RetryOptions = retry_options.RetryOptions(5)
client = document_client.DocumentClient(Test_retry_policy_tests.host, {'masterKey': Test_retry_policy_tests.masterKey}, connection_policy)
document_definition = { 'id': 'doc',
'name': 'sample document',
'key': 'value'}
client.CreateDocument(self.created_collection['_self'], document_definition)
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunction
try:
list(client.QueryDocuments(
self.created_collection['_self'],
{
'query': 'SELECT * FROM root r WHERE r.id=@id',
'parameters': [
{ 'name':'@id', 'value':document_definition['id'] }
]
}))
except errors.HTTPFailure as e:
self.assertEqual(e.status_code, 429)
self.assertEqual(connection_policy.RetryOptions.MaxRetryAttemptCount, client.last_response_headers[http_constants.HttpHeaders.ThrottleRetryCount])
self.assertGreaterEqual(client.last_response_headers[http_constants.HttpHeaders.ThrottleRetryWaitTimeInMs], connection_policy.RetryOptions.MaxRetryAttemptCount * self.retry_after_in_milliseconds)
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
def ObtainClient():
connection_policy = documents.ConnectionPolicy()
connection_policy.SSLConfiguration = documents.SSLConfiguration()
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
connection_policy.SSLConfiguration.SSLCaCerts = False
return document_client.DocumentClient(HOST, {'masterKey': MASTER_KEY},
connection_policy)
def test_globaldb_endpoint_discovery_retry_policy(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = True
write_location_client = document_client.DocumentClient(Test_globaldb_mock_tests.write_location_host, {'masterKey': Test_globaldb_mock_tests.masterKey}, connection_policy)
self.assertEqual(write_location_client._global_endpoint_manager.WriteEndpoint, Test_globaldb_mock_tests.write_location_host)
self.MockCreateDatabase(write_location_client, { 'id': 'mock database' })
self.assertEqual(write_location_client._global_endpoint_manager.WriteEndpoint, Test_globaldb_mock_tests.read_location_host)
def main():
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = True
#Reading json document - Event Data
with open('eventDocument.json') as eventDoc:
eventRecord = json.load(eventDoc)
#print(eventRecord)
LogEvents(eventRecord)
def ObtainClient():
connection_policy = documents.ConnectionPolicy()
connection_policy.SSLConfiguration = documents.SSLConfiguration()
# Try to setup the cacert.pem
# connection_policy.SSLConfiguration.SSLCaCerts = CaCertPath
# Else, disable verification
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
connection_policy.SSLConfiguration.SSLCaCerts = False
return document_client.DocumentClient(HOST, {'masterKey': MASTER_KEY},
connection_policy)
def test_globaldb_read_write_endpoints(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = False
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
document_definition = { 'id': 'doc',
'name': 'sample document',
'key': 'value'}
# When EnableEndpointDiscovery is False, WriteEndpoint is set to the endpoint passed while creating the client instance
created_document = client.CreateDocument(self.test_coll['_self'], document_definition)
self.assertEqual(client.WriteEndpoint, Test_globaldb_tests.host)
# Delay to get these resources replicated to read location due to Eventual consistency
time.sleep(5)
client.ReadDocument(created_document['_self'])
content_location = str(client.last_response_headers[http_constants.HttpHeaders.ContentLocation])
content_location_url = urlparse(content_location)
host_url = urlparse(Test_globaldb_tests.host)
# When EnableEndpointDiscovery is False, ReadEndpoint is set to the endpoint passed while creating the client instance
self.assertEqual(str(content_location_url.hostname), str(host_url.hostname))
self.assertEqual(client.ReadEndpoint, Test_globaldb_tests.host)
connection_policy.EnableEndpointDiscovery = True
document_definition['id'] = 'doc2'
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
# When EnableEndpointDiscovery is True, WriteEndpoint is set to the write endpoint
created_document = client.CreateDocument(self.test_coll['_self'], document_definition)
self.assertEqual(client.WriteEndpoint, Test_globaldb_tests.write_location_host)
# Delay to get these resources replicated to read location due to Eventual consistency
time.sleep(5)
client.ReadDocument(created_document['_self'])
content_location = str(client.last_response_headers[http_constants.HttpHeaders.ContentLocation])
content_location_url = urlparse(content_location)
write_location_url = urlparse(Test_globaldb_tests.write_location_host)
# If no preferred locations is set, we return the write endpoint as ReadEndpoint for better latency performance
self.assertEqual(str(content_location_url.hostname), str(write_location_url.hostname))
self.assertEqual(client.ReadEndpoint, Test_globaldb_tests.write_location_host)
def test_globaldb_database_account_unavailable(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = True
client = document_client.DocumentClient(Test_globaldb_mock_tests.host, {'masterKey': Test_globaldb_mock_tests.masterKey}, connection_policy)
self.assertEqual(client._global_endpoint_manager.WriteEndpoint, Test_globaldb_mock_tests.write_location_host)
self.assertEqual(client._global_endpoint_manager.ReadEndpoint, Test_globaldb_mock_tests.write_location_host)
global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub
client._global_endpoint_manager.DatabaseAccountAvailable = False
client._global_endpoint_manager.RefreshEndpointList()
self.assertEqual(client._global_endpoint_manager.WriteEndpoint, Test_globaldb_mock_tests.host)
self.assertEqual(client._global_endpoint_manager.ReadEndpoint, Test_globaldb_mock_tests.host)
class _test_config(object):
host = os.getenv('ACCOUNT_HOST', 'https://localhost:443')
masterKey = os.getenv('ACCOUNT_KEY', 'C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==')
connectionPolicy = documents.ConnectionPolicy()
connectionPolicy.DisableSSLVerification = True
global_host = '[YOUR_GLOBAL_ENDPOINT_HERE]'
write_location_host = '[YOUR_WRITE_ENDPOINT_HERE]'
read_location_host = '[YOUR_READ_ENDPOINT_HERE]'
read_location2_host = '[YOUR_READ_ENDPOINT2_HERE]'
global_masterKey = '[YOUR_KEY_HERE]'
write_location = '[YOUR_WRITE_LOCATION_HERE]'
read_location = '[YOUR_READ_LOCATION_HERE]'
read_location2 = '[YOUR_READ_LOCATION2_HERE]'
def test_default_retry_policy_for_query(self):
connection_policy = documents.ConnectionPolicy()
client = document_client.DocumentClient(
Test_retry_policy_tests.host,
{'masterKey': Test_retry_policy_tests.masterKey},
connection_policy)
document_definition_1 = {
'id': 'doc1',
'name': 'sample document',
'key': 'value'
}
document_definition_2 = {
'id': 'doc2',
'name': 'sample document',
'key': 'value'
}
client.CreateDocument(self.created_collection['_self'],
document_definition_1)
client.CreateDocument(self.created_collection['_self'],
document_definition_2)
self.OriginalExecuteFunction = retry_utility._ExecuteFunction
retry_utility._ExecuteFunction = self._MockExecuteFunctionConnectionReset
docs = client.QueryDocuments(self.created_collection['_self'],
"Select * from c", {'maxItemCount': 1})
result_docs = list(docs)
self.assertEqual(result_docs[0]['id'], 'doc1')
self.assertEqual(result_docs[1]['id'], 'doc2')
self.assertEqual(self.counter, 12)
self.counter = 0
retry_utility._ExecuteFunction = self.OriginalExecuteFunction
client.DeleteDocument(result_docs[0]['_self'])
client.DeleteDocument(result_docs[1]['_self'])
def test_globaldb_endpoint_assignments(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = False
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
# When EnableEndpointDiscovery is set to False, both Read and Write Endpoints point to endpoint passed while creating the client instance
self.assertEqual(client._global_endpoint_manager.WriteEndpoint, Test_globaldb_tests.host);
self.assertEqual(client._global_endpoint_manager.ReadEndpoint, Test_globaldb_tests.host);
connection_policy.EnableEndpointDiscovery = True
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
# If no preferred locations is set, we return the write endpoint as ReadEndpoint for better latency performance, write endpoint is set as expected
self.assertEqual(client._global_endpoint_manager.WriteEndpoint, Test_globaldb_tests.write_location_host);
self.assertEqual(client._global_endpoint_manager.ReadEndpoint, Test_globaldb_tests.write_location_host);
connection_policy.PreferredLocations = [Test_globaldb_tests.read_location2]
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
# Test that the preferred location is set as ReadEndpoint instead of default write endpoint when no preference is set
self.assertEqual(client._global_endpoint_manager.WriteEndpoint, Test_globaldb_tests.write_location_host);
self.assertEqual(client._global_endpoint_manager.ReadEndpoint, Test_globaldb_tests.read_location2_host);
import pydocumentdb
from pydocumentdb import document_client
from pydocumentdb import documents
# These values can be found in the Azure Portal under the Cosmos DB Connection Strings
masterKey = 'Insert your masterKey here'
host = 'Insert your host here'
connectionPolicy = documents.ConnectionPolicy()
connectionPolicy.EnableEndpointDiscovery
connectionPolicy.PreferredLocations = ["Central US", "East US 2", "Southeast Asia", "Western Europe","Canada Central"]
client = document_client.DocumentClient(host, {'masterKey': masterKey}, connectionPolicy)
databaseId = 'Insert Database ID'
collectionId = 'Insert Collection ID'
dbLink = 'dbs/' + databaseId
collLink = dbLink + '/colls/' + collectionId
querystr = "SELECT * from c"
query = client.QueryDocuments(collLink, querystr, options= { 'enableCrossPartitionQuery': True }, partition_key=None)
cosmoselements = list(query)
cosmos_data_df = spark.createDataFrame(cosmoselements)
def test_globaldb_update_locations_cache(self):
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey})
writable_locations = [{'name' : Test_globaldb_tests.write_location, 'databaseAccountEndpoint' : Test_globaldb_tests.write_location_host}]
readable_locations = [{'name' : Test_globaldb_tests.read_location, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location_host}, {'name' : Test_globaldb_tests.read_location2, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location2_host}]
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# If no preferred locations is set, we return the write endpoint as ReadEndpoint for better latency performance, write endpoint is set as expected
self.assertEqual(write_endpoint, Test_globaldb_tests.write_location_host)
self.assertEqual(read_endpoint, Test_globaldb_tests.write_location_host)
writable_locations = []
readable_locations = []
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# If writable_locations and readable_locations are empty, both Read and Write Endpoints point to endpoint passed while creating the client instance
self.assertEqual(write_endpoint, Test_globaldb_tests.host)
self.assertEqual(read_endpoint, Test_globaldb_tests.host)
writable_locations = [{'name' : Test_globaldb_tests.write_location, 'databaseAccountEndpoint' : Test_globaldb_tests.write_location_host}]
readable_locations = []
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# If there are no readable_locations, we use the write endpoint as ReadEndpoint
self.assertEqual(write_endpoint, Test_globaldb_tests.write_location_host)
self.assertEqual(read_endpoint, Test_globaldb_tests.write_location_host)
writable_locations = []
readable_locations = [{'name' : Test_globaldb_tests.read_location, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location_host}]
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# If there are no writable_locations, both Read and Write Endpoints point to endpoint passed while creating the client instance
self.assertEqual(write_endpoint, Test_globaldb_tests.host)
self.assertEqual(read_endpoint, Test_globaldb_tests.host)
writable_locations = [{'name' : Test_globaldb_tests.write_location, 'databaseAccountEndpoint' : Test_globaldb_tests.write_location_host}]
readable_locations = [{'name' : Test_globaldb_tests.read_location, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location_host}, {'name' : Test_globaldb_tests.read_location2, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location2_host}]
connection_policy = documents.ConnectionPolicy()
connection_policy.PreferredLocations = [Test_globaldb_tests.read_location2]
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# Test that the preferred location is set as ReadEndpoint instead of default write endpoint when no preference is set
self.assertEqual(write_endpoint, Test_globaldb_tests.write_location_host)
self.assertEqual(read_endpoint, Test_globaldb_tests.read_location2_host)
writable_locations = [{'name' : Test_globaldb_tests.write_location, 'databaseAccountEndpoint' : Test_globaldb_tests.write_location_host}, {'name' : Test_globaldb_tests.read_location2, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location2_host}]
readable_locations = [{'name' : Test_globaldb_tests.read_location, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location_host}]
connection_policy = documents.ConnectionPolicy()
connection_policy.PreferredLocations = [Test_globaldb_tests.read_location2]
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# Test that the preferred location is chosen from the WriteLocations if it's not present in the ReadLocations
self.assertEqual(write_endpoint, Test_globaldb_tests.write_location_host)
self.assertEqual(read_endpoint, Test_globaldb_tests.read_location2_host)
writable_locations = [{'name' : Test_globaldb_tests.write_location, 'databaseAccountEndpoint' : Test_globaldb_tests.write_location_host}]
readable_locations = [{'name' : Test_globaldb_tests.read_location, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location_host}, {'name' : Test_globaldb_tests.read_location2, 'databaseAccountEndpoint' : Test_globaldb_tests.read_location2_host}]
connection_policy.EnableEndpointDiscovery = False
client = document_client.DocumentClient(Test_globaldb_tests.host, {'masterKey': Test_globaldb_tests.masterKey}, connection_policy)
write_endpoint, read_endpoint = client._global_endpoint_manager.UpdateLocationsCache(writable_locations, readable_locations)
# If EnableEndpointDiscovery is False, both Read and Write Endpoints point to endpoint passed while creating the client instance
self.assertEqual(write_endpoint, Test_globaldb_tests.host)
self.assertEqual(read_endpoint, Test_globaldb_tests.host)
