def test_update_releases(self, mock_now, mock_list_containers): # Mock time mock_now.return_value = pendulum.datetime(year=2020, month=5, day=1, minute=10) # Mock containers containers_in = make_containers() mock_list_containers.return_value = containers_in # Mock fetching of containers client = MagArchiverClient(account_name=self.account_name, account_key=self.account_key) containers = client.list_containers(last_modified_thresh=1) try: # Update releases based on containers num_updated, num_errors = client.update_releases(containers) self.assertEqual(num_updated, 2) self.assertEqual(num_errors, 0) finally: # Clean up service = TableService(account_name=self.account_name, account_key=self.account_key) for container in containers: service.delete_entity(MagRelease.TABLE_NAME, 'mag', container.name.replace("mag-", ""))
class AzureTableProvider: def __init__(self, account_name, key, table_name): self.target_table = table_name if not account_name or not key: raise Exception('Account or key not specified') self.table_service = TableService(account_name=account_name, account_key=key) self.table_service.create_table(self.target_table) def get_all(self): return self.table_service.query_entities(self.target_table) def remove(self, item): query_str = "Link eq '%s'" % item tasks = self.table_service.query_entities(self.target_table, filter=query_str) if any(tasks): for task in tasks: self.table_service.delete_entity(self.target_table, task.PartitionKey, task.RowKey) return True return False def add(self, item): track = { 'PartitionKey': 'MusicBotEntry', 'RowKey': str(uuid.uuid4()), 'Link': item } self.table_service.insert_entity(self.target_table, track)
def main(req: func.HttpRequest) -> func.HttpResponse: logging.info('Unregister request recieved.') server_id = req.params.get('server-id') region = req.params.get('region') if not server_id: return http_utils.create_function_response( {'message': 'Missing required parameter: server-id'}, status_code=400) if not region: return http_utils.create_function_response( {'message': 'Missing required parameter: region'}, status_code=400) try: table_name = 'servers' table_service = TableService( connection_string=os.environ['AzureWebJobsStorage']) table_service.delete_entity(table_name, region, server_id) return http_utils.create_function_response( {'message': f'Server {server_id} successfully unregistered.'}, 200) except: return http_utils.create_function_response( {'message': f'Server {server_id} not found.'}, 400)
def delete_orphan_record(partition_key, short_id, instance_name, account_key): start = time.time() table_service = TableService(account_name=instance_name, account_key=account_key) table_service.delete_entity('orphans', partition_key, short_id) print "delete_orphan_record completed in {} sec".format(time.time()-start)
class AzureCosmosDb(TableStorage): def __init__(self, config: AzureCosmosDbConfig): self._tableService = TableService(account_name=config.account_name, account_key=config.account_key) self._tableName = config.table_name def check_entry_exists(self, entry): try: self.query(entry['PartitionKey'], entry['RowKey']) return True except: return False def write(self, entry): prepared = entry_storage.EntryOperations.prepare_entry_for_insert( entry) if not self.check_entry_exists(prepared): self._tableService.insert_entity(self._tableName, prepared) else: self._tableService.update_entity(self._tableName, prepared) def query(self, partitionkey, rowkey): task = self._tableService.get_entity(self._tableName, partitionkey, rowkey) return task def delete(self, partitionkey, rowkey): self._tableService.delete_entity(self._tableName, partitionkey, rowkey)
def migrate_notification_keys(table_service: TableService) -> None: table_name = "Notification" notifications = table_service.query_entities( table_name, select="PartitionKey,RowKey,config") partitionKey = None count = 0 for entry in notifications: try: UUID(entry.PartitionKey) continue except ValueError: pass table_service.insert_or_replace_entity( table_name, { "PartitionKey": entry.RowKey, "RowKey": entry.PartitionKey, "config": entry.config, }, ) table_service.delete_entity(table_name, entry.PartitionKey, entry.RowKey) count += 1 print("migrated %s rows" % count)
def store_predictions_in_table(predictions, times, table_name="predictedSoilMoistureMessages"): # Connect to account table_service = TableService( account_name='soilhumiditydata293s', account_key= '4PSsEO1xBAIdq3/MppWm+t6eYHi+CWhVn6xNZ6i4mLVgm50K8+NK6lA94v8MxG0bvVEfYCvsv1suxCyCnUYd0A==' ) # Delete existing table predictions table = table_service.query_entities(table_name) for entry in table: table_service.delete_entity(table_name, entry['PartitionKey'], entry['RowKey']) #'tasksSeattle', '001') # Store values in table for i in range(len(predictions)): new_entity = Entity() new_entity.PartitionKey = datetime.strftime(times[i], "%Y-%m-%d %H:%M:%S") new_entity.RowKey = str(i) new_entity['soilmoistureprediction'] = str(predictions[i]) table_service.insert_entity(table_name, new_entity)
def main(req: func.HttpRequest) -> func.HttpResponse: logging.info('Python HTTP Submit trigger received a request') logging.debug('Creating blob service') table_service = TableService( account_name=os.getenv('AZURE_STORAGE_ACCOUNT'), account_key=os.getenv('AZURE_STORAGE_ACCESS_KEY') ) headers_dict = { "Access-Control-Allow-Credentials": "true", "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "Post" } schema = getjob_schema.GetJobSchema() try: getjob_dict = schema.loads(req.get_body()) except ValidationError: return func.HttpResponse(f'Failed to validate getjob schema', headers=headers_dict, status_code=400 ) if not getjob_dict['num_messages'] == 1: return func.HttpResponse(f'Number of messages should be 1', headers=headers_dict, status_code=400 ) table_name = os.getenv('AZURE_TABLE_NAME') entity = None entities = table_service.query_entities(table_name, filter="PartitionKey eq 'await'") for entity in entities: break if not entity: return func.HttpResponse(f'No job found', headers=headers_dict, status_code=400 ) message = {} message['crop'] = entity.crop message['geometry'] = json.loads(entity.geometry) message['irrigated'] = entity.irrigated message['guid'] = entity.RowKey message['area_name'] = entity.area_name message['planting_date'] = entity.planting_date message['fraction'] = entity.fraction table_service.delete_entity(table_name, entity.PartitionKey, entity.RowKey) entity.PartitionKey = 'processing' table_service.insert_entity(table_name, entity) return func.HttpResponse(json.dumps(message), headers=headers_dict, mimetype='application/json' )
class AzureTable(): def __init__(self, account_name, account_key): self.table_service = TableService(account_name=account_name, account_key=account_key) def create_table(self, table_name): return self.table_service.create_table(table_name) def exists_table(self, table_name): return self.table_service.exists(table_name) def insert_or_replace_entity(self, table_name, partition_key, row_key, **kwargs): try: entity = self.table_service.get_entity(table_name, partition_key, row_key) except Exception: # Insert a new entity entity = {'PartitionKey': partition_key, 'RowKey': row_key} for (k, v) in kwargs.items(): entity[k] = v return self.table_service.insert_or_replace_entity(table_name, entity) def insert_or_replace_entity2(self, table_name, entity): return self.table_service.insert_or_replace_entity(table_name, entity) def insert_entity(self, table_name, entity): return self.table_service.insert_entity(table_name, entity) def update_entity(self, table_name, entity): return self.table_service.update_entity(table_name, entity) def get_entity(self, table_name, partition_key, row_key): return self.table_service.get_entity(table_name, partition_key, row_key) def delete_entity(self, table_name, partition_key, row_key): self.table_service.delete_entity(table_name, partition_key, row_key) def delete_table(self, table_name): return self.table_service.delete_table(table_name) def get_entities(self, table_name, partition_key): filter = "PartitionKey eq '{0}'".format(partition_key) return self.table_service.query_entities(table_name, filter)
def main(req: func.HttpRequest) -> func.HttpResponse: logging.info('Python HTTP Submit trigger received a request') logging.debug('Creating blob service') headers_dict = { "Access-Control-Allow-Credentials": "true", "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "Post" } schema = put_schema.DoneSchema() try: done_dict = schema.loads(req.get_body()) except ValidationError: error = f'Failed to validate the done message' return func.HttpResponse(error, headers=headers_dict, status_code=400 ) table_name = os.getenv('AZURE_TABLE_NAME') table_service = TableService( account_name=os.getenv('AZURE_STORAGE_ACCOUNT'), account_key=os.getenv('AZURE_STORAGE_ACCESS_KEY') ) try: entity = table_service.get_entity(table_name, 'processing', done_dict['guid']) except AzureMissingResourceHttpError: error = f'Failed to put done message' return func.HttpResponse(error, headers=headers_dict, status_code=400 ) if not done_dict['error']: entity.Error = "" else: entity.Error = done_dict['error'] table_service.delete_entity(table_name, entity.PartitionKey, entity.RowKey) entity.PartitionKey = 'done' table_service.insert_entity(table_name, entity) return func.HttpResponse('Message was successfully inserted into Done queue', headers=headers_dict )
def main(req: func.HttpRequest) -> func.HttpResponse: logging.info('Starting delete rows.') table_name = req.headers.get('name') column = req.headers.get('column') pattern = req.headers.get('pattern') if not table_name: #If name wasnt added as header, search for it in the parameters table_name = req.params.get('name') if not column: #If column wasnt added as header, search for it in the parameters column = req.params.get('column') if not pattern: #If pattern wasnt added as header, search for it in the parameters pattern = req.params.get('pattern') ret = dict() if table_name: retrieved_secret = getConnectionString() table_service = TableService(connection_string=retrieved_secret.value) df = get_dataframe_from_table_storage_table(table_service, table_name) #Filter dataframe by pattern to_delete = df[df[column].str.contains(pattern).fillna(value=False)] #Loop over the dataframe and delete records for i, o in to_delete.iterrows(): logging.info('Deleting {}'.format(i+1)) table_service.delete_entity(table_name, partition_key=o['PartitionKey'], row_key=o['RowKey']) ret['result'] = "Deleted {} rows!".format(to_delete.shape[0]) return func.HttpResponse( json.dumps(ret), status_code=200 ) else: ret['result'] = "Please pass a table name!!" return func.HttpResponse( json.dumps(ret), status_code=400 )
class AzureTable(object): def __init__(self, account_name: str, account_key: str, table_name: str, partition_key_field: str, clustering_key_field: str): self.table = TableService(account_name=account_name, account_key=account_key) self.table_name = self.table_name self.partition_key_field = partition_key_field self.clustering_key_field = clustering_key_field @property def partition_key_name(self) -> str: return 'PartitionKey' @property def clustering_key_name(self) -> str: return 'RowKey' def get_payload(self, payload: dict): item = deepcopy(payload) partition_key = payload.get(self.partition_key_field) clustering_key = payload.get(self.clustering_key_field) if partition_key is None: raise PartitionKeyNotFoundError( 'payload={} does not have a partition key') if clustering_key is None: raise ClusteringKeyNotFoundError( 'payload={} does not have a clustering key') item.update({ self.partition_key_name: partition_key, self.clustering_key_name: clustering_key }) return item def create(self): return self.table.create_table(self.table_name) def insert(self, item: dict): return self.table.insert_entity(self.table_name, self.get_payload(item)) def update(self, item: dict): return self.table.update_entity(self.table_name, self.get_payload(item)) def upsert(self, item: dict): return self.table.insert_or_replace_entity(self.table_name, self.get_payload(item)) def delete(self, partition_key: str, clustering_key: str): return self.table.delete_entity(self.table_name, partition_key=partition_key, row_key=clustering_key) def read(self, partition_key: str, clustering_key: str): return self.table.get_entity(self.table_name, partition_key=partition_key, row_key=clustering_key) def insert_batch(self, items: list): batch = TableBatch() for item in items: batch.insert_entity(self.get_payload(item)) return self.table.commit_batch(self.table_name, batch) def get(self, partition_key: str, clustering_key: str): return self.table.get_entity(self.table_name, partition_key, clustering_key) def get_by_partition(self, partition_key: str) -> list: return self.table.query_entities(self.table_name, filter="{} eq '{}'".format( self.partition_key_name, partition_key))
class CosmosData: # Setup initial variables def __init__(self, account_key, identifier): self.account_key = account_key self.identifier = identifier self.account_name = self.identifier + "-cosmos-db" self.table_endpoint = "https://" + self.identifier + "-cosmos-db.table.cosmos.azure.com:443/" self.connection_string = "DefaultEndpointsProtocol=https;AccountName=" + self.account_name + ";AccountKey=" + self.account_key + ";TableEndpoint=" + self.table_endpoint + ";" self.table_service = TableService( endpoint_suffix="table.cosmos.azure.com", connection_string=self.connection_string) self.table_name = self.identifier + "-cosmos-table" # Get a list of values def get_list(self, list_name, provider=None): pkey = "" if provider is not None: pkey = "|" + provider main_list = [] if list_name: try: items = self.table_service.query_entities( self.table_name, filter="PartitionKey eq '" + list_name + "'") except ValueError: pass else: for item in items: if provider is None: main_list.append(str(item.RowKey)) elif str(item.RowKey).endswith(provider): main_list.append(str(item.RowKey).replace(pkey, "")) return main_list # Add a value to a list def add_to_list(self, list_name, value, provider=None): pkey = "" if provider is not None: pkey = "|" + provider item = Entity() item.PartitionKey = list_name item.RowKey = value + pkey main_list = self.get_list(list_name, provider) try: self.table_service.insert_entity(self.table_name, item) except ValueError: pass else: main_list.append(value) return main_list # Remove a value from a list def remove_from_list(self, list_name, value, provider=None): pkey = "" if provider is not None: pkey = "|" + provider main_list = self.get_list(list_name, provider) try: self.table_service.delete_entity(self.table_name, list_name, value + pkey) except ValueError: pass else: main_list.remove(value) return main_list # Get a value from a K/V pair def get_value(self, key): retval = "" try: items = self.table_service.query_entities( self.table_name, filter="PartitionKey eq '" + key + "'") except ValueError: pass else: for item in items: retval = item.RowKey return retval # Set a value for a K/V pair def set_value(self, key, value): item = Entity() item.PartitionKey = key item.RowKey = value retval = False try: entries = self.table_service.query_entities( self.table_name, filter="PartitionKey eq '" + key + "'") old_value = "invalid" for entry in entries: old_value = entry.RowKey self.table_service.delete_entity(self.table_name, key, old_value) self.table_service.insert_entity(self.table_name, item) except ValueError: pass else: retval = True return retval # Delete all data and load in the default data set def data_load(self): data_set = { "required-modules": [ "custom-vnet|azurerm", "custom-sg|azurerm", "custom-blob|azurerm", "custom-vpc|aws", "custom-sg|aws" ], "approved-instances": [ "Standard_A1_v2|azurerm", "Standard_A2_v2|azurerm", "Standard_A4_v2|azurerm", "Standard_A8_v2|azurerm", "t3.micro|aws", "t3.small|aws", "t3.medium|aws", "t3.large|aws" ], "prohibited-resources": [ "azurerm_resource_group|azurerm", "azurerm_virtual_network|azurerm", "azurerm_network_security_group|azurerm", "azurerm_subnet_network_security_group_association|azurerm", "aws_internet_gateway|aws", "aws_route|aws", "aws_route_table|aws", "aws_route_table_association|aws", "aws_subnet|aws", "aws_vpc|aws", "aws_security_group|aws" ], "allowed-resources": [ "azurerm_virtual_machine|azurerm", "azurerm_network_interface|azurerm", "azurerm_public_ip|azurerm", "azurerm_storage_account|azurerm", "aws_instance|aws", "aws_s3_bucket|aws", "aws_s3_bucket_policy|aws" ], "prevent-deletion": ["true"], "default-provider": ["azurerm"], "mandatory-tags": ["Department", "Environment"], "max-cost": ["100"], "ddb-encryption": ["true"], "no-star-access": ["true"] } # delete all entries items = self.table_service.query_entities(self.table_name) for itm in items: self.table_service.delete_entity(self.table_name, itm.PartitionKey, itm.RowKey) # add all entries for category in data_set: for value in data_set[category]: item = Entity() item.PartitionKey = category item.RowKey = value self.table_service.insert_entity(self.table_name, item) return True
print(task.description) print(task.priority) # Query a set of entities print("Query set of entities...") tasks = table_service.query_entities( 'tasktable', filter="PartitionKey eq 'tasksSeattle'") for task in tasks: print(task.description) print(task.priority) try: print(task.newCol) except AttributeError: print("No newCol.") # Query a subset of entity properties print("Query a subset of entity properties...") tasks = table_service.query_entities( 'tasktable', filter="PartitionKey eq 'tasksSeattle'", select='description') for task in tasks: print(task.description) # Delete an entity print("Delete an entity...") table_service.delete_entity('tasktable', 'tasksSeattle', '001') # Delete a table print("Delete a table...") table_service.delete_table('tasktable')
class AzureOperationsStorage(BasicOperationStorage): """ Implementation of :class:`.interface.IOperationStorage` with Azure Table Storage using the default implementation :class:`.interface.BasicOperationStorage` On creating a connection to the storage is initialized and all needed tables are created. If a purge is necessary, tables are not deleted but simple the content removed. Table creation can take a while with Azure Table Storage. As Azure Table Storage only supports two indices, the operations are inserted multiple times in different tables to enable multi-index queries. """ def get_retry_exceptions(self): return (NewConnectionError) @retry_auto_reconnect def __init__(self, azure_config, purge=False): super(AzureOperationsStorage, self).__init__() if not azure_config: raise Exception("No azure table storage configuration provided!") self._azure_config = azure_config # ensure defaults self._azure_config["operation_table"] = self._azure_config.get( "operation_table", "operations") self._azure_config["address_table"] = self._azure_config.get( "address_table", "address") self._azure_config["status_table"] = self._azure_config.get( "status_table", "status") self._azure_config["balances_table"] = self._azure_config.get( "balances_table", "balances") if not self._azure_config["account"]: raise Exception( "Please include the azure account name in the config") if not self._azure_config["key"]: raise Exception( "Please include the azure account key in the config") self._service = TableService( account_name=self._azure_config["account"], account_key=self._azure_config["key"]) # if tables doesnt exist, create it self._create_operations_storage(purge) self._create_status_storage(purge) self._create_address_storage(purge) self._create_balances_storage(purge) def _debug_print(self, operation): from pprint import pprint pprint(operation) def _create_address_storage(self, purge): _varients = ["balance", "historyfrom", "historyto"] for variant in _varients: tablename = self._azure_config["address_table"] + variant if purge: try: for item in self._service.query_entities(tablename): self._service.delete_entity(tablename, item["PartitionKey"], item["RowKey"]) except AzureHttpError: pass except AzureMissingResourceHttpError: pass while not self._service.exists(tablename): self._service.create_table(tablename) time.sleep(0.1) def _create_status_storage(self, purge): if purge: try: tablename = self._azure_config["status_table"] for item in self._service.query_entities(tablename): self._service.delete_entity(tablename, item["PartitionKey"], item["RowKey"]) except AzureMissingResourceHttpError: pass while not self._service.exists(self._azure_config["status_table"]): self._service.create_table(self._azure_config["status_table"]) time.sleep(0.1) def _create_balances_storage(self, purge): if purge: try: tablename = self._azure_config["balances_table"] for item in self._service.query_entities(tablename): self._service.delete_entity(tablename, item["PartitionKey"], item["RowKey"]) except AzureMissingResourceHttpError: pass while not self._service.exists(self._azure_config["balances_table"]): self._service.create_table(self._azure_config["balances_table"]) time.sleep(0.1) def _create_operations_storage(self, purge): self._operation_varients = [ "incident", "statuscompleted", "statusfailed", "statusinprogress" ] # "customer" self._operation_tables = {} for variant in self._operation_varients: self._operation_tables[ variant] = self._azure_config["operation_table"] + variant self._operation_prep = { "statusinprogress": lambda op: { "PartitionKey": self._short_digit_hash(op["chain_identifier"]), "RowKey": op["chain_identifier"] }, "statuscompleted": lambda op: { "PartitionKey": self._short_digit_hash(op["chain_identifier"]), "RowKey": op["chain_identifier"] }, "statusfailed": lambda op: { "PartitionKey": self._short_digit_hash(op["chain_identifier"]), "RowKey": op["chain_identifier"] }, "customer": lambda op: { "PartitionKey": op["customer_id"], "RowKey": op["chain_identifier"] }, "incident": lambda op: { "PartitionKey": self._short_digit_hash(op["incident_id"]), "RowKey": op["incident_id"] } } for variant in self._operation_varients: if purge: try: tablename = self._operation_tables[variant] for item in self._service.query_entities(tablename): self._service.delete_entity(tablename, item["PartitionKey"], item["RowKey"]) except AzureMissingResourceHttpError: pass while not self._service.exists(self._operation_tables[variant]): self._service.create_table(self._operation_tables[variant]) time.sleep(0.1) def _get_with_ck(self, variant, operation): with_ck = operation.copy() with_ck.update(self._operation_prep[variant](with_ck)) return with_ck def _short_digit_hash(self, value): hash_type = Config.get("operation_storage", "key_hash", "type", default="crc32") if hash_type == "crc32": short_hash = hex(zlib.crc32(value.encode(encoding='UTF-8'))) short_hash = short_hash[2:len(short_hash)] elif hash_type == "sha256": checker = hashlib.sha256() checker.update(value.encode(encoding='UTF-8')) short_hash = checker.hexdigest() return short_hash[0:Config. get("operation_storage", "key_hash", "digits", 3)] @retry_auto_reconnect def track_address(self, address, usage="balance"): address = ensure_address_format(address) try: short_hash = self._short_digit_hash(address) logging.getLogger(__name__).debug("track_address with " + str(address) + ", hash " + str(short_hash)) self._service.insert_entity( self._azure_config["address_table"] + usage, { "PartitionKey": short_hash, "RowKey": address, "address": address, "usage": usage }) except AzureConflictHttpError: raise AddressAlreadyTrackedException @retry_auto_reconnect def untrack_address(self, address, usage="balance"): address = ensure_address_format(address) try: short_hash = self._short_digit_hash(address) logging.getLogger(__name__).debug("untrack_address with " + str(address) + ", hash " + str(short_hash)) self._service.delete_entity( self._azure_config["address_table"] + usage, short_hash, address) try: self._delete_balance(address) except AzureMissingResourceHttpError: pass except AzureMissingResourceHttpError: raise AddressNotTrackedException() @retry_auto_reconnect def _get_address(self, address, usage="balance"): try: short_hash = self._short_digit_hash(address) logging.getLogger(__name__).debug("_get_address with " + str(address) + ", hash " + str(short_hash)) return self._service.get_entity( self._azure_config["address_table"] + usage, short_hash, address) except AzureMissingResourceHttpError: raise AddressNotTrackedException() def _update(self, operation, status=None): try: mapping = { "in_progress": "statusinprogress", "completed": "statuscompleted", "failed": "statusfailed" } operation = self._get_with_ck("incident", operation.copy()) new_operation = operation if status: tmp = self.get_operation(operation["incident_id"]) new_operation["timestamp"] = tmp["timestamp"] new_operation["status"] = status new_operation = self._get_with_ck("incident", new_operation) logging.getLogger(__name__).debug( "_update: Table " + self._operation_tables["incident"] + " PartitionKey " + new_operation["PartitionKey"] + " " + new_operation["RowKey"]) self._service.update_entity(self._operation_tables["incident"], new_operation) operation = self._get_with_ck("statuscompleted", operation.copy()) new_operation = operation if status: tmp = self.get_operation(operation["incident_id"]) new_operation["timestamp"] = tmp["timestamp"] new_operation["status"] = status new_operation = self._get_with_ck("statuscompleted", new_operation) self._service.update_entity( self._operation_tables["statuscompleted"], new_operation) logging.getLogger(__name__).debug( "_update: Table " + self._operation_tables["statuscompleted"] + " PartitionKey " + new_operation["PartitionKey"] + " " + new_operation["RowKey"]) if status: # needs delete and insert try: self._service.delete_entity( self._operation_tables[mapping[operation["status"]]], operation["PartitionKey"], operation["RowKey"]) except AzureMissingResourceHttpError: pass try: self._service.insert_entity( self._operation_tables[mapping[ new_operation["status"]]], new_operation) except AzureConflictHttpError: # already exists, try update self._service.update_entity( self._operation_tables[mapping[ new_operation["status"]]], new_operation) else: self._service.update_entity( self._operation_tables[mapping[new_operation["status"]]], new_operation) except AzureMissingResourceHttpError: raise OperationNotFoundException() def _insert(self, operation): try: for variant in self._operation_varients: to_insert = operation.copy() to_insert.update(self._operation_prep[variant](to_insert)) if not to_insert["PartitionKey"]: raise AzureMissingResourceHttpError() if not to_insert["RowKey"]: raise AzureMissingResourceHttpError() logging.getLogger(__name__).debug( "_insert: Table " + self._operation_tables[variant] + " PartitionKey " + to_insert["PartitionKey"] + " " + to_insert["RowKey"]) self._service.insert_entity(self._operation_tables[variant], to_insert) except AzureConflictHttpError: raise DuplicateOperationException() def _delete(self, operation): try: for variant in self._operation_varients: to_delete = operation.copy() to_delete.update(self._operation_prep[variant](to_delete)) self._service.delete_entity(self._operation_tables[variant], to_delete["PartitionKey"], to_delete["RowKey"]) except AzureMissingResourceHttpError: raise OperationNotFoundException() @retry_auto_reconnect def flag_operation_completed(self, operation): # do basics operation = super(AzureOperationsStorage, self).flag_operation_completed(operation) self._update(operation, status="completed") self._ensure_balances(operation) @retry_auto_reconnect def flag_operation_failed(self, operation, message=None): # do basics operation = super(AzureOperationsStorage, self).flag_operation_failed(operation) operation["message"] = message self._update(operation, status="failed") @retry_auto_reconnect def insert_operation(self, operation): # do basics operation = super(AzureOperationsStorage, self).insert_operation(operation) error = None try: self._insert(operation) except DuplicateOperationException as e: error = e try: # always check if balances are ok if operation["status"] == "completed": self._ensure_balances(operation) except BalanceConcurrentException as e: if error is None: error = e if error is not None: raise error @retry_auto_reconnect def _delete_balance(self, address, if_match='*'): self._service.delete_entity(self._azure_config["balances_table"], self._short_digit_hash(address), address, if_match=if_match) @retry_auto_reconnect def _ensure_balances(self, operation): affected_address = get_tracking_address(operation) logging.getLogger(__name__).debug("_ensure_balances: with " + operation["chain_identifier"] + " for address " + str(affected_address)) try: self._get_address(affected_address) except AddressNotTrackedException: # delte if exists and return try: self._delete_balance(affected_address) except AzureMissingResourceHttpError: pass return try: balance_dict = self._service.get_entity( self._azure_config["balances_table"], self._short_digit_hash(affected_address), affected_address) insert = False except AzureMissingResourceHttpError as e: balance_dict = {"address": affected_address} balance_dict["PartitionKey"] = self._short_digit_hash( balance_dict["address"]) balance_dict["RowKey"] = balance_dict["address"] insert = True if operation["block_num"] < balance_dict.get("blocknum", 0): raise BalanceConcurrentException() elif operation["block_num"] == balance_dict.get("blocknum", 0) and\ operation["txnum"] < balance_dict.get("txnum", 0): raise BalanceConcurrentException() elif operation["block_num"] == balance_dict.get("blocknum", 0) and\ operation["txnum"] == balance_dict.get("txnum", 0) and\ operation["opnum"] <= balance_dict.get("opnum", 0): raise BalanceConcurrentException() balance_dict["blocknum"] = max(balance_dict.get("blocknum", 0), operation["block_num"]) balance_dict["txnum"] = max(balance_dict.get("txnum", 0), operation["tx_in_block"]) balance_dict["opnum"] = max(balance_dict.get("opnum", 0), operation["op_in_tx"]) total = 0 addrs = split_unique_address(affected_address) asset_id = "balance" + operation["amount_asset_id"].split("1.3.")[1] if addrs["account_id"] == operation["from"]: # internal transfer and withdraw # negative balance = balance_dict.get(asset_id, 0) balance_dict[asset_id] = balance - operation["amount_value"] # fee as well asset_id = operation["fee_asset_id"] balance = balance_dict.get(asset_id, 0) balance_dict[asset_id] = balance - operation["fee_value"] elif addrs["account_id"] == operation["to"]: # deposit # positive balance = balance_dict.get(asset_id, 0) balance_dict[asset_id] = balance + operation["amount_value"] # fees were paid by someone else else: raise InvalidOperationException() for key, value in balance_dict.items(): if key.startswith("balance"): total = total + value if total == 0: if not insert: try: self._delete_balance(affected_address, if_match=balance_dict.etag) except AzureMissingResourceHttpError: pass return # may be updated or inserted, total > 0 if (insert): try: self._service.insert_entity( self._azure_config["balances_table"], balance_dict) except AzureMissingResourceHttpError: raise OperationStorageException( "Critical error in database consistency") else: try: self._service.update_entity( self._azure_config["balances_table"], balance_dict, if_match=balance_dict.etag) except AzureConflictHttpError: raise OperationStorageException( "Critical error in database consistency") @retry_auto_reconnect def insert_or_update_operation(self, operation): # do basics operation = super(AzureOperationsStorage, self).insert_operation(operation) # check if this is from in_progress to complete (for withdrawals we need to find incident id as its # not stored onchain) try: logging.getLogger(__name__).debug( "insert_or_update_operation: check if in_progress with " + str(operation["chain_identifier"]) + " exists") existing_operation = self.get_operation_by_chain_identifier( "in_progress", operation["chain_identifier"]) logging.getLogger(__name__).debug( "insert_or_update_operation: found existing in_progress operation" ) if not existing_operation["incident_id"] == operation["incident_id"] and\ operation["incident_id"] == operation["chain_identifier"]: logging.getLogger(__name__).debug( "insert_or_update_operation: using preset incident_id " + str(existing_operation["incident_id"])) operation["incident_id"] = existing_operation["incident_id"] except OperationNotFoundException: existing_operation = None if existing_operation is None: try: logging.getLogger(__name__).debug( "insert_or_update_operation: attempting insert") error = None try: self._insert(operation) except DuplicateOperationException as e: error = e try: # always check if balances are ok if operation["status"] == "completed": self._ensure_balances(operation) except BalanceConcurrentException as e: if error is None: error = e if error is not None: raise error except DuplicateOperationException as ex: logging.getLogger(__name__).debug( "insert_or_update_operation: fallback to update") # could be an update to completed ... if operation.get("block_num"): try: operation.pop("status") self.flag_operation_completed(operation) except OperationNotFoundException: raise ex else: raise ex else: logging.getLogger(__name__).debug( "insert_or_update_operation: attempting update") if operation.get("block_num"): try: operation.pop("status") self.flag_operation_completed(operation) except OperationNotFoundException: raise ex @retry_auto_reconnect def delete_operation(self, operation_or_incident_id): # do basics operation = super(AzureOperationsStorage, self).delete_operation(operation_or_incident_id) if type(operation_or_incident_id) == str: operation = self.get_operation(operation_or_incident_id) else: operation = operation_or_incident_id self._delete(operation) @retry_auto_reconnect def get_operation_by_chain_identifier(self, status, chain_identifier): mapping = { "in_progress": "statusinprogress", "completed": "statuscompleted", "failed": "statusfailed" } try: operation = self._service.get_entity( self._operation_tables[mapping[status]], self._short_digit_hash(chain_identifier), chain_identifier) operation.pop("PartitionKey") operation.pop("RowKey") operation.pop("Timestamp") operation.pop("etag") except AzureMissingResourceHttpError: raise OperationNotFoundException() return operation @retry_auto_reconnect def get_operation(self, incident_id): try: short_hash = self._short_digit_hash(incident_id) logging.getLogger(__name__).debug("get_operation with " + str(incident_id) + ", hash " + str(short_hash)) operation = self._service.get_entity( self._operation_tables["incident"], short_hash, incident_id) operation.pop("PartitionKey") operation.pop("RowKey") operation.pop("Timestamp") operation.pop("etag") except AzureMissingResourceHttpError: raise OperationNotFoundException() return operation @retry_auto_reconnect def get_balances(self, take, continuation=None, addresses=None, recalculate=False): if recalculate: raise Exception( "Currently not supported due to memo change on withdraw") return self._get_balances_recalculate(take, continuation, addresses) else: if continuation is not None: try: continuation_marker = json.loads(continuation) except TypeError: raise InputInvalidException() except JSONDecodeError: raise InputInvalidException() balances = self._service.query_entities( self._azure_config["balances_table"], num_results=take, marker=continuation_marker) else: balances = self._service.query_entities( self._azure_config["balances_table"], num_results=take) return_balances = {} for address_balance in balances: return_balances[address_balance["address"]] = { "block_num": address_balance["blocknum"] } for key, value in address_balance.items(): if key.startswith("balance"): asset_id = "1.3." + key.split("balance")[1] return_balances[ address_balance["address"]][asset_id] = value return_balances["continuation"] = None if balances.next_marker: return_balances["continuation"] = json.dumps( balances.next_marker) return return_balances @retry_auto_reconnect def _get_balances_recalculate(self, take, continuation=None, addresses=None): address_balances = collections.defaultdict( lambda: collections.defaultdict()) if not addresses: if continuation is not None: try: continuation_marker = json.loads(continuation) except TypeError: raise InputInvalidException() except JSONDecodeError: raise InputInvalidException() addresses = self._service.query_entities( self._azure_config["address_table"] + "balance", num_results=take, marker=continuation_marker) else: addresses = self._service.query_entities( self._azure_config["address_table"] + "balance", num_results=take) if addresses.next_marker: address_balances["continuation"] = json.dumps( addresses.next_marker) addresses = [x["address"] for x in addresses] if type(addresses) == str: addresses = [addresses] for address in addresses: addrs = split_unique_address(address) max_block_number = 0 for operation in self.get_operations_completed( filter_by={"customer_id": addrs["customer_id"]}): this_block_num = operation["block_num"] asset_id = operation["amount_asset_id"] if addrs["account_id"] == operation["from"]: # negative balance = address_balances[address].get(asset_id, 0) address_balances[address][asset_id] =\ balance - operation["amount_value"] # fee as well asset_id = operation["fee_asset_id"] balance = address_balances[address].get(asset_id, 0) address_balances[address][asset_id] =\ balance - operation["fee_value"] elif addrs["account_id"] == operation["to"]: # positive balance = address_balances[address].get(asset_id, 0) address_balances[address][asset_id] =\ balance + operation["amount_value"] else: raise InvalidOperationException() max_block_number = max(max_block_number, this_block_num) if max_block_number > 0: address_balances[address]["block_num"] = max_block_number # do not return default dicts for key, value in address_balances.items(): if type(value) == collections.defaultdict: address_balances[key] = dict(value) return dict(address_balances) def _parse_filter(self, filter_by): if filter_by: if filter_by.get("customer_id"): return {"customer_id": filter_by.pop("customer_id")} if filter_by.get("address"): addrs = split_unique_address(filter_by.pop("address")) return {"customer_id": addrs["customer_id"]} if filter_by.get("from"): addrs = split_unique_address(filter_by.pop("from")) return {"from": addrs["account_id"]} if filter_by.get("to"): addrs = split_unique_address(filter_by.pop("to")) return {"to": addrs["account_id"]} if filter_by: raise Exception("Filter not supported") return {} def _filter_dict_to_string(self, filter_dict, partition_key=None): filter_str = None for key, value in filter_dict.items(): if partition_key == key: key = "PartitionKey" if filter_str is not None: delimiter = " and " delimiter = "" filter_str = delimiter + key + " eq '" + value + "'" return filter_str @retry_auto_reconnect def get_operations_in_progress(self, filter_by=None): mapping = { "in_progress": "statusinprogress", "completed": "statuscompleted", "failed": "statusfailed" } filter_dict = {} filter_dict.update(self._parse_filter(filter_by)) filter_str = self._filter_dict_to_string(filter_dict, "status") return list( self._service.query_entities( self._operation_tables[mapping["in_progress"]], filter_str)) @retry_auto_reconnect def get_operations_completed(self, filter_by=None): mapping = { "in_progress": "statusinprogress", "completed": "statuscompleted", "failed": "statusfailed" } filter_dict = {} filter_dict.update(self._parse_filter(filter_by)) filter_str = self._filter_dict_to_string(filter_dict, "status") return list( self._service.query_entities( self._operation_tables[mapping["completed"]], filter_str)) @retry_auto_reconnect def get_operations_failed(self, filter_by=None): mapping = { "in_progress": "statusinprogress", "completed": "statuscompleted", "failed": "statusfailed" } filter_dict = {} filter_dict.update(self._parse_filter(filter_by)) filter_str = self._filter_dict_to_string(filter_dict, "status") return list( self._service.query_entities( self._operation_tables[mapping["failed"]], filter_str)) @retry_auto_reconnect def get_last_head_block_num(self): try: document = self._service.get_entity( self._azure_config["status_table"], "head_block_num", "last") return document["last_head_block_num"] except AzureMissingResourceHttpError: return 0 @retry_auto_reconnect def set_last_head_block_num(self, head_block_num): current_last = self.get_last_head_block_num() if current_last >= head_block_num: raise Exception("Marching backwards not supported! Last: " + str(current_last) + " New: " + str(head_block_num)) self._service.insert_or_replace_entity( self._azure_config["status_table"], { "PartitionKey": "head_block_num", "RowKey": "last", "last_head_block_num": head_block_num })
config = configparser.ConfigParser() config.read('app.ini') identifier = config['App']['Identifier'] account_name = identifier + "-cosmos-db" account_key = config['App']['AccountKey'] table_endpoint = "https://" + identifier + "-cosmos-db.table.cosmos.azure.com:443/" connection_string = "DefaultEndpointsProtocol=https;AccountName=" + account_name + ";AccountKey=" + account_key + ";TableEndpoint=" + table_endpoint + ";" table_service = TableService(endpoint_suffix="table.cosmos.azure.com", connection_string=connection_string) table_name = identifier + "-cosmos-table" # Start over items = table_service.query_entities(table_name) for itm in items: table_service.delete_entity(table_name, itm.PartitionKey, itm.RowKey) # Add new entries item = Entity() item.PartitionKey = "required-modules" item.RowKey = "custom-vnet" table_service.insert_entity('kevincx-cosmos-table', item) item = Entity() item.PartitionKey = "required-modules" item.RowKey = "custom-sg" table_service.insert_entity('kevincx-cosmos-table', item) item = Entity() item.PartitionKey = "required-modules" item.RowKey = "custom-blob"
class AzureCosmosDb(TableStorage): """Azure CosmosDB provider for Table Storage.""" def __init__(self, config: AzureCosmosDbConfig): self._table_service = TableService(account_name=config.account_name, account_key=config.account_key) self._tableName = config.table_name def check_entry_exists(self, entry): """Check if entry exists in table. :param entry: Dictionary with PartitionKey and RowKey fields :return: True if entry exists """ try: self.query(entry['PartitionKey'], entry['RowKey']) return True except Exception: return False def write(self, resource): """Write resource to table. :param resource: Expecting Resource object (see Common.Contracts.Resource) :return: None """ entry = resource.to_dict() prepared = entry_storage.EntryOperations.prepare_entry_for_insert( entry) if not self.check_entry_exists(prepared): self._table_service.insert_entity(self._tableName, prepared) else: self._table_service.update_entity(self._tableName, prepared) def query(self, partition_key, row_key): """Get entry with specified partition and row keys. :param partition_key: Partition key for entry :param row_key: Row key for entry :return: Entity if found, None otherwise """ task = self._table_service.get_entity(self._tableName, partition_key, row_key) return task def query_list(self): """Get entities from table. :return: List of entities from table """ return self._table_service.query_entities(self._tableName) def delete(self, partition_key, row_key): """Delete entry with specified partition and row keys. :param partition_key: Partition key for entry :param row_key: Row key for entry :return: None """ self._table_service.delete_entity(self._tableName, partition_key, row_key) @staticmethod def create(): """Initialize AzureCosmosDb service. :return: AzureCosmosDb service object """ config = AzureConfig() cosmos_storage = AzureCosmosDb(config.cosmos_storage_config) return cosmos_storage
class AzureStorageHandler(): def __init__(self, kv): try: self.table_service = TableService( account_name=kv.get_keyvault_secret("storageAccount-name"), account_key=kv.get_keyvault_secret("storageAccount-key")) # Quick start code goes here except Exception as ex: print('Exception:') print(ex) def insert_submission_entry(self, entry): submission = Entity() submission.PartitionKey = entry.subreddit submission.RowKey = entry.id submission.author = entry.author submission.created_utc = entry.created_utc submission.flair = entry.flair submission.title = entry.title # Flatten list of keywords into comma separated string submission.title_keywords = ','.join(map(str, entry.title_keywords)) submission.title_sentiment = entry.title_sentiment try: submission.body_keywords = ','.join(map(str, entry.body_keywords)) submission.body_sentiment = entry.body_sentiment except AttributeError: submission.body_keywords = "" submission.body_sentiment = "" self.table_service.insert_entity('submissions', submission) def insert_comment_entry(self, entries): for entry in entries: comment = Entity() comment.PartitionKey = entry.link_id comment.RowKey = entry.id comment.author = entry.author comment.body = entry.body comment.created_utc = entry.created_utc comment.parent_id = entry.parent_id comment.score = entry.score comment.subreddit = entry.subreddit comment.subreddit_id = entry.subreddit_id comment.total_awards_received = entry.total_awards_received comment.sentiment = entry.sentiment # Flatten list of keywords into comma separated string comment.keywords = ','.join(map(str, entry.keywords)) self.table_service.insert_entity('comments', comment) def insert_recommendation_entry(self, entries): for entry in entries: recommendation = Entity() recommendation.PartitionKey = "{0}_{1}".format( entry.subreddit, entry.query_word) recommendation.RowKey = entry.keyword recommendation.subreddit = entry.subreddit recommendation.query_word = entry.query_word recommendation.post_id = ','.join(map(str, entry.post_id)) recommendation.comment_id = ','.join(map(str, entry.comment_id)) recommendation.sentiment = entry.sentiment recommendation.count = entry.count try: self.table_service.insert_entity('recommendations', recommendation) except AzureConflictHttpError as error: # print(error) subreddit_query_word = recommendation.PartitionKey.split('_') print( "The recommendation entry with subreddit = '{0}', search term = '{1}', and keyword = '{2}' already exists in the database. Updating it..." .format(subreddit_query_word[0], subreddit_query_word[1], recommendation.RowKey)) self.table_service.update_entity('recommendations', recommendation) def insert_sub_date_entry(self, entry): sub_date = Entity() sub_date.PartitionKey = entry.subreddit sub_date.RowKey = entry.title sub_date.created_utc = entry.created_utc sub_date.post_id = entry.post_id try: self.table_service.insert_or_replace_entity( 'mostrecentsubdate', sub_date) except TypeError as error: print(error) print( f"The mostrecentsubdate object is formatted incorrectly and was not updated. One of the parameters is not an int, str, bool or datetime, or defined custom EntityProperty. Continuing..." ) def get_entry(self, table, partition_key, row_key): return self.table_service.get_entity(table, partition_key, row_key) def filter_entries(self, table, filter_string): return self.table_service.query_entities(table, filter_string) def update_entry(self, table, entity): return self.table_service.update_entity(table, entity) def delete_entry(self, table, partition_key, row_key): return self.table_service.delete_entity(table, partition_key, row_key)
def main(req: func.HttpRequest) -> func.HttpResponse: logging.info("Python HTTP trigger function processed a request.") def get_param(req: func.HttpRequest, param: str) -> str: param_value = req.params.get(param) if not param_value: try: req_body = req.get_json() except ValueError: pass else: param_value = req_body.get(param) return param_value path = get_param(req, "path") logging.info( f"Parsed the following parameters from the HTTP request: path={path}.") if path: rowKey = path partitionKey = "links" tableName = "links" # connect to table storage_account_name = os.environ["STORAGE_ACCOUNT_NAME"] storage_account_key = os.environ["STORAGE_ACCOUNT_KEY"] table_service = TableService(account_name=storage_account_name, account_key=storage_account_key) # try to delete the entity - if it does not exist, throw try: table_service.delete_entity(tableName, partitionKey, rowKey) logging.info( f"Deleted the following from Azure Table Storage: PartitionKey={partitionKey}, RowKey={rowKey}" ) resp = { "Status": "Ok", "Message": f"Deleted the following from Azure Table Storage: PartitionKey={partitionKey}, RowKey={rowKey}", } except: logging.info( f"Skipped deletion of the following from Azure Table Storage as the entity does not exist: PartitionKey={partitionKey}, RowKey={rowKey}" ) resp = { "Status": "Ok", "Message": f"Skipped deletion of the following from Azure Table Storage as it does not exist: PartitionKey={partitionKey}, RowKey={rowKey}", } pass resp_json = json.dumps(resp) return func.HttpResponse(body=resp_json, mimetype="application/json") else: resp = { "Status": "Error", "Message": "Did not delete any entities from Azure Table Storage", } resp_json = json.dumps(resp) return func.HttpResponse(body=resp_json, mimetype="application/json", status_code=400)
class AzureTableDatabase(object): def __init__(self): self.connection = TableService(account_name=storage_account, account_key=table_connection_string) self.table_name = table_name def _update_entity(self, record): record.LastModified = datetime.now() self.connection.update_entity(self.table_name, record) def create_table(self): if not self.connection.exists(self.table_name): self.connection.create_table(self.table_name) def raw_table(self, limit=100): """ Retrieve a list of rows in the table. """ calls = self.connection.query_entities(self.table_name, num_results=limit) return calls def list_calls(self, limit=100, select='PartitionKey'): """ Retrieve a set of records that need a phone call """ calls = self.connection.query_entities(self.table_name, num_results=limit, select=select) return [c.PartitionKey for c in calls] def reset_stale_calls(self, time_limit): """ Retrieve calls that are not done and whose last modified time was older than the limit. """ records = self.connection.query_entities(self.table_name, filter="LastModified lt datetime'{0}' and Status ne '{1}'".format(time_limit.date(), Statuses.extracting_done)) if not records.items: raise NoRecordsToProcessError() num_records = len(records.items) for record in records: if 'LastErrorStep' in record: record.Status = record.LastErrorStep del record.LastErrorStep record.Status = Statuses.reset_map.get(record.Status, record.Status) self._update_entity(record) return num_records def retrieve_next_record_for_call(self): """ Retrieve a set of records that need a phone call """ records = self.connection.query_entities(self.table_name, num_results=1, filter="Status eq '{0}'".format(Statuses.new)) if len(records.items) == 0: raise NoRecordsToProcessError() record = records.items[0] record.Status = Statuses.calling self._update_entity(record) return record.PartitionKey def set_error(self, partition_key, step): """ Reset a row from error state """ record = self.connection.get_entity(self.table_name, partition_key, partition_key) record.Status = Statuses.error record['LastErrorStep'] = step self._update_entity(record) def retrieve_next_record_for_transcribing(self): records = self.connection.query_entities( self.table_name, num_results=1, filter="Status eq '{0}'".format(Statuses.recording_ready), ) if not records.items: raise NoRecordsToProcessError() record = records.items[0] record.Status = Statuses.transcribing self._update_entity(record) return record.CallUploadUrl, record.PartitionKey def update_transcript(self, partition_key, transcript, status): record = self.connection.get_entity( self.table_name, partition_key, partition_key, ) if status == TranscriptionStatus.success: record.CallTranscript = transcript record.Status = Statuses.transcribing_done record.TranscribeTimestamp = datetime.now() self._update_entity(record) elif status == TranscriptionStatus.request_error: self.set_error(partition_key, Statuses.transcribing) else: record.Status = Statuses.transcribing_failed self._update_entity(record) def change_status(self, original_status, new_status): records = self.connection.query_entities( self.table_name, filter="Status eq '{0}'".format(original_status), ) if not records.items: return for record in records.items: record.Status = new_status self.connection.update_entity(self.table_name, record) def query(self, column, value, limit=1): records = self.connection.query_entities(self.table_name, num_results=limit, filter="{0} eq '{1}'".format(column, value)) return records def retrieve_next_record_for_extraction(self): records = self.connection.query_entities(self.table_name, num_results=1, filter="Status eq '{0}'".format(Statuses.transcribing_done)) if not records.items: raise NoRecordsToProcessError() record = records.items[0] record.Status = Statuses.extracting self._update_entity(record) return record.CallTranscript, record.PartitionKey def retrieve_next_record_for_extraction(self): records = self.connection.query_entities(self.table_name, num_results=1, filter="Status eq '{0}'".format(Statuses.transcribing_done)) if not records.items: raise NoRecordsToProcessError() record = records.items[0] record.Status = Statuses.extracting self.connection.update_entity(self.table_name, record) return record.CallTranscript, record.PartitionKey def update_location_date(self, case_number, city, location_confidence, state, zipcode, date): record = self.connection.get_entity(self.table_name, case_number, case_number) record.City = city record.LocationConfidence = location_confidence record.State = state record.Zipcode = zipcode record.CourtHearingDate = date record.Status = Statuses.extracting_done self.connection.update_entity(self.table_name, record) def upload_new_requests(self, request_ids): """ Upload new request ids to the database """ for request_id in request_ids: record = {'PartitionKey': request_id, 'RowKey': request_id, 'Status': Statuses.new, 'LastModified': datetime.now()} try: self.connection.insert_entity(self.table_name, record) except AzureConflictHttpError: pass # already exists. silently ignore. def update_call_id(self, alien_registration_id, call_id): record = self.connection.get_entity(self.table_name, alien_registration_id, alien_registration_id) record.CallID = call_id record.Status = Statuses.calling record.CallTimestamp = datetime.now() self._update_entity(record) def update_azure_path(self, alien_registration_id, azure_path): record = self.connection.get_entity(self.table_name, alien_registration_id, alien_registration_id) record.Status = Statuses.recording_ready record.CallUploadUrl = azure_path self._update_entity(record) def delete_ain(self, ain): return self.connection.delete_entity(self.table_name, ain, ain) def get_ain(self, ain): return self.connection.get_entity(self.table_name, ain, ain)
def test_list_releases(self, mock_now, mock_list_containers): # Mock time mock_now.return_value = pendulum.datetime(year=2020, month=5, day=1, hour=1) # Mock containers containers_in = make_containers() mock_list_containers.return_value = containers_in # Mock fetching of containers client = MagArchiverClient(account_name=self.account_name, account_key=self.account_key) containers = client.list_containers(last_modified_thresh=1) try: # Update releases based on containers num_updated, num_errors = client.update_releases(containers) self.assertEqual(num_updated, 3) self.assertEqual(num_errors, 0) # Two releases start_date = pendulum.datetime(year=2020, month=4, day=17) end_date = pendulum.datetime(year=2020, month=5, day=1) releases = client.list_releases(start_date=start_date, end_date=end_date, state=MagState.discovered, date_type=MagDateType.release) self.assertEqual(len(releases), 2) # 1 release start_date = pendulum.datetime(year=2020, month=4, day=17, minute=1) end_date = pendulum.datetime(year=2020, month=5, day=1) releases = client.list_releases(start_date=start_date, end_date=end_date, state=MagState.discovered, date_type=MagDateType.release) self.assertEqual(len(releases), 1) # Three releases start_date = pendulum.datetime(year=2020, month=4, day=17) end_date = pendulum.datetime(year=2020, month=5, day=1, minute=1) releases = client.list_releases(start_date=start_date, end_date=end_date, state=MagState.discovered, date_type=MagDateType.release, reverse=False) self.assertEqual(len(releases), 3) # Sorting reverse=False self.assertEqual(releases[0].row_key, '2020-04-17') self.assertEqual(releases[1].row_key, '2020-04-24') self.assertEqual(releases[2].row_key, '2020-05-01') # Sorting reverse=True releases = client.list_releases(start_date=start_date, end_date=end_date, state=MagState.discovered, date_type=MagDateType.release, reverse=True) self.assertEqual(releases[0].row_key, '2020-05-01') self.assertEqual(releases[1].row_key, '2020-04-24') self.assertEqual(releases[2].row_key, '2020-04-17') finally: # Clean up service = TableService(account_name=self.account_name, account_key=self.account_key) for container in containers: service.delete_entity(MagRelease.TABLE_NAME, 'mag', container.name.replace("mag-", ""))
class Db(object): ts = None def __init__(self): """Init connection with cosmosdb""" self.ts = TableService(account_name=ACCOUNT_NAME, account_key=ACCOUNT_KEY) def migrate(self): """ Create tabel if not exists""" if not self.ts.exists(USER_TABLE_NAME): self.ts.create_table(USER_TABLE_NAME) if not self.ts.exists(MESSAGE_TABLE_NAME): self.ts.create_table(MESSAGE_TABLE_NAME) def get_all_users(self): """select email from user""" return [ i['PartitionKey'] for i in self.ts.query_entities(USER_TABLE_NAME) ] def create_user(self, data=None): bjson = ep( EdmType.BINARY, dumps({ 'email': data['email'], 'password': sha224(bytes(data['password'], encoding='utf-8')).hexdigest(), 'full_name': data['full_name'] })) user = { 'PartitionKey': data['email'], 'RowKey': row_key, 'info': bjson } if (self.ts.insert_or_replace_entity(USER_TABLE_NAME, user)): return {'success': True} def delete_user(self, email=None): if (self.ts.delete_entity(USER_TABLE_NAME, email, row_key)): return {'success': True} def create_message(self, email=None, message=None): """ Create message in protobuf""" proto_message = message_pb2.Message() proto_message.title = message['title'] proto_message.content = message['content'] proto_message.magic_number = message['magic_number'] details = ep(EdmType.BINARY, str(proto_message)) bmessage = { 'PartitionKey': email, 'RowKey': row_key, 'details': details, } if (self.ts.insert_or_replace_entity(MESSAGE_TABLE_NAME, bmessage)): return {'success': True} def get_user(self, email=''): return self.ts.get_entity(USER_TABLE_NAME, email, row_key) def get_message(self, email=''): return self.ts.get_entity(MESSAGE_TABLE_NAME, email, row_key) def get_messages(self): messages = self.ts.query_entities(MESSAGE_TABLE_NAME) return list(messages)