Added detection capability and sender orchectration
This commit is contained in:
Pankaj Katariya
@@ -4,6 +4,8 @@ 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__)
|
||||
@@ -40,8 +42,9 @@ def load_config_from_s3(bucket_name, file_key):
|
||||
# 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)
|
||||
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)
|
||||
@@ -70,12 +73,50 @@ def get_pdf_files_list_from_s3(source_bucket, source_folder):
|
||||
|
||||
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:
|
||||
@@ -89,7 +130,7 @@ def start_textract_analysis_job(
|
||||
},
|
||||
'FeatureTypes': analysis_feature_type, # Customize based on requirements
|
||||
'ClientRequestToken': 'unique-token-'+str(generate_unix_timestamp()), # Use a unique token for each request
|
||||
'JobTag': 'healthcare-contract', # Use a tag to identify your job
|
||||
'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
|
||||
@@ -150,58 +191,70 @@ def lambda_handler(event, context):
|
||||
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))
|
||||
|
||||
# List S3 Object & iterate (as per max files allowed)
|
||||
files_list = get_pdf_files_list_from_s3(S3_BUCKET_NAME,SOURCE_LOCATION)
|
||||
|
||||
if len(files_list):
|
||||
# File count
|
||||
file_count = 0
|
||||
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']:
|
||||
|
||||
for s3_file_key in files_list:
|
||||
# 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 = s3_file_key
|
||||
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)
|
||||
|
||||
# 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_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(s3_file_key) + " Job Id: " + str(job_id))
|
||||
|
||||
if SENDER_MAX_FILES == file_count:
|
||||
break
|
||||
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:
|
||||
message = 'No files found'
|
||||
logger.error(message)
|
||||
return {
|
||||
'statusCode': 500,
|
||||
'body': message
|
||||
}
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user