add more typing, add tag to file upload, and rename form generation service

This commit is contained in:
brmnjsh 2023-09-03 21:38:46 +00:00
parent 386ab25515
commit d90f51faca
3 changed files with 13 additions and 23 deletions

View File

@ -1,5 +1,5 @@
import boto3
import json
import io
from lib.application_configs import ApplicationConfigs
@ -16,21 +16,21 @@ else:
s3 = session.client('s3', region_name=ApplicationConfigs.region_name)
sqs = session.client('sqs', region_name=ApplicationConfigs.region_name)
def get_key_from_message(body):
def get_key_from_message(body: dict):
return body['Records'][0]['s3']['object']['key']
def get_bucket_from_message(body):
def get_bucket_from_message(body: dict):
return body['Records'][0]['s3']['bucket']['name']
def get_object(key, bucket):
def get_object(key: str, bucket: str):
return s3.get_object(
Bucket=bucket,
Key=key,
)['Body'].read()
def get_object_tag(key, bucket, tag_key):
def get_object_tag(key: str, bucket: str, tag_key: str):
tags = get_object_tags(key, bucket)
tag_index = 0
@ -41,20 +41,10 @@ def get_object_tag(key, bucket, tag_key):
else:
return None
def get_object_tags(key, bucket):
def get_object_tags(key: str, bucket: str):
tags = s3.get_object_tagging(Bucket=bucket, Key=key)['TagSet']
return tags
def file_stream_upload(buffer, name, bucket):
return s3.upload_fileobj(buffer, bucket, name)
def receive_message(queue, message_num=1, wait_time=1):
return sqs.receive_message(QueueUrl=queue,
MaxNumberOfMessages=message_num,
WaitTimeSeconds=wait_time)
def delete_message(queue, receipt):
return sqs.delete_message(QueueUrl=queue, ReceiptHandle=receipt)
def file_stream_upload(buffer: io.BytesIO, name: str, bucket: str, action: str = None):
return s3.upload_fileobj(buffer, bucket, name, ExtraArgs={'Tagging': f'action={action}'})

View File

@ -1,7 +1,7 @@
import os, sys, logging
from lib.application_configs import ApplicationConfigs
from services.loft_service import LoftService
from services.form_generation_service import FormGenerationService
from helpers import aws_helper
from daemonize import Daemonize
@ -18,14 +18,14 @@ class ServiceListener(Consumer):
# gather/manage/process data based on the particular needs
logging.info('Incoming message: %s', body)
service = LoftService(body)
service = FormGenerationService(body)
service.process()
logging.info('Process complete for %s', service.file_name)
def main():
logging.info('Starting Solver Service: That Was Rash (v1.4.0)...')
logging.info('Starting IRT Service: That Was Rash (v1.4.0)...')
# ToDo: Figure out a much better way of doing this.
# LocalStack wants 'endpoint_url', while prod doesnt :(

View File

@ -15,7 +15,7 @@ from models.target import Target
from services.base import Base
class LoftService(Base):
class FormGenerationService(Base):
def process(self):
try:
@ -256,4 +256,4 @@ class LoftService(Base):
# upload generated file to s3 and return result
return aws_helper.file_stream_upload(
solution_file, self.file_name,
ApplicationConfigs.s3_processed_bucket)
ApplicationConfigs.s3_processed_bucket, 'formGeneration')