Files
doczyai-pipelines/ops_scripts/generic/t_drive_search.py
T
Umang Shailesh Mistry c210052952 Merged in feature/ops_scripts (pull request #333)
Feature/ops scripts

* Added comments to Aryan's script and added some more scripts

* Search and Copy Script uploaded as a Python Notebook - with comments and markdown

* Merged main into feature/ops_scripts

* Merged main into feature/ops_scripts


Approved-by: Michael McGuinness
Approved-by: Chris Stobie
2024-12-17 11:18:41 +00:00

187 lines
7.2 KiB
Python

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()