covid_data.py - AWS Code Sample

covid_data.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) to write and retrieve Amazon DynamoDB data as part of a REST API. This file is uploaded to AWS Lambda as part of the serverless deployment package created by AWS Chalice. """ import datetime import os import random import boto3 from boto3.dynamodb.conditions import Key class Storage: """ Handles basic storage functions, backed by an Amazon DynamoDB table. """ STATES = { 'Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut', 'Delaware', 'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois', 'Indiana', 'Iowa', 'Kansas', 'Kentucky', 'Louisiana', 'Maine', 'Maryland', 'Massachusetts', 'Michigan', 'Minnesota', 'Mississippi', 'Missouri', 'Montana', 'Nebraska', 'Nevada', 'New Hampshire', 'New Jersey', 'New Mexico', 'New York', 'North Carolina', 'North Dakota', 'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island', 'South Carolina', 'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington', 'West Virginia', 'Wisconsin', 'Wyoming' } def __init__(self, table): self._table = table @classmethod def from_env(cls): """ Creates a Storage object that contains a table identified by the TABLE_NAME environment variable. :return: The newly created Storage object. """ table_name = os.environ.get('TABLE_NAME', '') table = boto3.resource('dynamodb').Table(table_name) return cls(table) @staticmethod def _generate_random_data(state): """ Generates some random data for the demo. :param state: The state for which to create the data. :return: The newly created data. """ return { 'state': state, 'date': datetime.date.today().isoformat(), 'cases': random.randint(1, 1000), 'deaths': random.randint(1, 100) } def get_state_data(self, state): """ Gets the data records for the specified state. If there are no records, a new one is generated with random values for today's date and stored in the table before it is returned. :param state: The state to retrieve. :return: The retrieved data. """ response = self._table.query( KeyConditionExpression=Key('state').eq(state) ) items = response.get('Items', []) if len(items) == 0: items.append(self._generate_random_data(state)) self._table.put_item(Item=items[0]) return items def put_state_data(self, state, state_data): """ Puts data for a state into the table. :param state: The state for which to store the data. :param state_data: The data record to store. """ self._table.put_item(Item=state_data) def delete_state_data(self, state): """ Deletes all records for a state from the table. :param state: The state to delete. """ response = self._table.query( KeyConditionExpression=Key('state').eq(state) ) items = response.get('Items', []) with self._table.batch_writer() as batch: for item in items: batch.delete_item(Key={'state': item['state'], 'date': item['date']}) def post_state_data(self, state, state_data): """ Puts data for a state into the table. :param state: The state for which to store the data. :param state_data: The data record to store. """ self._table.put_item(Item=state_data) def get_state_date_data(self, state, date): """ Gets a single record for the specified state and date. :param state: The state of the record to retrieve. :param date: The date of the record to retrieve. :return: The retrieved record, or None if no record exists. """ response = self._table.get_item(Key={'state': state, 'date': date}) item = response.get('Item', None) return item def delete_state_date_data(self, state, date): """ Deletes the record for the specified state and date. :param state: The state of the record to remove. :param date: The date of the record to remove. """ self._table.delete_item(Key={'state': state, 'date': date})