Merged in dev_pkatariya_sprint_3 (pull request #14)

Dev pkatariya sprint 3

lambda/receiver/lambda_function.py

@@ -4,6 +4,7 @@ import boto3
 import os
 from configparser import ConfigParser
 from botocore.exceptions import ClientError
+from urllib.parse import urlencode
 
 # Initialize logger
 logger = logging.getLogger(__name__)
@@ -44,6 +45,57 @@ def generate_output_json_path(src_folder,dest_folder, s3_object):
     if txt.lower().endswith(".pdf"):
         return dest_folder+txt[0:-4]+".json"
 
+# Function to get Textract document analysis
+def get_textract_document_detection(job_id, textract_client):
+    # Initialize an empty list to store blocks
+    all_blocks = []
+    next_token = None
+    response = {}
+    flag = True
+    logger.info("Get document detection")
+    try:
+        while True:
+            if flag:
+                # Call the Textract API to get document analysis
+                response = textract_client.get_document_text_detection(
+                    JobId=job_id
+                )
+                flag = False
+            else:
+                # Call the Textract API to get document analysis
+                response = textract_client.get_document_text_detection(
+                    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 get Textract document analysis
 def get_textract_document_analysis(job_id, textract_client):
     # Initialize an empty list to store blocks
@@ -51,7 +103,7 @@ def get_textract_document_analysis(job_id, textract_client):
     next_token = None
     response = {}
     flag = True
-    
+    logger.info("Get document analysis")
     try:
         while True:
             if flag:
@@ -96,7 +148,7 @@ def get_textract_document_analysis(job_id, textract_client):
         return final_response
 
 # Function to save Textract response to S3
-def upload_response_to_s3(response, bucket_name, object_key, s3_client):
+def upload_response_to_s3(response, bucket_name, object_key, s3_client, tags):
     # Convert the response to JSON
     response_json = json.dumps(response)
     
@@ -106,7 +158,8 @@ def upload_response_to_s3(response, bucket_name, object_key, s3_client):
             Bucket=bucket_name,
             Key=object_key,
             Body=response_json,
-            ContentType='application/json'
+            ContentType='application/json',
+            Tagging=tags
         )
         logger.info(f"Saved analysis response to S3: {object_key}")
     except Exception as e:
@@ -125,6 +178,27 @@ def move_file_within_s3(source_bucket, source_path, destination_path):
     except Exception as e:
         logger.error(f"Error moving file: {e}")
 
+# Function to get s3 object tags
+def get_s3_object_tags(bucket_name, object_key):
+
+    try:
+        # Get object tags
+        response = s3_client.get_object_tagging(
+            Bucket=bucket_name,
+            Key=object_key
+        )
+
+        # Extract tags from the response and convert to dictionary
+        tags_list = response['TagSet']
+        tags_dict = {tag['Key']: tag['Value'] for tag in tags_list}
+
+        return tags_dict
+
+    except Exception as e:
+        logger.exception(f"Error: {e}")
+        print(f"Error: {e}")
+        return None
+    
 # AWS Lambda handler function
 def lambda_handler(event, context):
 
@@ -154,11 +228,13 @@ def lambda_handler(event, context):
         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)
+        PROCESS_TYPE = str(config_dict['OTHERS']['PROCESS_TYPE']).upper()
 
         logger.info('STAGGING_LOCATION: ' + STAGING_LOCATION)
         logger.info('OUTPUT_LOCATION: ' + OUTPUT_LOCATION)
         logger.info('PROCESSED_LOCATION: ' + PROCESSED_LOCATION)
         logger.info('UNPROCESSED_LOCATION: ' + UNPROCESSED_LOCATION)
+        logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE))
 
         # Process each message from the SQS event
         for record in event['Records']:
@@ -172,16 +248,30 @@ def lambda_handler(event, context):
                 job_id = message_body.get('JobId')
                 document_location = message_body.get('DocumentLocation')
                 s3_object_name = document_location.get('S3ObjectName')
