"""Common dynamoDB methods.""" import datetime import time import os import boto3 from boto3.dynamodb import conditions from botocore import config as botocore_config DDB_WRITE_MAX_ATTEMPTS = int(os.environ.get('DDB_WRITE_MAX_ATTEMPTS', 10)) def get_dynamodb(max_attempts=None): """Get Boto DynamoDB resource. Args: max_attempts (int): optional max retry attempts Returns: boto3.resources.factory.dynamodb.ServiceResource: DynamoDB resource """ if max_attempts is None: max_attempts = DDB_WRITE_MAX_ATTEMPTS config_instance = botocore_config.Config( retries={'max_attempts': max_attempts}) return boto3.resource('dynamodb', config=config_instance) def get_dynamodb_table(table_name): """Get the DynamoDB table. Args: table_name (str): DynamoDB table name Returns: boto3.resources.factory.dynamodb.Table: DynamoDB table """ dynamodb = get_dynamodb() return dynamodb.Table(table_name) def batch_write_data(table_name, items): """Batch write items into DynamoDB table. Args: table_name (str): Table name items (list): List with item dictionaries """ table = get_dynamodb_table(table_name) with table.batch_writer() as batch_writer: for item in items: batch_writer.put_item( Item=item ) def batch_delete_data(table_name, keys): """Delete items from DynamoDB table in batch. Args: table_name (str): table name keys (list): keys to delete """ table = get_dynamodb_table(table_name) with table.batch_writer() as batch_writer: for key in keys: batch_writer.delete_item(Key=key) def put_item(table_name, item): """Put an item into DynamoDB table. Args: table_name (str): name of the table item (dict): item to put """ table = get_dynamodb_table(table_name) table.put_item(Item=item, ReturnValues='NONE') def get_ttl_value_for_item(time_delta): """Calculate the TTL value for the item. Args: time_delta (datetime.timedelta): expected time to live Returns: int: TTL in seconds from epoch """ ttl_datetime = datetime.datetime.now() + time_delta return int(ttl_datetime.timestamp()) def get_scan_results(table, scan_kwargs): """Yield scan responses with respect to LastEvaluatedKey. Args: table (boto3.resources.factory.dynamodb.Table): DynamoDB table to scan scan_kwargs (dict): scan key word arguments Returns: dict: individual scan response """ key = True while key: resp = table.scan(**scan_kwargs) key = resp.get('LastEvaluatedKey') if key: scan_kwargs['ExclusiveStartKey'] = key start = time.time() yield resp # sleep only if less then a second has passed if time.time() - start < 1: time.sleep(1) def update_existing_item(table, update_data): """Update existin DynamoDB record. Args: table (boto3.resources.factory.dynamodb.Table): DynamoDB table update update_data (dict): update parameters, expected to have: Key, UpdateExpression and ExpressionAttributeValues Returns: dict: DynamoDB response dict """ pk = update_data['Key']['label_type_id_period_id'] sort_key = update_data['Key']['file_name'] condition = ( conditions.Key('label_type_id_period_id').eq(pk) & conditions.Key('file_name').eq(sort_key) ) update_kwargs = { 'ConditionExpression': condition, } update_data.update(update_kwargs) resp = table.update_item(**update_data) return resp