def create_state_store_status() -> (StatusEnum, str):
status = StatusEnum.ok
message = ""
debug = True if config.DEBUG == "true" else False
try:
primary_master_key = get_store_key()
cosmos_client = CosmosClient(config.STATE_STORE_ENDPOINT,
primary_master_key,
connection_verify=debug)
list(cosmos_client.list_databases())
except exceptions.ServiceRequestError:
status = StatusEnum.not_ok
message = strings.STATE_STORE_ENDPOINT_NOT_RESPONDING
except: # noqa: E722 flake8 - no bare excepts
status = StatusEnum.not_ok
message = strings.UNSPECIFIED_ERROR
return status, message
class AzureHelper:
def __init__(self):
pass
def generate_mgmt_credentials(self):
logger = logging.getLogger(__name__)
logger.info(
'LOGGER: generating Azure Management plane credentials via Credential Wrapper for azure identity'
)
try:
self._mgmt_credentials = CredentialWrapper()
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
def generate_data_credentials(self):
logger = logging.getLogger(__name__)
logger.info(
'LOGGER: generating Azure Data plane credentials via Default Credentials for azure identity'
)
try:
self._data_credentials = DefaultAzureCredential()
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
def get_storage_mgmt_client(self, subscription: str):
logger = logging.getLogger(__name__)
logger.info(
'LOGGER: acquiring storage resource management client connection')
if subscription:
self._subscription = subscription
try:
self._storage_mgmt_client = StorageManagementClient(
self._mgmt_credentials, self._subscription)
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._storage_mgmt_client
def get_resource_mgmt_client(self, subscription: str):
logger = logging.getLogger(__name__)
logger.info(
'LOGGER: acquiring storage resource management client connection')
if subscription:
self._subscription = subscription
try:
self._resource_mgmt_client = ResourceManagementClient(
self._mgmt_credentials, self._subscription)
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._resource_mgmt_client
def get_cosmosdb_mgmt_client(self, subscription: str):
logger = logging.getLogger(__name__)
logger.info(
'LOGGER: acquiring storage resource management client connection')
if subscription:
self._subscription = subscription
try:
self._cosmosdb_mgmt_client = CosmosDBManagementClient(
self._mgmt_credentials, self._subscription)
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._cosmosdb_mgmt_client
def get_storage_account_keys(self,
resource_group: str,
storage_account: str,
storage_mgmt_client: object = None):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring storage account keys')
if storage_mgmt_client:
self._storage_mgmt_client = storage_mgmt_client
try:
self._storage_account_keys = self._storage_mgmt_client.storage_accounts.list_keys(
resource_group, storage_account).keys
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._storage_account_keys
def get_cosmosdb_account_keys(self,
resource_group_name: str,
cosmosdb_account_name: str,
cosmosdb_mgmt_client: object = None):
logger = logging.getLogger(__name__)
logger.info('LOGGER: Acquiring Cosmos account keys')
if not self._cosmosdb_mgmt_client:
self._cosmosdb_mgmt_client = cosmosdb_mgmt_client
try:
self._cosmosdb_account_keys = cosmosdb_mgmt_client.database_accounts.list_keys(
resource_group_name, cosmosdb_account_name)
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return self._cosmosdb_account_keys
def generate_storage_conn_string(self, storage_key: str,
storage_account: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring storage account keys')
try:
self._storage_conn_string = f'DefaultEndpointsProtocol=https;AccountName={storage_account};AccountKey={storage_key};EndpointSuffix=core.windows.net'
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._storage_conn_string
def get_storage_client(self, storage_conn_string: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring storage resource client connection')
if storage_conn_string:
self.storage_conn_string = storage_conn_string
try:
self._storage_client = BlobServiceClient.from_connection_string(
self.storage_conn_string, self._data_credentials)
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._storage_client
def get_cosmos_client(self, cosmos_url: str, cosmos_key: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring Cosmos client connection')
try:
self._cosmos_client = CosmosClient(cosmos_url, cosmos_key)
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return self._cosmos_client
def get_container_client(self,
container_name: str,
storage_client: object = None):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring container client connection')
if storage_client:
self._storage_client = storage_client
try:
self._container_client = self._storage_client.get_container_client(
container_name)
except Exception as e:
logger.error(f'LOGGER: Exception occurred {e}')
raise e
return self._container_client
def get_api_versions(self,
resource_id: str,
resource_mgmt_client: object = None):
logger = logging.getLogger(__name__)
logger.info('LOGGER: acquiring API version')
if resource_mgmt_client:
self._resource_mgmt_client = resource_mgmt_client
try:
namespace = self.get_provider_namespace
service = self.get_service_name
provider = self._resource_mgmt_client.providers.get(namespace)
resource_types = next((item for item in provider.resource_types
if item.resource_type == service))
api = resource_types.api_versions
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return api
def get_resource_group_name(self, resource_id: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: Acquiring resource group name')
try:
id_list = resource_id.lower().split('/')
resource_group_name_idx = id_list.index('resourcegroups') + 1
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return resource_id.split('/')[resource_group_name_idx]
def get_provider_namespace(self, resource_id: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: Acquiring resource/account name')
try:
id_list = resource_id.lower().split('/')
provider_namespace_idx = id_list.index('providers') + 1
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return resource_id.split('/')[provider_namespace_idx]
def get_service_name(self, resource_id: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: Acquiring resource/account name')
try:
id_list = resource_id.lower().split('/')
service_name_idx = id_list.index('providers') + 2
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return resource_id.split('/')[service_name_idx]
def get_resource_name(self, resource_id: str):
logger = logging.getLogger(__name__)
logger.info('LOGGER: Acquiring resource/account name')
try:
id_list = resource_id.lower().split('/')
resource_name_idx = id_list.index('providers') + 3
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
raise e
return resource_id.split('/')[resource_name_idx]
def get_cosmosdb_throughputs(self, cosmos_client: object = None):
logger = logging.getLogger(__name__)
if cosmos_client:
self._cosmos_client = cosmos_client
existing_dbs = [
item['id'] for item in self._cosmos_client.list_databases()
]
db_dict = {}
for db in existing_dbs:
try:
logger.info(
'LOGGER: Acquiring Cosmos database-level throughput')
db_client = cosmos_client.get_database_client(db)
try:
db_throughput = db_client.read_offer(
).properties['content']['offerThroughput']
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
logger.info(
'LOGGER: No database throughput option detected. Setting at none'
)
db_throughput = 'N/A'
existing_containers = [
item['id'] for item in db_client.list_containers()
]
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
finally:
container_dict = {}
if existing_containers:
for container in existing_containers:
try:
logger.info(
'LOGGER: Acquiring Cosmos container-level throughput'
)
container_throughput = db_client.get_container_client(
container).read_offer(
).properties['content']['offerThroughput']
container_dict[container] = container_dict.get(
container, container_throughput)
except Exception as e:
logger.error(f'LOGGER: Exception occurred: {e}')
db_dict[db] = {
'containers': container_dict,
'throughput': db_throughput
}
else:
db_dict[db] = {
'containers': container_dict,
'throughput': db_throughput
}
return db_dict
# --------------------------------------------------------------------------------
# Demo 2:
from azure.cosmos import exceptions, CosmosClient, PartitionKey
#Cosmos DB - SQL API
print("Demo / Cosmos DB - SQL API ... Start")
client = CosmosClient('https://CosmosDBAccountName.documents.azure.com:443/',
'CosmosDBAccountKey')
database = client.create_database_if_not_exists("testdb", offer_throughput=400)
container1 = database.create_container_if_not_exists(
id="cont1", partition_key=PartitionKey(path="/pk"))
container2 = database.create_container_if_not_exists(
id="cont2", partition_key=PartitionKey(path="/pk"))
for x in client.list_databases():
print(x)
print("collection : ")
for y in database.list_containers():
print(y)
for i in range(1, 10):
container1.upsert_item({
'id': 'item{0}'.format(i),
'productName': 'Widget',
'productModel': 'Model {0}'.format(i),
'pk': 'test'
})
#create
container1.create_item({
