Pankaj Katariya
286 lines
11 KiB
Python
286 lines
11 KiB
Python
import boto3
|
|
import time
|
|
from configparser import ConfigParser
|
|
import logging
|
|
import os
|
|
from botocore.exceptions import ClientError
|
|
import json
|
|
from urllib.parse import unquote_plus
|
|
|
|
# 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 generate a Unix timestamp
|
|
def generate_unix_timestamp():
|
|
# Get the current time in seconds since the epoch
|
|
unix_timestamp = int(time.time())
|
|
return unix_timestamp
|
|
|
|
# Function to retrieve configuration values 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 move a file from source to destination in S3
|
|
def move_file_within_s3(source_bucket, source_key, destination_key):
|
|
try:
|
|
tags = "env=dev"
|
|
# Copy the file to the destination folder
|
|
s3_client.copy_object(Bucket=source_bucket, CopySource={'Bucket': source_bucket, 'Key': source_key}, Key=destination_key, Tagging=f'{tags}')
|
|
|
|
# Delete the file from the source folder
|
|
s3_client.delete_object(Bucket=source_bucket, Key=source_key)
|
|
|
|
logger.info(f"File moved from {source_key} to {destination_key}")
|
|
except ClientError as e:
|
|
logger.error(f"Error moving file: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error moving file: {e}")
|
|
|
|
# Function to get a list of PDF files in a given S3 folder
|
|
def get_pdf_files_list_from_s3(source_bucket, source_folder):
|
|
file_list = []
|
|
# List S3 Object & iterate (as per max files allowed)
|
|
s3_list_response = s3_client.list_objects_v2(Bucket=source_bucket, Prefix=source_folder)
|
|
|
|
if s3_list_response and s3_list_response['ResponseMetadata']['HTTPStatusCode'] == 200 and s3_list_response['KeyCount'] != 0:
|
|
objects = s3_list_response['Contents']
|
|
|
|
for s3_object in objects:
|
|
# Skip non-PDF files
|
|
if not s3_object['Key'].lower().endswith('.pdf'):
|
|
continue
|
|
|
|
file_list.append(s3_object['Key'])
|
|
|
|
return file_list
|
|
|
|
def start_textract_detection_job( bucket_name,
|
|
document_file_name,
|
|
sns_topic_arn,
|
|
sns_role_arn,
|
|
job_tag,):
|
|
try:
|
|
# Define the parameters for the start_document_analysis API
|
|
start_document_detection_params = {
|
|
'DocumentLocation': {
|
|
'S3Object': {
|
|
'Bucket': bucket_name,
|
|
'Name': document_file_name
|
|
}
|
|
},
|
|
'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()), # Use a unique token for each request
|
|
'JobTag': job_tag, # Use a tag to identify your job
|
|
'NotificationChannel': {
|
|
'SNSTopicArn': sns_topic_arn,
|
|
'RoleArn': sns_role_arn # Role to allow Textract service to notify SNS topic when response is ready
|
|
}
|
|
}
|
|
|
|
logger.info('start_document_detection_params ' + str(start_document_detection_params))
|
|
|
|
# Send the request to start document detection
|
|
textract_response = textract_client.start_document_text_detection(**start_document_detection_params)
|
|
|
|
job_id = textract_response["JobId"]
|
|
logger.info(
|
|
"Started text detection job %s on %s.", job_id, document_file_name
|
|
)
|
|
except ClientError:
|
|
logger.exception("Couldn't detect text in %s.", document_file_name)
|
|
raise
|
|
else:
|
|
return job_id
|
|
|
|
def start_textract_analysis_job(
|
|
bucket_name,
|
|
document_file_name,
|
|
analysis_feature_type,
|
|
sns_topic_arn,
|
|
sns_role_arn,
|
|
job_tag,
|
|
):
|
|
|
|
try:
|
|
# Define the parameters for the start_document_analysis API
|
|
start_document_analysis_params = {
|
|
'DocumentLocation': {
|
|
'S3Object': {
|
|
'Bucket': bucket_name,
|
|
'Name': document_file_name
|
|
}
|
|
},
|
|
'FeatureTypes': analysis_feature_type, # Customize based on requirements
|
|
'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()), # Use a unique token for each request
|
|
'JobTag': job_tag, # Use a tag to identify your job
|
|
'NotificationChannel': {
|
|
'SNSTopicArn': sns_topic_arn,
|
|
'RoleArn': sns_role_arn # Role to allow Textract service to notify SNS topic when response is ready
|
|
}
|
|
}
|
|
|
|
logger.info('start_document_analysis_params ' + str(start_document_analysis_params))
|
|
|
|
# Send the request to start document analysis
|
|
textract_response = textract_client.start_document_analysis(**start_document_analysis_params)
|
|
|
|
job_id = textract_response["JobId"]
|
|
logger.info(
|
|
"Started text analysis job %s on %s.", job_id, document_file_name
|
|
)
|
|
except ClientError:
|
|
logger.exception("Couldn't analyze text in %s.", document_file_name)
|
|
raise
|
|
else:
|
|
return job_id
|
|
|
|
# AWS Lambda handler function
|
|
def lambda_handler(event, context):
|
|
|
|
try:
|
|
# Extract AWS account ID and region from the Lambda ARN
|
|
aws_account_id = context.invoked_function_arn.split(":")[4]
|
|
aws_region = context.invoked_function_arn.split(":")[3]
|
|
|
|
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', '')
|
|
|
|
# Read config.properties
|
|
file_path_array = property_file_path.split("/")
|
|
|
|
# Valid if file_path_array has more than 2 elements
|
|
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))
|
|
|
|
# Extract configuration values
|
|
SOURCE_LOCATION = config_dict['FOLDER_LOCATIONS']['SOURCE_LOCATION'].format(batch_id) # SOURCE_LOCATION
|
|
STAGING_LOCATION = config_dict['FOLDER_LOCATIONS']['STAGING_LOCATION'].format(batch_id)
|
|
ANALYSIS_FEATURE_TYPE = config_dict['OTHERS']['ANALYSIS_FEATURE_TYPE'].split(",") # Analysis FeatureType
|
|
SENDER_MAX_FILES = int(config_dict['OTHERS']['SENDER_MAX_FILES'])
|
|
SNS_TOPIC_ARN = config_dict['RESOURCES']['SNS_TOPIC_ARN'].replace("{aws_region}",aws_region).replace("{aws_account_id}",aws_account_id)
|
|
TEXTRACT_ROLE_ARN = config_dict['RESOURCES']['TEXTRACT_ROLE_ARN'].replace("{aws_account_id}",aws_account_id) # Textract IAM Role ARN to publish to SNS
|
|
JOB_TAG = config_dict['OTHERS']['JOB_TAG']
|
|
PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper()
|
|
|
|
logger.info('SOURCE_LOCATION: ' + SOURCE_LOCATION)
|
|
logger.info('STAGING_LOCATION: ' + STAGING_LOCATION)
|
|
logger.info('ANALYSIS_FEATURE_TYPE: ' + str(ANALYSIS_FEATURE_TYPE))
|
|
logger.info('SNS_TOPIC_ARN: ' + SNS_TOPIC_ARN)
|
|
logger.info('TEXTRACT_ROLE_ARN: ' + TEXTRACT_ROLE_ARN)
|
|
logger.info('SENDER_MAX_FILES: ' + str(SENDER_MAX_FILES))
|
|
logger.info('JOB_TAG: ' + str(JOB_TAG))
|
|
logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE))
|
|
|
|
# File count
|
|
file_count = 0
|
|
|
|
# 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'])
|
|
|
|
#logger.info('Message Count: ', str(len(record_body['Records'])) )
|
|
|
|
for sqs_record in record_body['Records']:
|
|
|
|
# Construct the source and destination paths
|
|
source_path = unquote_plus(sqs_record['s3']['object']['key'])
|
|
destination_path = STAGING_LOCATION + source_path.replace(SOURCE_LOCATION,"")
|
|
|
|
# Move file to stagging
|
|
move_file_within_s3(S3_BUCKET_NAME, source_path, destination_path)
|
|
|
|
job_id = ""
|
|
|
|
if PROCESS_TYPE == "ANALYSIS":
|
|
# Start Textract analysis job
|
|
job_id = start_textract_analysis_job (
|
|
S3_BUCKET_NAME,
|
|
destination_path,
|
|
ANALYSIS_FEATURE_TYPE,
|
|
SNS_TOPIC_ARN,
|
|
TEXTRACT_ROLE_ARN,
|
|
JOB_TAG,
|
|
)
|
|
elif PROCESS_TYPE == "DETECTION":
|
|
# Start Textract detection job
|
|
job_id = start_textract_detection_job (
|
|
S3_BUCKET_NAME,
|
|
destination_path,
|
|
SNS_TOPIC_ARN,
|
|
TEXTRACT_ROLE_ARN,
|
|
JOB_TAG,
|
|
)
|
|
|
|
file_count = file_count + 1
|
|
logger.info(str(file_count) + '. ' + str(source_path) + " Job Id: " + str(job_id))
|
|
|
|
success_message = 'Total files sent to textract : '+ str(file_count)
|
|
logger.info(success_message)
|
|
return {
|
|
'statusCode': 200,
|
|
'body': success_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
|
|
}
|
|
|
|
except ClientError as e:
|
|
# Handle specific Textract client errors
|
|
error_message = f"Error in Textract operation: {e}"
|
|
logger.error(error_message)
|
|
return {
|
|
'statusCode': 500,
|
|
'body': error_message
|
|
}
|
|
|
|
except Exception as e:
|
|
# Handle other exceptions
|
|
error_message = f"Unexpected error: {e}"
|
|
logger.error(error_message)
|
|
return {
|
|
'statusCode': 500,
|
|
'body': error_message
|
|
}
|
|
|
|
|