Truncate Dynamodb Tables

lambda_function.py

import json
import boto3 as boto3
import os

dynamodb = boto3.resource('dynamodb')


def truncate_table(table_name):
    table = dynamodb.Table(table_name)
    scan_kwargs = {
        'ProjectionExpression': "#pk, #sk",
        'ExpressionAttributeNames': {"#pk": "<<PK>>", "#sk": "<<SK>>"}
    }

    done = False
    start_key = None
    print('--' * 20 + '>>')
    print(f'\nBegin Deletion Of Table: {table_name}')
    print('--' * 20 + '>>')
    iter_counter = 1
    while not done:
        if start_key:
            scan_kwargs['ExclusiveStartKey'] = start_key
        response = table.scan(**scan_kwargs)
        deletion_rec_count = len(response.get('Items', []))

        print(f'Table: {table_name} Deletion Iteration Number: {iter_counter}')
        print(f'Table: {table_name} Deleted Record count: {deletion_rec_count}')

        if deletion_rec_count > 0:
            params = {
                'Keys': response.get('Items', []),
                'TableName': table_name
            }
            batch_delete_item(params)

        start_key = response.get('LastEvaluatedKey', None)
        done = start_key is None
        iter_counter += 1
    print('<<' + '--' * 20)
    print(f'Completed Deletion of Table: {table_name}\n\n')


def batch_delete_item(params):
    table_name = params['TableName']
    deletion_keys = params['Keys']

    table = dynamodb.Table(table_name)

    try:
        with table.batch_writer() as writer:
            for key in deletion_keys:
                writer.delete_item(Key=key)
    except ClientError:
        raise


def lambda_handler(event, context):
    truncate_table_list = os.environ.get('TRUNCATE_TABLE_NAME_LIST_CSV').split(',')
    for table in truncate_table_list:
        truncate_table(table.strip())

    # TODO implement
    return {
        'statusCode': 200,
        'body': json.dumps(f'TRUNCATION COMPLETED FOR TABLES: {truncate_table_list}!')
    }