doczyai-pipelines/archive/ops_scripts/generic/textract_file_reconciliation.py

import boto3
import os
import pandas as pd
from datetime import datetime
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed


"""

This script is used to move files from one S3 bucket to another S3 bucket based on the batch_ids provided in the CSV file.
The script will move the text files, stuck files, PDF files and invalid files to the destination bucket.
The copy operation is done using multithreading to speed up the process. Use the max_workers parameter to adjust the number of threads.

This script is NOT CLIENT SPECIFIC. It is a generic script that can be used for any client.

Usage: 
1. Update the configuration section with the required parameters.
2. Run the script.

Note:
- Ensure you have the necessary permissions to read from the source bucket and write to the destination bucket (Admin role is preferred).

"""

def log_error_to_file(key, error_message, error_log_file):
    # Log error to file
    # Create new file if it doesnt exist 
    if not os.path.exists(error_log_file):
        with open(error_log_file, 'w') as f:
            f.write("Error Log\n")
            f.write("---------\n")

    with open(error_log_file, 'a') as f:
        f.write(f"Error with file: {key} - {error_message}\n")

def list_filtered_files(bucket_name, prefix, start_date, end_date, profile_name):
    start_date = datetime.fromisoformat(start_date.replace('Z', '+00:00'))
    end_date = datetime.fromisoformat(end_date.replace('Z', '+00:00'))

    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client('s3')

    paginator = s3_client.get_paginator('list_objects_v2')
    page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    filtered_files = []
    for page in page_iterator:
        if 'Contents' in page:
            for obj in page['Contents']:
                last_modified = obj['LastModified']
                if start_date <= last_modified <= end_date:
                    filtered_files.append(obj['Key'])

    return filtered_files

def copy_single_file(s3_client, source_bucket, destination_bucket, key, destination_prefix, error_log_file):
    try:
        copy_source = {'Bucket': source_bucket, 'Key': key}
        destination_key = os.path.join(destination_prefix, os.path.basename(key))
        print(f"Copying {key} to {destination_key}")
        s3_client.copy(copy_source, destination_bucket, destination_key)
    except Exception as e:
        error_message = str(e)
        print(f"Error copying file {key}: {error_message}")
        # Log error to file
        log_error_to_file(key, error_message, error_log_file)

def copy_files_to_destination(source_bucket, destination_bucket, file_keys, profile_name, destination_prefix, max_workers, error_log_file):
    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client('s3')

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(copy_single_file, s3_client, source_bucket, destination_bucket, key, destination_prefix, error_log_file) for key in file_keys]
        
        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Error in thread: {e}")

    print("Copy completed.")


def main():
    

    # ------------------------------ Configuration ------------------------------


    # Error log file path
    # There is usually no errors in this script, but in case of any errors, they will be logged to this file
    error_log_file = "error_log_batch5.txt"


    # Main script
    # Fetch the batch_ids csv file from the client bucket ('centene-national-contracting-files')/processed_batch_ids/ folder
    csv_file_path = "batch5_diff_files.csv"

    # Source and destination bucket names. Where source bucket is the textract pipeline where all files reside in their batch folders
    source_bucket = 'doczyai-use2-u-cn1-s3-textract-processing-001'

    # Destination bucket is the client bucket where we move the files to satge it for DS team
    destination_bucket = 'centene-national-contracting-files'

    # Depending on what batch files we want to move, we can change the destination prefix
    destination_prefix = 'batch_5_priority_files/txt_files/'  

    # If the batch size in scheduler is large, there is a possibility of throttling and files getting stuck in the textract pipeline
    # We move these stuck files to a separate folder so that it can be sent for reprocessing
    stuck_destination_prefix = 'batch_5_priority_files/throttled_files_final/'

    # PDF files are required for the MCS team to review the contracts
    pdf_files_destination = f'batch_5_priority_files/pdf_files/'

    # Less than 1% of the files are invalid files. These files are moved to a separate folder
    invalid_files_destination = f'batch_5_priority_files/invalid_files/'

    # Date range for files to be moved so that we can limit the list files in bucket operation 
    start_date = '2024-07-07T00:00:00Z'  # ISO 8601 format
    end_date = '2024-12-31T23:59:59Z'  # ISO 8601 format

    # If you are using AWS CLI profiles, provide the profile name here to authenticate
    profile_name = 'doczy_uat'
    max_workers = 50  # Adjust the number of threads as needed


    # ------------------------------ Script ------------------------------

    batch_ids_df = pd.read_csv(csv_file_path)
    batch_ids = batch_ids_df['batch_id'].unique()

    print(f"Total number of batches: {len(batch_ids)}")
    print(f"Batch IDs: {batch_ids}")

    text_count = 0
    stuck_count = 0
    pdf_count = 0
    invalid_count = 0
        
    # Bulk copy files from S3 for all batch_ids
    for batch_id in batch_ids:
        prefix = f'contract-text-file/{batch_id}/'
        stuck_files_path = f'textract-sender-staging-pdfs/{batch_id}/'
        pdf_files_path = f'textract-receiver-processed-pdfs/{batch_id}'
        invalid_files_path = f'invalid-pdfs/{batch_id}'

        # List filtered text files
        filtered_files = list_filtered_files(source_bucket, prefix, start_date, end_date, profile_name)
        print(f"Batch_id: {batch_id} Filtered files: {len(filtered_files)}")
        text_count += len(filtered_files)

        # Throttled files list
        stuck_files = list_filtered_files(source_bucket, stuck_files_path, start_date, end_date, profile_name)
        print(f"Batch_id: {batch_id} Stuck files: {len(stuck_files)}")
        stuck_count += len(stuck_files)

        # PDF file list
        pdf_files = list_filtered_files(source_bucket, pdf_files_path, start_date, end_date, profile_name)
        print(f"Batch_id: {batch_id} PDF files: {len(pdf_files)}")
        pdf_count += len(pdf_files)

        # Invalid files list
        invalid_files = list_filtered_files(source_bucket, invalid_files_path, start_date, end_date, profile_name)
        print(f"Batch_id: {batch_id} Invalid files: {len(invalid_files)}")
        invalid_count += len(invalid_files)

        # Copy files to destination using multithreading
        copy_files_to_destination(source_bucket, destination_bucket, filtered_files, profile_name, destination_prefix, max_workers, error_log_file)
        copy_files_to_destination(source_bucket, destination_bucket, stuck_files, profile_name, stuck_destination_prefix, max_workers, error_log_file)
        copy_files_to_destination(source_bucket, destination_bucket, pdf_files, profile_name, pdf_files_destination, max_workers, error_log_file)
        copy_files_to_destination(source_bucket, destination_bucket, invalid_files, profile_name, invalid_files_destination, max_workers, error_log_file)

    print(f"Total text files copied: {text_count}")
    print(f"Total stuck files copied: {stuck_count}")
    print(f"Total pdf files copied: {pdf_count}")
    print(f"Total invalid files copied: {invalid_count}")


if __name__ == '__main__':
    main()