diff --git a/python/src/functions/generate_treasury_report.py b/python/src/functions/generate_treasury_report.py index 1cd4dc0b..b941f050 100644 --- a/python/src/functions/generate_treasury_report.py +++ b/python/src/functions/generate_treasury_report.py @@ -14,7 +14,11 @@ from pydantic import BaseModel from src.lib.logging import get_logger, reset_contextvars -from src.lib.s3_helper import download_s3_object, upload_generated_file_to_s3 +from src.lib.s3_helper import ( + delete_file_from_s3, + download_s3_object, + upload_generated_file_to_s3, +) from src.lib.treasury_generation_common import ( OrganizationObj, OutputFileType, @@ -54,6 +58,7 @@ class ProjectLambdaPayload(BaseModel): ProjectType: str uploadsToAdd: Dict[AgencyId, UploadObj] uploadsToRemove: Dict[AgencyId, UploadObj] + forceRegenerate: bool = False @reset_contextvars @@ -132,6 +137,16 @@ def process_event( organization = payload.organization + # If we need to force-regenerate the treasury report, we delete the treasury + # report from s3 so that it starts with the template file. + if payload.forceRegenerate: + delete_output_file( + s3_client=s3_client, + organization=organization, + project_use_code=project_use_code, + logger=logger, + ) + # If the treasury report file exists, download it and store the data # If it doesn't exist, download the output template @@ -267,6 +282,29 @@ def process_event( return {"statusCode": 200, "body": "Success"} +def delete_output_file( + s3_client: S3Client, + organization: OrganizationObj, + project_use_code: str, + logger: structlog.stdlib.BoundLogger, +): + try: + delete_file_from_s3( + client=s3_client, + bucket=os.environ["REPORTING_DATA_BUCKET_NAME"], + key=get_generated_output_file_key( + file_type=OutputFileType.XLSX, + project=project_use_code, + organization=organization, + ), + ) + except ClientError as e: + error = e.response.get("Error") or {} + if error.get("Code") == "404": + logger.exception("Expected to find an existing treasury output report") + raise + + def download_output_file( s3_client: S3Client, output_file: IO[bytes], diff --git a/python/src/lib/s3_helper.py b/python/src/lib/s3_helper.py index bca00002..ac8dc95d 100644 --- a/python/src/lib/s3_helper.py +++ b/python/src/lib/s3_helper.py @@ -29,6 +29,31 @@ def download_s3_object( raise +def delete_file_from_s3(client: S3Client, bucket: str, key: str) -> None: + """Deletes file from S3. + + Args: + client: Client facilitating delete from S3 + bucket: bucket file should be deleted from + key: S3 key for file + """ + logger = get_logger() + + logger = logger.bind(upload={"s3": {"bucket": bucket, "key": key}}) + try: + client.delete_object( + Bucket=bucket, + Key=unquote_plus(key), + ) + except: + logger.exception( + f"failed to delete file from S3, bucket {bucket} and key {key}" + ) + raise + + logger.info("successfully deleted file from s3") + + def upload_generated_file_to_s3( client: S3Client, bucket: str, @@ -45,7 +70,7 @@ def upload_generated_file_to_s3( """ logger = get_logger() - logger = logger.bind(upload={"s3": {"bucket": bucket, "key": key}}) + logger = logger.bind(delete={"s3": {"bucket": bucket, "key": key}}) try: file.seek(0) client.put_object(