src/aws/s3.py

from io import StringIO

import boto3
import pandas as pd
from botocore.exceptions import NoCredentialsError, PartialCredentialsError
from mypy_boto3_s3.client import S3Client

from src.utils import setup_logger

logger = setup_logger(__name__)

trial_bucket_name = "pharma-doc-trials"
base_change_request_history_file_path_in_trial_bucket = (
    "bankyo/change_request_history.csv"
)
csv_extension = ".csv"
err_msg_file_is_not_csv = "指定したファイルはcsvではありません。"


def get_client() -> S3Client:
    return boto3.client("s3")


def get_csv_as_pd_dataframe_from_s3(
    client: S3Client, bucket_name: str, file_name: str
) -> pd.DataFrame:
    if not file_name.lower().endswith(csv_extension):
        raise ValueError(err_msg_file_is_not_csv)
    try:
        obj = client.get_object(Bucket=bucket_name, Key=file_name)
        csv_content = obj["Body"].read().decode("utf-8")
        data = pd.read_csv(StringIO(csv_content))
        return data
    except client.exceptions.NoSuchKey as e:
        logger.info(f'No such key: {e.response["Error"]["Key"]}')
        raise e
    except NoCredentialsError as e:
        logger.info("Error: No AWS credentials found.")
        raise e
    except PartialCredentialsError as e:
        logger.info("Error: Incomplete AWS credentials found.")
        raise e
    except client.exceptions.ClientError as e:
        logger.info(f"Unexpected error: {e}")
        raise e


def save_pd_dataframe_as_csv_to_s3(
    client: S3Client, data: pd.DataFrame, bucket_name: str, file_name: str
) -> None:
    if not file_name.lower().endswith(csv_extension):
        raise ValueError(err_msg_file_is_not_csv)
    csv_buffer = StringIO()
    data.to_csv(csv_buffer, index=False)
    try:
        client.put_object(Bucket=bucket_name, Key=file_name, Body=csv_buffer.getvalue())
        logger.info(f"Successfully saved {file_name} to {bucket_name}.")
    except NoCredentialsError as e:
        logger.info("Error: No AWS credentials found.")
        raise e
    except PartialCredentialsError as e:
        logger.info("Error: Incomplete AWS credentials found.")
        raise e
    except client.exceptions.ClientError as e:
        logger.info(f"Unexpected error: {e}")
        raise e