def get_cosmos_client():
"""Get the client for all CosmosDB queries.
Returns:
CosmosClient
"""
return cosmos_client.CosmosClient(
url_connection=os.getenv('CosmosDBHost'),
auth={'masterKey': os.getenv('CosmosDBMasterKey')})
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 = cosmos_client.CosmosClient(cls.host, cls.masterKey, connection_policy=cls.connectionPolicy)
cls.created_collection = test_config._test_config.create_multi_partition_collection_with_custom_pk_if_not_exist(cls.client)
def cleanUpTestDatabase(cls):
global client
client = cosmos_client.CosmosClient(cls.host,
{'masterKey': cls.masterKey}, cls.connectionPolicy)
query_iterable = client.QueryDatabases('SELECT * FROM root r WHERE r.id=\'' + cls.testDbName + '\'')
it = iter(query_iterable)
test_db = next(it, None)
if test_db is not None:
client.DeleteDatabase(test_db['_self'])
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 cosmos_client.CosmosClient(HOST, {'masterKey': MASTER_KEY}, connection_policy)
def SetupCosmosDB(container=None):
#Azure Function Setup
config = CommonFramework.RetrieveConfigOptions("cosmosdb")
cosmosclient = cosmos_client.CosmosClient(url_connection=config['ENDPOINT'],auth={'masterKey': config['PRIMARYKEY']})
db = next((data for data in cosmosclient.ReadDatabases() if data['id'] == config['DATABASE']))
if container is None:
coll = next((coll for coll in cosmosclient.ReadContainers(db['_self']) if coll['id'] == config['CONTAINER']))
else:
coll = next((coll for coll in cosmosclient.ReadContainers(db['_self']) if coll['id'] == str(container)))
return cosmosclient, coll
def test_globaldb_endpoint_discovery(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = False
read_location_client = cosmos_client.CosmosClient(
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(StatusCodes.FORBIDDEN,
SubStatusCodes.WRITE_FORBIDDEN,
read_location_client.CreateItem,
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 = cosmos_client.CosmosClient(
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.CreateItem(
self.test_coll['_self'], document_definition)
self.assertEqual(created_document['id'], document_definition['id'])
def getdata():
config = {
'ENDPOINT': 'https://estruscosmosdb.documents.azure.com:443/',
'PRIMARYKEY':
'R2M6nKvkUKf4vYohKEr2hph74zDMVD9i4mQT2wTyKSE8kYdZ3NAHNlth1Fz0BZ1vaYnCgamLRnYhOZ5B33A8nQ==',
'DATABASE': 'estruscosmosdb',
'CONTAINER': 'EstrusStatus'
}
# Initialize the Cosmos client
client = cosmos_client.CosmosClient(
url_connection=config['ENDPOINT'],
auth={'masterKey': config['PRIMARYKEY']})
# Create a database
db = list(
client.QueryDatabases("select * from c where c.id='estruscosmosdb'"))
# Create container options
options = {'offerThroughput': 400}
container_definition = {'id': config['CONTAINER']}
# Create a container
#container = client.CreateContainer(db['_self'], container_definition, options)
container = list(
client.QueryContainers(db[0]['_self'],
"select * from c where c.id='EstrusStatus'",
options))
# Query these items in SQL
query = {'query': "SELECT * FROM c"}
options = {}
options['enableCrossPartitionQuery'] = True
options['maxItemCount'] = 2
result_iterable = client.QueryItems(container[0]['_self'], query, options)
data = []
for item in iter(result_iterable):
row = {}
row['cowId'] = item['cowId']
row['probability'] = item['probability']
row['label'] = item['estrusLabel']
row['hour_temp'] = item['hour_stomach_temp_celcius']
row['activity'] = item['hour_animal_activity']
row['hour_low_pass_over_activity'] = item[
'hour_low_pass_over_activity']
row['hour_temp_without_drink_cycles'] = item[
'hour_temp_without_drink_cycles']
data.append(row)
data.sort(key=lambda a: -a['probability'])
return data
def get_conn(self):
"""
Return a cosmos db client.
"""
if self.cosmos_client is not None:
return self.cosmos_client
# Initialize the Python Azure Cosmos DB client
self.cosmos_client = cosmos_client.CosmosClient(self.endpoint_uri, {'masterKey': self.master_key})
return self.cosmos_client
def __init__(self):
self.config = {
'ENDPOINT': 'https://nasa-hackathon-db.documents.azure.com:443/',
'PRIMARYKEY':
'XwAxG7ma3R2KUAWHAQJXVFD95FDSAnVANaeiOlMeQwOgyCl4v0i2jnW8uWkQ2n1ssHQai7yh4ScdKJgPlgD02Q==',
'DATABASE': 'NewsDatabase',
'CONTAINER': 'NewsContainer'
}
self.client = cosmos_client.CosmosClient(
url_connection=self.config['ENDPOINT'],
auth={'masterKey': self.config['PRIMARYKEY']})
def setUp(self):
RoutingMapEndToEndTests.cleanUpTestDatabase()
self.client = cosmos_client.CosmosClient(RoutingMapEndToEndTests.host, {'masterKey': RoutingMapEndToEndTests.masterKey}, RoutingMapEndToEndTests.connectionPolicy)
self.created_db = self.client.CreateDatabase({ 'id': 'sample database' })
self.created_collection = self.create_collection(self.client, self.created_db)
self.collection_link = self.GetDocumentCollectionLink(self.created_db, self.created_collection)
# sanity check:
partition_key_ranges = list(self.client._ReadPartitionKeyRanges(self.collection_link))
self.assertGreaterEqual(len(partition_key_ranges), 5)
def __init__(self):
# Initialize the Cosmos client
self.client = cosmos_client.CosmosClient(url_connection=config['ENDPOINT'],
auth={'masterKey': config['PRIMARYKEY']})
# Read a database
self.database_link = 'dbs/' + config['DATABASE']
# self.database_link = self.client.ReadDatabase(self.database_link)['_self']
# Read a container
self.container_link = self.database_link + '/colls/{0}'.format(config['CONTAINER'])
def test_failure_with_wrong_proxy(self):
connection_policy.ProxyConfiguration.Port = self.serverPort + 1
try:
# client does a getDatabaseAccount on initialization, which fails
client = cosmos_client.CosmosClient(self.host,
{'masterKey': self.masterKey},
connection_policy)
self.fail("Client instantiation is not expected")
except Exception as e:
self.assertTrue(type(e) is ProxyError,
msg="Error is not a ProxyError")
def __init__(self, cosmosdb_uri, cosmos_key, throughput, db_id,
collection_id):
master_key = "masterKey"
self.cosmos_db_client = cosmos_client.CosmosClient(
url_connection=cosmosdb_uri, auth={master_key: cosmos_key})
self.db_id = db_id
self.collection_id = collection_id
self.throughput = throughput
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 = cosmos_client.CosmosClient(cls.host, {'masterKey': cls.masterKey}, "Session", cls.connectionPolicy)
cls.created_collection = test_config._test_config.create_single_partition_collection_if_not_exist(cls.client)
cls.retry_after_in_milliseconds = 1000
def __init__(self):
self.dbKey = "bH7iSiCXWS1TlQbWskc37u29oDpFjAUe5VTPon0RttLXA3wgMrko5QNYE244KHU3KOGdXtjMWBb4rcuwYDn8cQ=="
self.dbURL = "https://study-talles-cosmos.documents.azure.com:443/"
self.dbId = "study-talles"
self.client = cosmosClient.CosmosClient(self.dbURL,
{'masterKey': self.dbKey})
self.db = database = self.client.get_database_client(self.dbId)
self.personRepository = PersonRepository(self.db, "person")
def __init__(self, cosmosdb_uri, cosmos_key, db_id, collection_id, rows):
master_key = "masterKey"
self.cosmos_db_client = cosmos_client.CosmosClient(
url_connection=cosmosdb_uri, auth={master_key: cosmos_key})
self.collection_link = "dbs/" + db_id + "/colls/" + collection_id
self.db_id = db_id
self.collection_id = collection_id
self.rows = rows
def connect(self) -> cosmos_client.CosmosClient:
"""Connects to the Cosmos Database.
Returns:
----
cosmos_client.CosmosClient: A cosmos client object.
"""
client = cosmos_client.CosmosClient(
url=self.account_uri, credential={"masterKey": self.account_key})
return client
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 = cosmos_client.CosmosClient(cls.host,
{'masterKey': cls.masterKey},
cls.connectionPolicy)
cls.created_db = test_config._test_config.create_database_if_not_exist(
cls.client)
def __init__(self, opts):
self._opts = opts
self._dbname = None
self._dbproxy = None
self._ctrproxy = None
self._cname = None
self._query_metrics = True
self.reset_record_diagnostics()
print(self._opts)
url = opts['url']
key = opts['key']
self._client = cosmos_client.CosmosClient(url, {'masterKey': key})
def __init__(self, client_id, key, tenant_id, url, cosmos_key):
load_dotenv()
self.url = url
self.__comos_key = cosmos_key
self.cosmos_client = cosmos_client.CosmosClient(url, cosmos_key)
self.credentials = AzureConnections().authenticate_device_code(
CLIENT=client_id, KEY=key, TENANT_ID=tenant_id)
with change_dir(OPERATIONSPATH):
self.config = ConfigParser()
self.config.read(CONFVARIABLES)
def run_sample():
with IDisposable(
cosmos_client.CosmosClient(HOST,
{'masterKey': MASTER_KEY})) as client:
try:
# setup database for this sample
try:
db = client.create_database(id=DATABASE_ID)
except errors.HTTPFailure as e:
if e.status_code == 409:
pass
else:
raise errors.HTTPFailure(e.status_code)
# setup container for this sample
try:
container = db.create_container(id=CONTAINER_ID,
partition_key=PartitionKey(
path='/id', kind='Hash'))
print('Container with id \'{0}\' created'.format(CONTAINER_ID))
except errors.HTTPFailure as e:
if e.status_code == 409:
print('Container with id \'{0}\' was found'.format(
CONTAINER_ID))
else:
raise errors.HTTPFailure(e.status_code)
ItemManagement.CreateItems(container)
ItemManagement.ReadItem(container, 'SalesOrder1')
ItemManagement.ReadItems(container)
ItemManagement.QueryItems(container, 'SalesOrder1')
ItemManagement.ReplaceItem(container, 'SalesOrder1')
ItemManagement.UpsertItem(container, 'SalesOrder1')
ItemManagement.DeleteItem(container, 'SalesOrder1')
# cleanup database after sample
try:
client.delete_database(db)
except errors.CosmosError as e:
if e.status_code == 404:
pass
else:
raise errors.HTTPFailure(e.status_code)
except errors.HTTPFailure as e:
print('\nrun_sample has caught an error. {0}'.format(e.message))
finally:
print("\nrun_sample done")
def add_to_database(studentRecord):
try:
cosmos_cli = cosmos_client.CosmosClient(url_connection=COSMOS_ENDPOINT,
auth={"masterKey": COSMOS_KEY})
collection_link = get_container_link()
container = cosmos_cli.ReadContainer(collection_link)
query = "SELECT * FROM c WHERE c.userid = '{0}'".format(
studentRecord['userid'])
records = list(
cosmos_cli.QueryItems(
database_or_Container_link=container["_self"],
query=query,
options={'enableCrossPartitionQuery': True}))
if len(records) == 0:
studentRecord['iteration'] = 1
cosmos_cli.CreateItem(
database_or_Container_link=container["_self"],
document=studentRecord)
else:
existingRecord = records[0]
iteration = existingRecord['iteration']
studentRecord['happiness'] = get_mean_emotion(
existingRecord['happiness'], studentRecord['happiness'],
iteration)
studentRecord['anger'] = get_mean_emotion(existingRecord['anger'],
studentRecord['anger'],
iteration)
studentRecord['sadness'] = get_mean_emotion(
existingRecord['sadness'], studentRecord['sadness'], iteration)
studentRecord['neutral'] = get_mean_emotion(
existingRecord['neutral'], studentRecord['neutral'], iteration)
studentRecord['contempt'] = get_mean_emotion(
existingRecord['contempt'], studentRecord['contempt'],
iteration)
studentRecord['disgust'] = get_mean_emotion(
existingRecord['disgust'], studentRecord['disgust'], iteration)
studentRecord['surprise'] = get_mean_emotion(
existingRecord['surprise'], studentRecord['surprise'],
iteration)
studentRecord['fear'] = get_mean_emotion(existingRecord['fear'],
studentRecord['fear'],
iteration)
studentRecord['iteration'] = iteration + 1
studentRecord['id'] = existingRecord['id']
cosmos_cli.ReplaceItem(document_link=existingRecord['_self'],
new_document=studentRecord)
except errors.HTTPFailure as err:
logging.error(err)
raise
def __init__(self, config: CosmosDbConfig):
"""Create the storage object.
:param config:
"""
super(CosmosDbStorage, self).__init__()
self.config = config
self.client = cosmos_client.CosmosClient(
self.config.endpoint, {'masterKey': self.config.masterkey})
# these are set by the functions that check
# the presence of the db and container or creates them
self.db = None
self.container = None
def test_globaldb_endpoint_assignments(self):
connection_policy = documents.ConnectionPolicy()
connection_policy.EnableEndpointDiscovery = False
client = cosmos_client.CosmosClient(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 = cosmos_client.CosmosClient(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 = cosmos_client.CosmosClient(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)
def create_spy_client(self, use_multiple_write_locations,
enable_endpoint_discovery,
is_preferred_locations_list_empty):
preferred_locations = ["location1", "location2", "location3"]
connectionPolicy = documents.ConnectionPolicy()
connectionPolicy.DisableSSLVerification = True
connectionPolicy.PreferredLocations = [] if is_preferred_locations_list_empty else preferred_locations
connectionPolicy.EnableEndpointDiscovery = enable_endpoint_discovery
connectionPolicy.UseMultipleWriteLocations = use_multiple_write_locations
client = cosmos_client.CosmosClient(self.DEFAULT_ENDPOINT,
{'masterKey': "SomeKeyValue"},
connectionPolicy)
return client
def ensure_connected(self):
global CosmosClient
if not CosmosClient:
import os
host = os.environ["COSMOS_ENDPOINT"]
key = os.environ["COSMOS_KEY"]
if (not host) or (not key):
logging.error("Cosmos endpoint credentials is not set.")
print("cosmos endpoint credentials is not set")
raise Exception("cosmos endpoint credentials are not set")
CosmosClient = cosmos_client.CosmosClient(host, {'masterKey': key})
def main():
print("Running receiver")
with d.IDisposable(
cosmos_client.CosmosClient(d.HOST,
{'masterKey': d.MASTER_KEY})) as client:
try:
''''''
# setup database for this sample
try:
client.CreateDatabase({"id": d.DATABASE_ID})
except errors.HTTPFailure as e:
if e.status_code == 409:
pass
else:
raise errors.HTTPFailure(e.status_code)
# setup collection for this sample
try:
client.CreateContainer(d.database_link,
{"id": d.COLLECTION_ID})
print('Collection with id \'{0}\' created'.format(
d.COLLECTION_ID))
except errors.HTTPFailure as e:
if e.status_code == 409:
print('Collection with id \'{0}\' was found'.format(
d.COLLECTION_ID))
else:
raise errors.HTTPFailure(e.status_code)
# Create keys and queue
# d.DocumentManagement.CreateDocuments(client)
data = d.DocumentManagement.ReadDocument(
client, '0376af03-9620-4980-a905-dbfa6b189495')
publicKey = keys.Key().read_key(data['publicKey'])
privateKey = keys.Key().read_key_from_file('wxptlpygzv')
queue = data['queue']
print(publicKey)
print(privateKey)
print(queue)
except errors.HTTPFailure as e:
print('\nrun_sample has caught an error. {0}'.format(
e._http_error_message))
finally:
print("\nrun_sample done")
def __init__(self, config):
self._config = config
# Initialize the Cosmos client
self._client = cosmos_client.CosmosClient(
url_connection=config['ENDPOINT'],
auth={'masterKey': config['PRIMARYKEY']})
#self._options = {
# 'offerThroughput': 400
#}
#self._container_definition = {
# 'id': self._config['CONTAINER']
#}
self._collection_link = "/dbs/{}/colls/{}".format(
self._config['DATABASE'], self._config['CONTAINER'])
def intialise_client():
"""Initialise Client Routine
======================================
Initialise and return database connection client.
Args:
None.
Returns:
client (CosmosClient) - Database connection client.
"""
client = cosmos_client.CosmosClient(HOST, {'masterKey': MASTER_KEY})
return client
def run_sample():
with IDisposable(
cosmos_client.CosmosClient(HOST,
{'masterKey': MASTER_KEY})) as client:
try:
# setup database for this sample
try:
db = client.create_database(id=DATABASE_ID)
except errors.HTTPFailure as e:
if e.status_code == 409:
pass
else:
raise errors.HTTPFailure(e.status_code)
# query for a container
ContainerManagement.find_container(db, CONTAINER_ID)
# create a container
ContainerManagement.create_Container(db, CONTAINER_ID)
# get & change Offer Throughput of container
ContainerManagement.manage_offer_throughput(db, CONTAINER_ID)
# get a container using its id
ContainerManagement.read_Container(db, CONTAINER_ID)
# list all container on an account
ContainerManagement.list_Containers(db)
# delete container by id
ContainerManagement.delete_Container(db, CONTAINER_ID)
# cleanup database after sample
try:
client.delete_database(db)
except errors.CosmosError as e:
if e.status_code == 404:
pass
else:
raise errors.HTTPFailure(e.status_code)
except errors.HTTPFailure as e:
print('\nrun_sample has caught an error. {0}'.format(e.message))
finally:
print("\nrun_sample done")
