class MongoDataClient(DataClient): """ Data client for the mongodb nosql document data store. """ def __init__(self, **kwargs): """ Initialization for mongodb with support for TDD. """ super().__init__(**kwargs) self.db_name = Settings.MONGO_DB_NAME collection_name = kwargs.get(Constants.TABLE, Settings.SERVICE_NAME) self.host = kwargs.get(Constants.HOST, Settings.MONGO_URL) self.port = kwargs.get(Constants.PORT, Settings.MONGO_PORT) self.client = kwargs.get(Constants.CLIENT, MongoClient(self.host, self.port, serverSelectionTimeoutMS=20)) self.collection = None self.logger = Logger() self.database = self.client[self.db_name] self._select_collection(collection_name) def get_health(self): try: return self.database.command(Constants.MONGO_STATS_COMMAND) except ServerSelectionTimeoutError: return {Constants.HEALTH: Constants.MONGO_CONNECTION_HEALTH_ERROR.format(host=self.host, port=self.port)} @log_exception def _select_collection(self, collection_name): collection = self.get_collection_name(collection_name) self.collection = self.client[self.db_name][collection] @log_exception def rename_table(self, old_name, new_name): self.logger.debug(f'Rename table called on data client. old: {old_name} new: {new_name}') self.database[old_name].rename(new_name, dropTarget=True) @log_exception def table_exists(self, table_name): self.logger.debug(f'Existence check on table: {table_name}') return table_name in self.database.collection_names() @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_create')) def create(self, model): if isinstance(model, BaseModel): model = model.__dict__ result = self.collection.insert_one(model) return result.inserted_id is not None @log_exception def restore(self, data): self.logger.debug(f'Restoring data...') result = self.collection.insert_many(data) return result.inserted_ids is not None @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_get')) def get(self, key): self.logger.debug(f'Get item for id: {key}') return self.collection.find_one({Constants.DB_ID_KEY: ObjectId(key)}) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_update')) def update(self, key, model): self.logger.debug(f'Update item for id: {key}') if isinstance(model, BaseModel): model = model.__dict__ model.pop(Constants.ID) return self.collection.find_one_and_replace({Constants.DB_ID_KEY: ObjectId(key)}, model, upsert=True) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_get_specific')) def get_specific(self, id_value, collection): self.logger.debug(f'Get specific item for id: {id_value} from {collection}') self.collection = self.client[self.db_name][collection] return self.collection.find_one({Constants.DB_ID_KEY: id_value}) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_get_all')) def get_all(self): self.logger.debug(f'Get all called') return self.convert_cursor_to_array(self.collection.find()) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_find')) def find(self, key, value): self.logger.debug(f'Find called key: {key} val: {value}') return self.collection.find_one({key: value}) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_count')) def count(self): self.logger.debug('Count called') return self.collection.find().count(True) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_purge')) def purge(self): self.logger.debug('Purge called') return self.collection.drop() @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix(Settings.SERVICE_NAME, 'mongo_read')) def read(self, offset, limit, query): self.logger.debug(f'Read called offset: {offset} limit: {limit} query: {query}') count = self.collection.find(query).count() cursor = self.collection.find(query).skip(offset).limit(limit) return self.convert_cursor_to_array(cursor), offset, limit, count @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'mongo_delete')) def delete(self, key): self.logger.debug(f'Delete called id: {key}') result = self.collection.delete_one({Constants.DB_ID_KEY: ObjectId(key)}) return result.deleted_count == 1 @log_exception def delete_db(self): self.logger.debug('Drop db called.') self.client.drop_database(self.db_name) @log_exception def get_server_info(self): self.logger.debug('Queries DB for server info') return self.client.server_info() @staticmethod def convert_cursor_to_array(cursor): documents = [] if cursor is not None: for document in cursor: documents.append(document) return documents @staticmethod def get_collection_name(model): if model: if isinstance(model, str): return model return model.__name__ return 'not_set' def get_status(self): pass def filter(self, **kwargs): offset, limit, query = get_query(**kwargs) return self.read(offset, limit, query)
class DynamoDataClient(DataClient): def __init__(self, **kwargs): super().__init__(**kwargs) self.logger = Logger() self.table_name = kwargs.get(Constants.TABLE, Settings.SERVICE_NAME) self.client = kwargs.get(Constants.CLIENT, get_default_dynamo_client()) self.resource = kwargs.get(Constants.RESOURCE, get_default_dynamo_resource()) self._initialize_tables() @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_create')) def create(self, model): model = self.get_converted_model(model) result = self.resource.Table(self.table_name).put_item(Item=model) return self.check_status(result) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_update')) def update(self, key, model): model.id = key return self.create(model) def get_converted_model(self, model, nest=False): empty_keys = list() if not isinstance(model, dict): model = model.__dict__ for key, value in model.items(): if isinstance(value, float): model[key] = Decimal(str(value)) continue if not value: if key not in empty_keys: empty_keys.append(key) continue if isinstance(value, dict): model[key] = self.get_converted_model(value, True) if not model[key]: if key not in empty_keys: empty_keys.append(key) for key in empty_keys: model.pop(key) if not nest: if Constants.ID not in model.keys(): model[Constants.ID] = str(uuid4()) return model @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_get')) def get(self, key): self.logger.debug(f'Get item on dynamo for id: {key}') self.create_table(self.table_name) return self.resource.Table(self.table_name).get_item(Key={ Constants.ID: key }).get(Constants.DYNAMO_ITEM_KEY, None) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_get_specific')) def get_specific(self, id_value, table_name): self.logger.debug( f'Get specific item on dynamo for id: {id_value} from {table_name}' ) self.create_table(self.table_name) return self.resource.Table( self.table_name).get_item(Key={ Constants.ID: id_value }).get(Constants.DYNAMO_ITEM_KEY, None) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_find')) def find(self, key, value): self.logger.debug(f'Find called on dynamo key: {key} val: {value}') all_items = self.get_all() return self.find_item_in_list(all_items, **{key: value}) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_get_all')) def get_all(self): self.logger.debug(f'Get all called on dynamo') self.create_table(self.table_name) return self.resource.Table(self.table_name).scan().get( Constants.DYNAMO_ITEMS_KEY, []) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_count')) def count(self): return self.resource.Table(self.table_name).scan().get( Constants.DYNAMO_COUNT_KEY, 0) @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_purge')) def purge(self): self.logger.debug('Purge called on dynamo...') success = True items = self.get_all() for item in items: success = success and self.delete(item.get(Constants.ID)) return success @log_exception @monitor(monitor_name=MonitoringProvider.get_monitor_prefix( Settings.SERVICE_NAME, 'dynamo_delete')) def delete(self, key): self.logger.debug(f'Delete on dynamo called id: {key}') result = self.resource.Table( self.table_name).delete_item(Key={Constants.ID: key}) return self.check_status(result) @staticmethod def check_status(result): return result.get(Constants.DYNAMO_RESPONSE_METADATA, {}).get( Constants.DYNAMO_SUCCESS_STATUS, Constants.BAD_REQUEST_CODE) == Constants.SUCCESS_RESPONSE_CODE @log_exception def table_exists(self, table_name): try: return self.resource.Table(table_name).table_status == 'ACTIVE' except Exception: # pylint: disable=broad-except return False @log_exception def delete_table(self, table_name): try: return self.check_status( self.client.delete_table(TableName=table_name)) except Exception: # pylint: disable=broad-except return False @log_exception def filter(self, **kwargs): if not kwargs: items = self.get_all() return items, 0, len(items), len(items) offset, limit, query = get_query(**kwargs) all_items = self.get_all() return self.find_items_in_list(all_items, offset, limit, **query) @log_exception def delete_db(self): self.logger.debug('Drop db on dynamo called.') return self.delete_table(self.table_name) def create_table(self, table_name): try: if not self.table_exists(table_name): self.resource.create_table( AttributeDefinitions=[Constants.DYNAMO_ID_TABLE_ATTR], KeySchema=[Constants.DYNAMO_ID_TABLE_KEY], ProvisionedThroughput=Constants.DEFAULT_DYNAMO_THROUGHPUT, TableName=table_name) return True except Exception as ex: # pylint: disable=broad-except self.logger.info(f'Error creating dynamo table. {ex}', exc_info=True) return False @staticmethod def _get_dictionary_set(items_dictionary): converted_list = list() for key, value in items_dictionary.items(): converted_list.append(f'{key}|{value}') return set(converted_list) def find_item_in_list(self, all_items, **kwargs): args_set = self._get_dictionary_set(kwargs) match_list = list() for item in all_items: item_set = self._get_dictionary_set(item) if len(item_set.intersection(args_set)) == len(args_set): match_list.append(item) return match_list def find_items_in_list(self, all_items, offset, limit, **kwargs): self.logger.info(f'Find called on dynamo with {offset}, {limit}') match_list = self.find_item_in_list(all_items, **kwargs) return match_list, offset, limit, len(match_list) def _tag_table(self, table_arn): try: self.client.tag_resource(ResourceArn=table_arn, Tags=Constants.DYNAMO_TAGS) except Exception: # pylint: disable=broad-except pass def _initialize_tables(self): if self.create_table(self.table_name): self._tag_table(self.resource.Table(self.table_name).table_arn) def get_status(self): return self.table_exists(self.table_name)