[python][treasury-report] add force-regenerate parameter to allow for deletion of the existing treasury file before generation

api/src/graphql/uploads.sdl.ts

@@ -43,6 +43,7 @@ export const schema = gql`
     updateUpload(id: Int!, input: UpdateUploadInput!): Upload! @requireAuth
     deleteUpload(id: Int!): Upload! @requireAuth
     downloadUploadFile(id: Int!): String! @requireAuth
+    generateTreasuryReport(regenerate: Boolean!): Boolean! @requireAuth
     sendTreasuryReport: Boolean! @requireAuth
   }
 `

api/src/services/uploads/uploads.test.ts

@@ -22,6 +22,7 @@ import {
   Upload as UploadRelationResolver,
   getUploadsByExpenditureCategory,
   getValidUploadsInCurrentPeriod,
+  generateTreasuryReport,
   sendTreasuryReport,
   SubrecipientLambdaPayload,
   ProjectLambdaPayload,
@@ -380,6 +381,137 @@ describe('treasury report', () => {
     jest.resetAllMocks()
     process.env.TREASURY_STEP_FUNCTION_ARN = 'test-arn'
   })
+  scenario(
+    'successfully force regenerate a treasury report',
+    async (scenario: StandardScenario) => {
+      mockCurrentUser(scenario.user.one)
+      const mockOrganization = scenario.organization.one
+      const mockReportingPeriod = scenario.reportingPeriod.one
+      const mockUpload = scenario.upload.two
+      const mockUser = scenario.user.one
+
+      const projectPayload: ProjectLambdaPayload = {
+        '1A': {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+          user: {
+            email: mockUser.email,
+            id: mockUser.id,
+          },
+          outputTemplateId: mockReportingPeriod.outputTemplateId,
+          uploadsToAdd: {
+            [mockUpload.agencyId]: {
+              objectKey: `uploads/${mockOrganization.id}/${mockUpload.agencyId}/${mockReportingPeriod.id}/${mockUpload.id}/${mockUpload.filename}`,
+              createdAt: mockUpload.createdAt,
+              filename: mockUpload.filename,
+            },
+          },
+          uploadsToRemove: {},
+          forceRegenerate: true,
+          ProjectType: '1A',
+        },
+        '1B': {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+          user: {
+            email: mockUser.email,
+            id: mockUser.id,
+          },
+          outputTemplateId: mockReportingPeriod.outputTemplateId,
+          uploadsToAdd: {},
+          uploadsToRemove: {},
+          forceRegenerate: true,
+          ProjectType: '1B',
+        },
+        '1C': {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+          user: {
+            email: mockUser.email,
+            id: mockUser.id,
+          },
+          outputTemplateId: mockReportingPeriod.outputTemplateId,
+          uploadsToAdd: {},
+          uploadsToRemove: {},
+          forceRegenerate: true,
+          ProjectType: '1C',
+        },
+      }
+      const subrecipientPayload: SubrecipientLambdaPayload = {
+        Subrecipient: {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+          user: {
+            email: mockUser.email,
+            id: mockUser.id,
+          },
+          outputTemplateId: mockReportingPeriod.outputTemplateId,
+        },
+      }
+
+      const zipPayload: CreateArchiveLambdaPayload = {
+        zip: {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+        },
+      }
+      const emailPayload: EmailLambdaPayload = {
+        email: {
+          organization: {
+            id: mockOrganization.id,
+            preferences: {
+              current_reporting_period_id: mockReportingPeriod.id,
+            },
+          },
+          user: {
+            email: mockUser.email,
+            id: mockUser.id,
+          },
+        },
+      }
+
+      const input = JSON.stringify({
+        '1A': {},
+        '1B': {},
+        '1C': {},
+        Subrecipient: {},
+        zip: {},
+        email: {},
+        ...projectPayload,
+        ...subrecipientPayload,
+        ...zipPayload,
+        ...emailPayload,
+      })
+      const result = await generateTreasuryReport({ regenerate: true })
+
+      expect(result).toBe(true)
+      expect(startStepFunctionExecution).toHaveBeenCalledWith(
+        'test-arn',
+        `Force-kick-off-00000000-0000-0000-0000-000000000000`,
+        input
+      )
+    }
+  )
 
   scenario(
     'successfully sends a treasury report',
@@ -411,6 +543,7 @@ describe('treasury report', () => {
             },
           },
           uploadsToRemove: {},
+          forceRegenerate: false,
           ProjectType: '1A',
         },
         '1B': {
@@ -427,6 +560,7 @@ describe('treasury report', () => {
           outputTemplateId: mockReportingPeriod.outputTemplateId,
           uploadsToAdd: {},
           uploadsToRemove: {},
+          forceRegenerate: false,
           ProjectType: '1B',
         },
         '1C': {
@@ -443,6 +577,7 @@ describe('treasury report', () => {
           outputTemplateId: mockReportingPeriod.outputTemplateId,
           uploadsToAdd: {},
           uploadsToRemove: {},
+          forceRegenerate: false,
           ProjectType: '1C',
         },
       }

api/src/services/uploads/uploads.ts

@@ -186,6 +186,7 @@ type UploadInfoForProject = {
   ProjectType: string
   uploadsToAdd: Partial<Record<AgencyId, UploadPayload>>
   uploadsToRemove: Partial<Record<AgencyId, UploadPayload>>
+  forceRegenerate: boolean
 }
 type InfoForSubrecipient = {
   organization: OrganizationObj
@@ -216,7 +217,8 @@ export type EmailLambdaPayload = Record<'email', InfoForEmail>
 
 export const getUploadsByExpenditureCategory = async (
   organization: Organization,
-  reportingPeriod: ReportingPeriod
+  reportingPeriod: ReportingPeriod,
+  regenerate = false
 ): Promise<ProjectLambdaPayload> => {
   const validUploadsInPeriod: UploadsWithValidationsAndExpenditureCategory[] =
     await getValidUploadsInCurrentPeriod(organization, reportingPeriod)
@@ -236,6 +238,7 @@ export const getUploadsByExpenditureCategory = async (
     outputTemplateId: reportingPeriod.outputTemplateId,
     uploadsToAdd: {},
     uploadsToRemove: {},
+    forceRegenerate: regenerate,
   }
 
   const uploadsByEC: ProjectLambdaPayload = {
@@ -364,6 +367,57 @@ export const getEmailLambdaPayload = async (
   }
 }
 
+export const generateTreasuryReport: MutationResolvers['generateTreasuryReport'] =
+  async ({ regenerate }) => {
+    try {
+      const organization = await db.organization.findFirst({
+        where: { id: context.currentUser.agency.organizationId },
+      })
+      const reportingPeriod = await db.reportingPeriod.findFirst({
+        where: { id: organization.preferences['current_reporting_period_id'] },
+      })
+      const projectLambdaPayload: ProjectLambdaPayload =
+        await getUploadsByExpenditureCategory(
+          organization,
+          reportingPeriod,
+          regenerate
+        )
+      const subrecipientLambdaPayload: SubrecipientLambdaPayload =
+        await getSubrecipientLambdaPayload(
+          organization,
+          context.currentUser,
+          reportingPeriod
+        )
+      const createArchiveLambdaPayload: CreateArchiveLambdaPayload =
+        await getCreateArchiveLambdaPayload(organization)
+
+      const emailLambdaPayload: EmailLambdaPayload =
+        await getEmailLambdaPayload(organization, context.currentUser)
+
+      const input = {
+        '1A': {},
+        '1B': {},
+        '1C': {},
+        Subrecipient: {},
+        zip: {},
+        email: {},
+        ...projectLambdaPayload,
+        ...subrecipientLambdaPayload,
+        ...createArchiveLambdaPayload,
+        ...emailLambdaPayload,
+      }
+      await startStepFunctionExecution(
+        process.env.TREASURY_STEP_FUNCTION_ARN,
+        `Force-kick-off-${uuidv4()}`,
+        JSON.stringify(input)
+      )
+      return true
+    } catch (error) {
+      logger.error(error, 'Error sending Treasury Report')
+      throw new RedwoodError(error.message)
+    }
+  }
+
 export const sendTreasuryReport: MutationResolvers['sendTreasuryReport'] =
   async () => {
     try {

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_files(
+            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,59 @@ def process_event(
     return {"statusCode": 200, "body": "Success"}
 
 
+def _try_delete(
+    s3_client: S3Client,
+    bucket: str,
+    key: str,
+    logger: structlog.stdlib.BoundLogger,
+):
+    """Tries to delete files from s3.
+
+    No error is thrown if the file is not found."""
+    try:
+        delete_file_from_s3(
+            client=s3_client,
+            bucket=bucket,
+            key=key,
+        )
+    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")
+
+
+def delete_output_files(
+    s3_client: S3Client,
+    organization: OrganizationObj,
+    project_use_code: str,
+    logger: structlog.stdlib.BoundLogger,
+):
+    """Deletes the output files and the generated report."""
+    for output_file_type in [
+        OutputFileType.XLSX,
+        OutputFileType.CSV,
+        OutputFileType.JSON,
+    ]:
+        _try_delete(
+            s3_client=s3_client,
+            bucket=os.environ["REPORTING_DATA_BUCKET_NAME"],
+            key=get_generated_output_file_key(
+                file_type=output_file_type,
+                project=project_use_code,
+                organization=organization,
+            ),
+            logger=logger,
+        )
+
+    report_key = f"treasuryreports/{organization.id}/{organization.preferences.current_reporting_period_id}/report.zip"
+    _try_delete(
+        s3_client=s3_client,
+        bucket=os.environ["REPORTING_DATA_BUCKET_NAME"],
+        key=report_key,
+        logger=logger,
+    )
+
+
 def download_output_file(
     s3_client: S3Client,
     output_file: IO[bytes],

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(delete={"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,