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

import os
import re
import pandas as pd
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import partial
from tqdm import tqdm
import logging


"""
This script searches for files in a directory based on a list of file names in an Excel file.
The Excel file must contain a column with the file names to search for.

This is used mainly to find files that are missing in the batch staging files and cannot be located in the s3 bucket. 

Usually the flow of staging and prepping batches involve:
1. Staging the files in the s3 bucket
2. Running the excel_s3_diff.py script to compare the files in the s3 bucket with the files in the staging folder
3. The output of that script is an excel file that contains the missing files in the s3 bucket
4. This script is then used to search for the missing files in the staging folder
5. The output of this script is an excel file that contains the missing files in the staging folder
6. The output from this diff will have files names with encoding issue / no enocoding issue.    
    For files without any encoding issue usually are found in the t drive. We use the path to upoad it to a prefix in the s3 bucket for the client/batch_#
7. Then the staged pdfs are then processed with the scheduler to run the pdfs through the textract pipeline to extract the text from the pdfs.


"""



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

# Path to the Excel file
# EXCEL_FILE_PATH = "C:\\Doczy\\National contracting\\Allbatches\\diff_batch13.xlsx"  # <-- Update this path
EXCEL_FILE_PATH = 'C:\\Doczy\\National contracting\\Allbatches\\new_diff_batch14_3_modified.xlsx'

# Directory to search within
SEARCH_DIRECTORY = 'T:\\AArete Client Work\\Doczy-Production\\Restricted\\2024-06-28-pdf'      # <-- Update this path

# Output Excel file path
OUTPUT_FILE_PATH = 'C:\\Doczy\\National contracting\\Allbatches\\missing files analysis\\batch_14_search_output_3.xlsx'      # <-- Update this path

# Maximum search depth
MAX_DEPTH = 5

# Number of threads for multi-threading
NUM_THREADS = 100  # Start with 4 and adjust as needed

# Log file path
LOG_FILE_PATH = 'C:\\Doczy\\National contracting\\Allbatches\\missing files analysis\\batch_14_search_3.log'

# --------------------------- Logging Setup ---------------------------

logging.basicConfig(
    level=logging.INFO,  # Change to DEBUG for more detailed logs
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE_PATH),
        logging.StreamHandler()
    ]
)

# --------------------------- Helper Functions ---------------------------

def remove_suffix(file_name):
    """
    Removes suffixes like (1), (2), etc., from the file name.
    Example: "XYZ_DOCUMENT(1).pdf" -> "XYZ_DOCUMENT.pdf"
    """
    # Removing last 3 characters from the file name
    cleaned_name = file_name[:-3].strip() if file_name.endswith(('(1)', '(2)', '(3)', '(4)')) else file_name

    logging.debug(f"Removed suffix: '{file_name}' -> '{cleaned_name}'")
    return cleaned_name

def find_files(search_dir, target_name, max_depth, file_types):
    """
    Searches for files matching the target_name within search_dir up to max_depth.
    Returns a list of full file paths.
    """
    logging.info(f"Searching for '{target_name}' in '{search_dir}' with max depth {max_depth}")
    results = []
    target_name_lower = target_name.lower()
    search_dir = Path(search_dir).resolve()
    base_depth = len(search_dir.parts)

    for root, dirs, files in os.walk(search_dir, topdown=True):
        current_depth = len(Path(root).parts) - base_depth + 1
        if current_depth > max_depth:
            dirs[:] = []  # Prevent descending further
            logging.debug(f"Reached max depth at: {root}")
            continue

        for file in files:
            if Path(file).suffix.lower() in file_types:
                if file.lower() == target_name_lower:
                    file_path = Path(root) / file
                    logging.info(f"Found file: {file_path}")
                    results.append(file_path)

    logging.info(f"Search complete for '{target_name}'. Found {len(results)} file(s).")
    return results

def search_file(row, search_dir, max_depth, file_types):
    """
    Processes a single row to find the appropriate file path based on the Reason.
    """
    # file_name = row['File Name']
    # reason = row['Reason']
    file_name = row['Missing File Name']
    reason = row['Encoding Check']
    found_path = None
    
    # Remove suffix like (1), (2), etc.
    cleaned_file_name = remove_suffix(file_name)
    
    # Ensure the file has a .pdf extension
    
    cleaned_file_name += '.pdf'
    print(f"Cleaned file name: {cleaned_file_name}")

    # Perform a case-insensitive search for the cleaned file name
    found_files = find_files(search_dir, cleaned_file_name, max_depth, file_types)

    if found_files:
        # Select the file with the largest size if multiple are found
        largest_file = max(found_files, key=lambda f: f.stat().st_size)
        found_path = str(largest_file.resolve())
        logging.debug(f"Selected largest file: {found_path}")
    else:
        logging.warning(f"File '{cleaned_file_name}' not found for Reason: '{reason}'")

    return found_path

# --------------------------- Main Processing ---------------------------

def main():
    logging.info("Script started.")

    # Read the Excel file
    try:
        df = pd.read_excel(EXCEL_FILE_PATH, sheet_name='Sheet1', engine='openpyxl')
        logging.info(f"Excel file '{EXCEL_FILE_PATH}' read successfully.")
    except Exception as e:
        logging.error(f"Error reading Excel file: {e}")
        return

    # Ensure required columns exist
    # required_columns = {'File Name', 'Reason'}
    required_columns = {'Missing File Name', 'Encoding Check'}
    if not required_columns.issubset(df.columns):
        logging.error(f"Error: The Excel sheet must contain the following columns: {required_columns}")
        return

    # Initialize a new column for found paths
    df['Found Path'] = None

    # Prepare for multi-threaded processing
    file_types = {'.pdf'}
    search_partial = partial(search_file, search_dir=SEARCH_DIRECTORY, max_depth=MAX_DEPTH, file_types=file_types)

    with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
        # Submit all tasks
        futures = {executor.submit(search_partial, row): idx for idx, row in df.iterrows()}
        # Iterate through completed futures with a progress bar
        for future in tqdm(as_completed(futures), total=len(futures), desc="Processing"):
            idx = futures[future]
            try:
                found_path = future.result()
                df.at[idx, 'Found Path'] = found_path
                logging.debug(f"Row {idx} processed. Found Path: {found_path}")
            except Exception as e:
                df.at[idx, 'Found Path'] = f"Error: {e}"
                logging.error(f"Error processing row {idx}: {e}")

    # Save the results to a new Excel file
    try:
        df.to_excel(OUTPUT_FILE_PATH, index=False)
        logging.info(f"Processing complete. Results saved to '{OUTPUT_FILE_PATH}'")
    except Exception as e:
        logging.error(f"Error saving output Excel file: {e}")

    logging.info("Script finished.")

if __name__ == "__main__":
    main()