doczyai-pipelines/ops_scripts/generic/s3_prefix_search_and_copy.py

import csv
import boto3
from concurrent.futures import ThreadPoolExecutor, as_completed
import os

"""
This is an adhoc script that is used tot copy files from one prefix to another prefix within the same S3 bucket.

TYPICALLY USED FOR CENTENE-HEALTHNET 

The need for this script is when the client team requests re-run and the previous runs are stored in a different prefixes. 

The script reads a CSV file containing 'File Name' and 'Type' columns. It then searches for these files in the specified prefixes
and copies them to a new destination prefix within the same bucket.

The script also writes a CSV file containing the files that were not found in the specified prefixes.

The script is multithreaded to improve performance when searching for files in multiple prefixes.

This script can be used to just location files by commenting out the copy functionality.

Usage:
1. Update the configuration section below with the appropriate values.
2. Run the script using Python 3.

Note:
- Ensure that the CSV file contains 'File Name' and 'Type' columns.
- Ensure you have the required permissions to read from and write to the specified S3 bucket.

"""

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

# Path to your CSV file from client / DS team that need to be staged
CSV_FILE = 'C:\\Doczy\\Health net\\no_tin_filenames.csv' # Ensure this file has 'File Name' and 'Type' columns
# AWS Configuration
PROFILE_NAME = 'doczy_uat'       # Replace with your AWS profile name
BUCKET_NAME = 'centene-healthnet-files'   # Replace with your S3 bucket name

# List of prefixes to search within the bucket
PREFIXES = [
    'facility_rerun_092524/txt_files/',
    'facility-remaining-text-files/',
    'agreements/text_files/',
    'ammendments/text_files/',
    '1kfacilityrun/',
    'ancillary_1200_rerun/txt_files/',
    'facility_remaining_1k_rerun/',
    'facility-table-issue-files/',
    'facility-test-files/',
    'facility-textract-rerun/txt_files/',
    'HN-10-Ancillary-Mix/',
    'lessthan200pagesfiles/',
    'new_prof_and_ancillary_1200/',
    'professional_text_files_for_llm/',
    'text_files_for_llm/',
    'professional_paragraph_format_rerun/',
    'HN-10-Ancillary-Mix/',
    'facility_rerun_092524/',
    'adhoc/'
    # Add more prefixes as needed
]

# Destination prefix for copying files (within the same bucket)
DESTINATION_PREFIX = 'facility_rerun_112124/no_tin/'  # Replace with your desired destination prefix

# Number of threads for multithreading
MAX_WORKERS = 100

# Output CSV for files not found
NOT_FOUND_CSV = 'C:\\Doczy\\Health net\\files_not_found_no_tin.csv'


# --------------------- Functions ---------------------

def create_s3_client(profile_name):
    """
    Create a Boto3 S3 client using the specified AWS profile.
    """
    session = boto3.Session(profile_name=profile_name)
    s3_client = session.client('s3')
    return s3_client

def normalize_string(s):
    """
    Normalize a string by removing spaces and converting to lowercase.
    
    Parameters:
        s (str): The string to normalize.
    
    Returns:
        str: The normalized string.
    """
    return ''.join(s.split()).lower()