+
+                # Read tags from stagging file
+                tags_dict = get_s3_object_tags(S3_BUCKET_NAME,s3_object_name)
                 
+                # encode to URL Query parameter
+                tags = urlencode(tags_dict)
+
                 # Check if the status is "SUCCEEDED"
                 if message_body.get('Status') == 'SUCCEEDED':
 
-                    # Call the function to get document analysis using Textract
+                    document = {}
-                    document_analysis = get_textract_document_analysis(job_id, textract_client)
+                    if PROCESS_TYPE == "ANALYSIS":
+
+                        # Call the function to get document analysis using Textract
+                        document = get_textract_document_analysis(job_id, textract_client)
+
+                    elif PROCESS_TYPE == "DETECTION":
+
+                        # Call the function to get document analysis using Textract
+                        document = get_textract_document_detection(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)
+                    upload_response_to_s3(document, S3_BUCKET_NAME, s3_object_key, s3_client,tags)
 
                     # Construct the destination paths
                     destination_path = PROCESSED_LOCATION + s3_object_name.replace(STAGING_LOCATION,"")

lambda/sender/lambda_function.py

@@ -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
@@ -111,6 +152,27 @@ def start_textract_analysis_job(
     else:
         return job_id
 
+# Function to get s3 object tags
+def get_s3_object_tags(bucket_name, object_key):
+
+    try:
+        # Get object tags
+        response = s3_client.get_object_tagging(
+            Bucket=bucket_name,
+            Key=object_key
+        )
+
+        # Extract tags from the response and convert to dictionary
+        tags_list = response['TagSet']
+        tags_dict = {tag['Key']: tag['Value'] for tag in tags_list}
+
+        return tags_dict
+
+    except Exception as e:
+        logger.exception(f"Error: {e}")
+        print(f"Error: {e}")
+        return None
+
 # AWS Lambda handler function
 def lambda_handler(event, context):
 
@@ -150,58 +212,78 @@ 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))                
+            logger.info('SENDER_MAX_FILES: ' + str(SENDER_MAX_FILES))
-            
+            logger.info('JOB_TAG: ' + str(JOB_TAG))
-            # List S3 Object & iterate (as per max files allowed)
+            logger.info('PROCESS_TYPE: ' + str(PROCESS_TYPE))
-            files_list = get_pdf_files_list_from_s3(S3_BUCKET_NAME,SOURCE_LOCATION)
+
-            
+            # File count
-            if len(files_list):
+            file_count = 0
-                # 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
+                    # Read tags from file
-                    job_id = start_textract_analysis_job (
+                    tags_dict = get_s3_object_tags(S3_BUCKET_NAME,destination_path)
-                            S3_BUCKET_NAME,
+                    
-                            destination_path,
+                    # check if key exist in tags_dict
-                            ANALYSIS_FEATURE_TYPE,
+                    if "batch_id" in tags_dict.keys():
-                            SNS_TOPIC_ARN,
+                        JOB_TAG = JOB_TAG + "-" + tags_dict['batch_id']
-                            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))
+                    logger.info(str(file_count) + '. ' + str(source_path) + " Job Id: " + str(job_id))
-                    
+                    time.sleep(1)
-                    if SENDER_MAX_FILES == file_count:
-                        break
 
-                success_message = 'Total files sent to textract : '+ str(file_count)
+            success_message = 'Total files sent to textract : '+ str(file_count)
-                logger.info(success_message)
+            logger.info(success_message)
-                return {
+            return {
-                    'statusCode': 200,
+                'statusCode': 200,
-                    'body': success_message
+                'body': success_message
-                }
+            }
-            else:
+                
-                message = 'No files found'
-                logger.error(message)
-                return {
-                        'statusCode': 500,
-                        'body': message
-                    }
 
         else:
             error_message = 'Incorrect value for ENVIRONMENT VARIABLES: PROPERTY_FILE_S3_PATH\r' + str(property_file_path)