#The MIT License (MIT) #Copyright (c) 2014 Microsoft Corporation #Permission is hereby granted, free of charge, to any person obtaining a copy #of this software and associated documentation files (the "Software"), to deal #in the Software without restriction, including without limitation the rights #to use, copy, modify, merge, publish, distribute, sublicense, and/or sell #copies of the Software, and to permit persons to whom the Software is #furnished to do so, subject to the following conditions: #The above copyright notice and this permission notice shall be included in all #copies or substantial portions of the Software. #THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR #IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, #FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER #LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, #OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE #SOFTWARE. import os import time import uuid import azure.cosmos.documents as documents import azure.cosmos.exceptions as exceptions from azure.cosmos.http_constants import StatusCodes from azure.cosmos.database import DatabaseProxy from azure.cosmos.cosmos_client import CosmosClient from azure.cosmos.partition_key import PartitionKey from azure.cosmos.partition_key import NonePartitionKeyValue try: import urllib3 urllib3.disable_warnings() except: print("no urllib3") class _test_config(object): #[SuppressMessage("Microsoft.Security", "CS002:SecretInNextLine", Justification="Cosmos DB Emulator Key")] masterKey = os.getenv('ACCOUNT_KEY', 'C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==') host = os.getenv('ACCOUNT_HOST', 'https://localhost:8081/') connection_str = os.getenv('ACCOUNT_CONNECTION_STR', 'AccountEndpoint={};AccountKey={};'.format(host, masterKey)) connectionPolicy = documents.ConnectionPolicy() connectionPolicy.DisableSSLVerification = True global_host = os.getenv('GLOBAL_ACCOUNT_HOST', host) write_location_host = os.getenv('WRITE_LOCATION_HOST', host) read_location_host = os.getenv('READ_LOCATION_HOST', host) read_location2_host = os.getenv('READ_LOCATION_HOST2', host) global_masterKey = os.getenv('GLOBAL_ACCOUNT_KEY', masterKey) write_location = os.getenv('WRITE_LOCATION', host) read_location = os.getenv('READ_LOCATION', host) read_location2 = os.getenv('READ_LOCATION2', host) THROUGHPUT_FOR_5_PARTITIONS = 30000 THROUGHPUT_FOR_1_PARTITION = 400 TEST_DATABASE_ID = os.getenv('COSMOS_TEST_DATABASE_ID', "Python SDK Test Database " + str(uuid.uuid4())) TEST_DATABASE_ID_PLAIN = "COSMOS_TEST_DATABASE" TEST_THROUGHPUT_DATABASE_ID = "Python SDK Test Throughput Database " + str(uuid.uuid4()) TEST_COLLECTION_SINGLE_PARTITION_ID = "Single Partition Test Collection" TEST_COLLECTION_MULTI_PARTITION_ID = "Multi Partition Test Collection" TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_ID = "Multi Partition Test Collection With Custom PK" TEST_COLLECTION_MULTI_PARTITION_PARTITION_KEY = "id" TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_PARTITION_KEY = "pk" TEST_DATABASE = None TEST_COLLECTION_SINGLE_PARTITION = None TEST_COLLECTION_MULTI_PARTITION = None TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK = None IS_MULTIMASTER_ENABLED = False @classmethod def create_database_if_not_exist(cls, client): # type: (CosmosClient) -> Database if cls.TEST_DATABASE is not None: return cls.TEST_DATABASE cls.try_delete_database(client) cls.TEST_DATABASE = client.create_database(cls.TEST_DATABASE_ID) cls.IS_MULTIMASTER_ENABLED = client.get_database_account()._EnableMultipleWritableLocations return cls.TEST_DATABASE @classmethod def try_delete_database(cls, client): # type: (CosmosClient) -> None try: client.delete_database(cls.TEST_DATABASE_ID) except exceptions.CosmosHttpResponseError as e: if e.status_code != StatusCodes.NOT_FOUND: raise e @classmethod def create_single_partition_collection_if_not_exist(cls, client): # type: (CosmosClient) -> Container if cls.TEST_COLLECTION_SINGLE_PARTITION is None: cls.TEST_COLLECTION_SINGLE_PARTITION = cls.create_collection_with_required_throughput(client, cls.THROUGHPUT_FOR_1_PARTITION, False) cls.remove_all_documents(cls.TEST_COLLECTION_SINGLE_PARTITION, False) return cls.TEST_COLLECTION_SINGLE_PARTITION @classmethod def create_multi_partition_collection_if_not_exist(cls, client): # type: (CosmosClient) -> Container if cls.TEST_COLLECTION_MULTI_PARTITION is None: cls.TEST_COLLECTION_MULTI_PARTITION = cls.create_collection_with_required_throughput(client, cls.THROUGHPUT_FOR_5_PARTITIONS, False) cls.remove_all_documents(cls.TEST_COLLECTION_MULTI_PARTITION, False) return cls.TEST_COLLECTION_MULTI_PARTITION @classmethod def create_multi_partition_collection_with_custom_pk_if_not_exist(cls, client): # type: (CosmosClient) -> Container if cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK is None: cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK = cls.create_collection_with_required_throughput(client, cls.THROUGHPUT_FOR_5_PARTITIONS, True) cls.remove_all_documents(cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK, True) return cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK @classmethod def create_collection_with_required_throughput(cls, client, throughput, use_custom_partition_key): # type: (CosmosClient, int, boolean) -> Container database = cls.create_database_if_not_exist(client) if throughput == cls.THROUGHPUT_FOR_1_PARTITION: collection_id = cls.TEST_COLLECTION_SINGLE_PARTITION_ID partition_key = cls.TEST_COLLECTION_MULTI_PARTITION_PARTITION_KEY else: if use_custom_partition_key: collection_id = cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_ID partition_key = cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_PARTITION_KEY else: collection_id = cls.TEST_COLLECTION_MULTI_PARTITION_ID partition_key = cls.TEST_COLLECTION_MULTI_PARTITION_PARTITION_KEY document_collection = database.create_container( id=collection_id, partition_key=PartitionKey(path='/' + partition_key, kind='Hash'), offer_throughput=throughput) return document_collection @classmethod def remove_all_documents(cls, document_collection, use_custom_partition_key): # type: (Container, boolean) -> None while True: query_iterable = document_collection.query_items(query="Select * from c", enable_cross_partition_query=True) read_documents = list(query_iterable) try: for document in read_documents: partition_key = 'dummy_pk' if not use_custom_partition_key: partition_key = document[cls.TEST_COLLECTION_MULTI_PARTITION_PARTITION_KEY] else: if cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_PARTITION_KEY in document: partition_key = document[cls.TEST_COLLECTION_MULTI_PARTITION_WITH_CUSTOM_PK_PARTITION_KEY] else: partition_key = NonePartitionKeyValue document_collection.delete_item(item=document, partition_key=partition_key) if cls.IS_MULTIMASTER_ENABLED: # sleep to ensure deletes are propagated for multimaster enabled accounts time.sleep(2) break except exceptions.CosmosHttpResponseError as e: print("Error occurred while deleting documents:" + str(e) + " \nRetrying...") class FakeResponse: def __init__(self, headers): self.headers = headers self.reason = "foo" self.status_code = "bar"