def read_csv(csv_file):
    """
    Read the CSV file and extract the list of file names along with their types.
    
    Parameters:
        csv_file (str): Path to the CSV file.
    
    Returns:
        list of dict: List containing dictionaries with 'File Name' and 'Type'.
    """
    files = []
    try:
        with open(csv_file, mode='r', newline='', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                files.append({
                    'File Name': row['File Name'],
                    'Type': row['Type']
                })
    except FileNotFoundError:
        print(f"Error: The file {csv_file} does not exist.")
    except KeyError:
        print("Error: CSV file must contain 'File Name' and 'Type' columns.")
    except Exception as e:
        print(f"An unexpected error occurred while reading the CSV: {e}")
    return files

def list_objects_in_prefix(s3_client, bucket, prefix):
    """
    List all objects in a given prefix and return a mapping of normalized keys to actual keys.
    
    Parameters:
        s3_client (boto3.client): Boto3 S3 client.
        bucket (str): S3 bucket name.
        prefix (str): Prefix to list objects from.
    
    Returns:
        dict: Mapping of normalized keys to actual S3 keys.
    """
    normalized_key_mapping = {}
    continuation_token = None
    while True:
        try:
            if continuation_token:
                response = s3_client.list_objects_v2(
                    Bucket=bucket,
                    Prefix=prefix,
                    ContinuationToken=continuation_token
                )
            else:
                response = s3_client.list_objects_v2(
                    Bucket=bucket,
                    Prefix=prefix
                )
            if 'Contents' in response:
                for obj in response['Contents']:
                    actual_key = obj['Key']
                    normalized_key = normalize_string(os.path.basename(actual_key).replace('.txt', ''))
                    normalized_key_mapping[normalized_key] = actual_key
            if response.get('IsTruncated'):
                continuation_token = response.get('NextContinuationToken')
            else:
                break
        except Exception as e:
            print(f"Error listing objects in prefix {prefix}: {e}")
            break
    return normalized_key_mapping

def build_normalized_key_mapping(s3_client, bucket, prefixes):
    """
    Build a mapping from normalized filenames to actual S3 keys across all specified prefixes.
    
    Parameters:
        s3_client (boto3.client): Boto3 S3 client.
        bucket (str): S3 bucket name.
        prefixes (list): List of prefixes to search.
    
    Returns:
        dict: Mapping of normalized filenames to actual S3 keys.
    """
    normalized_key_mapping = {}
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        future_to_prefix = {
            executor.submit(list_objects_in_prefix, s3_client, bucket, prefix): prefix
            for prefix in prefixes
        }
        for future in as_completed(future_to_prefix):
            prefix = future_to_prefix[future]
            try:
                prefix_mapping = future.result()
                # Handle potential duplicate normalized keys
                for norm_key, actual_key in prefix_mapping.items():
                    if norm_key in normalized_key_mapping:
                        # This case is has never been hit in the past
                        # It is possible that the same file name is present in multiple prefixes (For Centene National Contracting)
                        # print(f"Warning: Duplicate normalized key '{norm_key}' found in prefix '{prefix}'.")
                        pass
                    else:
                        normalized_key_mapping[norm_key] = actual_key
            except Exception as exc:
                print(f"Prefix {prefix} generated an exception: {exc}")
    return normalized_key_mapping

def copy_file(s3_client, bucket, source_key, destination_key):
    """
    Copy a file from source_key to destination_key within the same bucket.
    
    Parameters:
        s3_client (boto3.client): Boto3 S3 client.
        bucket (str): S3 bucket name.
        source_key (str): Source S3 key.
        destination_key (str): Destination S3 key.
    """
    copy_source = {'Bucket': bucket, 'Key': source_key}
    try:
        s3_client.copy(copy_source, bucket, destination_key)
        print(f"Copied {source_key} to {destination_key}")
    except Exception as e:
        print(f"Failed to copy {source_key} to {destination_key}: {e}")

def write_not_found_csv(not_found_files, output_csv):
    """
    Write the list of not found files to a CSV file.
    
    Parameters:
        not_found_files (list of dict): List containing dictionaries with 'File Name' and 'Type'.
        output_csv (str): Path to the output CSV file.
    """
    try:
        with open(output_csv, mode='w', newline='', encoding='utf-8') as f:
            writer = csv.DictWriter(f, fieldnames=['File Name', 'Type'])
            writer.writeheader()
            for file in not_found_files:
                writer.writerow(file)
        print(f"Successfully wrote {len(not_found_files)} not found files to {output_csv}")
    except Exception as e:
        print(f"Failed to write not found files to CSV: {e}")

def main():
    # Create S3 client
    s3_client = create_s3_client(PROFILE_NAME)
    
    # Read filenames and types from CSV
    files = read_csv(CSV_FILE)
    if not files:
        print("No files to process. Exiting.")
        return
    print(f"Total files to search: {len(files)}")
    
    # Normalize filenames and build a mapping from normalized name to original file info
    normalized_files_mapping = {}
    for file in files:
        original_name = file['File Name']
        normalized_name = normalize_string(original_name) + '.txt'  # Append '.txt' as per original logic
        normalized_files_mapping[normalized_name] = file
    
    # Build normalized key mapping from S3
    print("Listing and mapping S3 objects...")
    normalized_key_mapping = build_normalized_key_mapping(s3_client, BUCKET_NAME, PREFIXES)
    print(f"Total S3 objects mapped: {len(normalized_key_mapping)}")
    
    # Match normalized filenames with normalized S3 keys
    found_files = {}  # Mapping of original file info to actual S3 keys
    for norm_filename, file_info in normalized_files_mapping.items():
        norm_key = normalize_string(norm_filename.replace('.txt', ''))  # Normalize without '.txt'
        actual_key = normalized_key_mapping.get(norm_key)
        if actual_key:
            found_files[file_info['File Name']] = actual_key
        else:
            # File not found; will be handled later
            pass
    
    # Display the results
    total_found = len(found_files)
    print(f"\nTotal files found: {total_found}")
    for fname, key in found_files.items():
        print(f"Found: {key}")
    
    # Identify files not found
    found_filenames = set(found_files.keys())
    not_found_files = [file for file in files if file['File Name'] not in found_filenames]
    
    # Write not found files to CSV
    if not_found_files:
        write_not_found_csv(not_found_files, NOT_FOUND_CSV)
    else:
        print("All files were found. No entries to write to not_found CSV.")
    
    # ----------------- Copy Functionality -----------------
    
    print("\nStarting to copy files to the destination prefix...")

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as copy_executor:
        copy_futures = []
        for filename, source_key in found_files.items():
            destination_key = os.path.join(DESTINATION_PREFIX, os.path.basename(source_key))
            copy_futures.append(
                copy_executor.submit(copy_file, s3_client, BUCKET_NAME, source_key, destination_key)
            )
        
        # Optionally, wait for all copy operations to complete
        for future in as_completed(copy_futures):
            try:
                future.result()
            except Exception as exc:
                print(f"Copy operation generated an exception: {exc}")
    
    print("File copying completed.")
    
    # ----------------------------------------------------------

if __name__ == '__main__':
    main()