Moved files within lambda folder
This commit is contained in:
Pankaj Katariya
@@ -0,0 +1,228 @@
|
||||
import json
|
||||
import logging
|
||||
import boto3
|
||||
import os
|
||||
from configparser import ConfigParser
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
# Initialize logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
# Initialize S3 & Textract clients
|
||||
s3_client = boto3.client('s3')
|
||||
textract_client = boto3.client('textract')
|
||||
|
||||
# Function to load configuration from S3
|
||||
def load_config_from_s3(bucket_name, file_key):
|
||||
|
||||
# Download the config file from S3
|
||||
response = s3_client.get_object(Bucket=bucket_name, Key=file_key)
|
||||
config_content = response['Body'].read().decode('utf-8')
|
||||
|
||||
# Parse the config file
|
||||
config_parser = ConfigParser()
|
||||
config_parser.read_string(config_content)
|
||||
|
||||
# Convert the configuration to a dictionary
|
||||
config_dict = {}
|
||||
for section in config_parser.sections():
|
||||
config_dict[section] = {key.upper(): value for key, value in config_parser.items(section)}
|
||||
|
||||
return config_dict
|
||||
|
||||
# Function to construct JSON path from S3 object path
|
||||
def generate_output_json_path(src_folder,dest_folder, s3_object):
|
||||
|
||||
# split string to remove source folder path
|
||||
txt = s3_object.split(src_folder)
|
||||
if len(txt) == 2:
|
||||
txt = txt[1]
|
||||
else:
|
||||
txt = txt[0]
|
||||
# change file extension - remove .pdf and add .json
|
||||
if txt.lower().endswith(".pdf"):
|
||||
return dest_folder+txt[0:-4]+".json"
|
||||
|
||||
# Function to get Textract document analysis
|
||||
def get_textract_document_analysis(job_id, textract_client):
|
||||
# Initialize an empty list to store blocks
|
||||
all_blocks = []
|
||||
next_token = None
|
||||
response = {}
|
||||
flag = True
|
||||
|
||||
try:
|
||||
while True:
|
||||
if flag:
|
||||
# Call the Textract API to get document analysis
|
||||
response = textract_client.get_document_analysis(
|
||||
JobId=job_id
|
||||
)
|
||||
flag = False
|
||||
else:
|
||||
# Call the Textract API to get document analysis
|
||||
response = textract_client.get_document_analysis(
|
||||
JobId=job_id,
|
||||
NextToken=next_token
|
||||
)
|
||||
|
||||
job_status = response["JobStatus"]
|
||||
logger.info("Job %s status is %s.", job_id, job_status)
|
||||
|
||||
# Merge the blocks from the current response
|
||||
all_blocks.extend(response.get('Blocks', []))
|
||||
|
||||
# Check if there are more blocks to retrieve
|
||||
next_token = response.get('NextToken')
|
||||
if not next_token:
|
||||
logger.info("No more Textract response to retrieve")
|
||||
break
|
||||
|
||||
except ClientError:
|
||||
logger.exception("Couldn't get data for job %s.", job_id)
|
||||
raise
|
||||
else:
|
||||
# Remove unnecessary keys from the last response
|
||||
last_response = response.copy()
|
||||
last_response.pop('Blocks', None)
|
||||
last_response.pop('ResponseMetadata', None)
|
||||
logger.info("Removed 'ResponseMetadata' key")
|
||||
|
||||
# Merge with {'Blocks': all_blocks}
|
||||
final_response = {'Blocks': all_blocks}
|
||||
final_response.update(last_response)
|
||||
logger.info("Final Textract response is contructed")
|
||||
return final_response
|
||||
|
||||
# Function to save Textract response to S3
|
||||
def upload_response_to_s3(response, bucket_name, object_key, s3_client):
|
||||
# Convert the response to JSON
|
||||
response_json = json.dumps(response)
|
||||
|
||||
try:
|
||||
# Upload the JSON response to S3
|
||||
s3_client.put_object(
|
||||
Bucket=bucket_name,
|
||||
Key=object_key,
|
||||
Body=response_json,
|
||||
ContentType='application/json'
|
||||
)
|
||||
logger.info(f"Saved analysis response to S3: {object_key}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving response to S3 - {object_key}: {str(e)}")
|
||||
|
||||
# Function to move file within S3
|
||||
def move_file_within_s3(source_bucket, source_path, destination_path):
|
||||
try:
|
||||
# Copy the file to the destination folder
|
||||
s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_path}, Key=destination_path)
|
||||
|
||||
# Delete the file from the source folder
|
||||
s3_client.delete_object(Bucket=source_bucket, Key=source_path)
|
||||
|
||||
logger.info(f"File moved from {source_path} to {destination_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error moving file: {e}")
|
||||
|
||||
# AWS Lambda handler function
|
||||
def lambda_handler(event, context):
|
||||
|
||||
logger.info('## ENVIRONMENT VARIABLES\r' + str(os.environ))
|
||||
|
||||
# Read environment variables
|
||||
property_file_path = os.environ.get('PROPERTY_FILE_S3_PATH', '')
|
||||
batch_id = os.environ.get('BATCH_ID', '')
|
||||
|
||||
# Validate environment variables
|
||||
file_path_array = property_file_path.split("/")
|
||||
if len(file_path_array) > 1:
|
||||
|
||||
# Extract BUCKET_NAME and config_file_path
|
||||
S3_BUCKET_NAME = file_path_array[0]
|
||||
CONFIG_FILE_PATH = "/".join(file_path_array[1:])
|
||||
|
||||
logger.info(f'S3_BUCKET_NAME: {S3_BUCKET_NAME}')
|
||||
logger.info(f'CONFIG_FILE_PATH: {CONFIG_FILE_PATH}')
|
||||
|
||||
# Load config file
|
||||
config_dict = load_config_from_s3(S3_BUCKET_NAME, CONFIG_FILE_PATH)
|
||||
|
||||
logger.info('## CONFIG DICTIONARY\r' + str(config_dict))
|
||||
|
||||
STAGING_LOCATION = config_dict['FOLDER_LOCATIONS']['STAGING_LOCATION'].format(batch_id)
|
||||
OUTPUT_LOCATION = config_dict['FOLDER_LOCATIONS']['OUTPUT_LOCATION'].format(batch_id)
|
||||
PROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['PROCESSED_LOCATION'].format(batch_id)
|
||||
UNPROCESSED_LOCATION = config_dict['FOLDER_LOCATIONS']['UNPROCESSED_LOCATION'].format(batch_id)
|
||||
|
||||
logger.info('STAGGING_LOCATION: ' + STAGING_LOCATION)
|
||||
logger.info('OUTPUT_LOCATION: ' + OUTPUT_LOCATION)
|
||||
logger.info('PROCESSED_LOCATION: ' + PROCESSED_LOCATION)
|
||||
logger.info('UNPROCESSED_LOCATION: ' + UNPROCESSED_LOCATION)
|
||||
|
||||
# Process each message from the SQS event
|
||||
for record in event['Records']:
|
||||
# Extract the message body from the record
|
||||
record_body = json.loads(record['body'])
|
||||
message_body = json.loads(record_body['Message'])
|
||||
logger.info('MESSAGE_BODY: ' + str(message_body))
|
||||
try:
|
||||
|
||||
# Extract relevant information from the message body
|
||||
job_id = message_body.get('JobId')
|
||||
document_location = message_body.get('DocumentLocation')
|
||||
s3_object_name = document_location.get('S3ObjectName')
|
||||
|
||||
# Check if the status is "SUCCEEDED"
|
||||
if message_body.get('Status') == 'SUCCEEDED':
|
||||
|
||||
# Call the function to get document analysis using Textract
|
||||
document_analysis = get_textract_document_analysis(job_id, textract_client)
|
||||
|
||||
# Save the document analysis response to S3
|
||||
s3_object_key = generate_output_json_path(STAGING_LOCATION,OUTPUT_LOCATION, s3_object_name)
|
||||
upload_response_to_s3(document_analysis, S3_BUCKET_NAME, s3_object_key, s3_client)
|
||||
|
||||
# Construct the destination paths
|
||||
destination_path = PROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"")
|
||||
|
||||
# Move file to processed folder
|
||||
move_file_within_s3(S3_BUCKET_NAME, s3_object_name, destination_path)
|
||||
|
||||
success_message = 'Processed file '+ str(s3_object_key)
|
||||
logger.info(success_message)
|
||||
return {
|
||||
'statusCode': 200,
|
||||
'body': success_message
|
||||
}
|
||||
else:
|
||||
error_message = f"Skipping message with JobId {message_body.get('JobId')} as Status is not 'SUCCEEDED'"
|
||||
logger.info(error_message)
|
||||
|
||||
# Construct the destination paths
|
||||
destination_path = UNPROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"")
|
||||
|
||||
# Move file to unprocessed folder
|
||||
move_file_within_s3(S3_BUCKET_NAME, s3_object_name, destination_path)
|
||||
return {
|
||||
'statusCode': 500,
|
||||
'body': error_message
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
error_message = f"Error processing message with JobId {message_body.get('JobId')}: {str(e)}"
|
||||
logger.error(error_message)
|
||||
return {
|
||||
'statusCode': 500,
|
||||
'body': error_message
|
||||
}
|
||||
|
||||
else:
|
||||
error_message = 'Incorrect value for ENVIRONMENT VARIABLES: PROPERTY_FILE_S3_PATH\r' + str(property_file_path)
|
||||
logger.error(error_message)
|
||||
|
||||
return {
|
||||
'statusCode': 500,
|
||||
'body': error_message
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user