storage.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Purpose

Shows how to use the AWS SDK for Python (Boto3) with Amazon DynamoDB to store work items
in a table.
"""

import logging
from uuid import uuid4
from boto3.dynamodb.conditions import Attr
from botocore.exceptions import ClientError


logger = logging.getLogger(__name__)


class StorageError(Exception):
    pass


class Storage:
    """
    Encapsulates work item data in a DynamoDB table.
    """
    def __init__(self, table):
        """
        :param table: A Boto3 DynamoDB Table object that represents an existing DynamoDB
                      table. This object is a high-level object that wraps low-level
                      DynamoDB service actions.
        """
        self.table = table

    def get_work_items(self, archived=None):
        """
        Gets work items currently stored in the table.

        :param archived: When specified, only archived or non-archived work items are
                         returned. Otherwise, all work items are returned.
        :return: A list of work items currently stored in the table.
        """
        try:
            if archived is None:
                work_items = self.table.scan().get('Items', [])
            else:
                work_items = self.table.scan(
                    FilterExpression=Attr('archived').eq(archived)).get('Items', [])
        except ClientError as err:
            logger.exception(
                "Couldn't get items from table %s with archived %s.", self.table.name,
                archived)
            raise StorageError(err)
        else:
            return work_items

    def get_work_item(self, iditem):
        """
        Gets a single work item from the table.

        :param iditem: The ID of the item to retrieve.
        :return: The requested item.
        """
        try:
            item = self.table.get_item(Key={'iditem': iditem}).get('Item', {})
        except ClientError as err:
            logger.exception(
                "Couldn't get item %s from table %s.", iditem, self.table.name)
            raise StorageError(err)
        else:
            return item

    def add_or_update_work_item(self, item):
        """
        Adds or updates an item in the table. When the item contains an iditem field, it
        is updated. Otherwise, it is added. When an item is added, a UUID is generated
        as its ID.

        :param item: The item to add or update.
        :return: The ID of the item.
        """
        try:
            if item.get('iditem') is None:
                item['iditem'] = str(uuid4())
                self.table.put_item(Item=item)
            else:
                attrs = {
                    key: {'Value': val, 'Action': 'PUT'}
                    for key, val in item.items() if key != 'iditem'
                }
                self.table.update_item(
                    Key={'iditem': item['iditem']},
                    AttributeUpdates=attrs)
        except ClientError as err:
            logger.exception("Couldn't add or update item %s in table %s.")
            raise StorageError(err)
        return item['iditem']