kron/persistence/dynamodb_request_log.py

#dynamodb access# 
from datetime import datetime
import boto3
from botocore.exceptions import ClientError

import logging
logger = logging.getLogger(__name__)

session = boto3.Session(
    region_name='us-west-2'
#    aws_access_key_id=AWS_ACCESS_KEY_ID,
#    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
)

logger.debug(f'region name {session.region_name}')
dynamodb = session.resource('dynamodb')

class RequestLog:
    """Encapsulates an Amazon DynamoDB table of request data."""
    def __init__(self, dyn_resource):
        """
        :param dyn_resource: A Boto3 DynamoDB resource.
        """
        self.dyn_resource = dyn_resource
        self.table = None

    def exists(self, table_name):
        """
        Determines whether a table exists. As a side effect, stores the table in
        a member variable.

        :param table_name: The name of the table to check.
        :return: True when the table exists; otherwise, False.
        """
        try:
            table = self.dyn_resource.Table(table_name)
            table.load()
            exists = True
        except ClientError as err:
            if err.response['Error']['Code'] == 'ResourceNotFoundException':
                exists = False
            else:
                logger.error(
                    "Couldn't check for existence of %s. Here's why: %s: %s",
                    table_name,
                    err.response['Error']['Code'], err.response['Error']['Message'])
                raise
        else:
            self.table = table
        return exists
    
    def create_table(self, table_name):
        """
        Creates an Amazon DynamoDB table that can be used to store request data.
        The table uses the release year of the movie as the partition key and the
        title as the sort key.

        :param table_name: The name of the table to create.
        :return: The newly created table.
        """
        try:
            self.table = self.dyn_resource.create_table(
                TableName=table_name,
                KeySchema=[
                    {'AttributeName': 'model', 'KeyType': 'HASH'},  # Partition key
                    {'AttributeName': 'timestamp', 'KeyType': 'RANGE'}  # Sort key
                ],
                AttributeDefinitions=[
                    {'AttributeName': 'model', 'AttributeType': 'S'},
                    {'AttributeName': 'timestamp', 'AttributeType': 'S'},
 #                   {'AttributeName': 'request', 'AttributeType': 'S'},
 #                   {'AttributeName': 'response', 'AttributeType': 'S'}
                ],
                ProvisionedThroughput={'ReadCapacityUnits': 10, 'WriteCapacityUnits': 10})
            self.table.wait_until_exists()
        except ClientError as err:
            logger.error(
                "Couldn't create table %s. Here's why: %s: %s", table_name,
                err.response['Error']['Code'], err.response['Error']['Message'])
            raise
        else:
            return self.table
        
    def log_request(self, req_timestamp_str, model, request_str, response_str, rating = 0):
        """
        Log a request to the table.

        # TODO
        :param title: The title of the movie.
        :param year: The release year of the movie.
        :param plot: The plot summary of the movie.
        :param rating: The quality rating of the movie.
        """
        try:
            self.table.put_item(
                Item={
                    'timestamp': req_timestamp_str,
                    'model': model,
                    'request': request_str,
                    'response': response_str,
                    'rating': rating,
                }
            )
        except ClientError as err:
            logger.error(
                "Couldn't add request log %s to table %s. Here's why: %s: %s",
                model, self.table.name,
                err.response['Error']['Code'], err.response['Error']['Message'])
            raise    
    
    def add_request_log_entry(self, query_model, req, resp, rating=0):
        """
            Logs the cuurent model, req and response
        """
        today = datetime.now()
        # Get current ISO 8601 datetime in string format
        iso_date = today.isoformat()        
        self.log_request(iso_date, query_model, req, resp, rating)

table_name = 'hf-spaces-request-log'

def get_request_log():
    request_log = RequestLog(dynamodb)
    request_log_exists = request_log.exists(table_name)
    if not request_log_exists:
        print(f"\nCreating table {table_name}...")
        request_log.create_table(table_name)
        print(f"\nCreated table {request_log.table.name}.")
    return request_log

#def add_request_log_entry(request_log, query_model, req, resp, rating=0):
#    today = datetime.now()
#    # Get current ISO 8601 datetime in string format
#    iso_date = today.isoformat()        
#    request_log.log_request(iso_date, query_model, req, resp, rating)