doczyai-pipelines/ops_scripts/CNC/s3_search_and_copy_bucket.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting up imports that we will need, particularly - csv, boto3, ThreadPoolExecutor and pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import boto3\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from botocore.exceptions import ClientError\n",
    "import pandas as pd\n",
    "\n",
    "s3 = boto3.Session(profile_name='temp_cred').client('s3')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function use pagination to iterate through all objects in the specified s3 bucket and builds a file cache with all s3 keys. This is meant to be run only-once and is expected to take 8-10 mins to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_all_files(bucket_name):\n",
    "    all_keys = []\n",
    "    try:\n",
    "        paginator = s3.get_paginator('list_objects_v2')\n",
    "        for page in paginator.paginate(Bucket=bucket_name):\n",
    "            if 'Contents' in page:\n",
    "                all_keys.extend([item['Key'] for item in page['Contents']])\n",
    "    except ClientError as e:\n",
    "        print(f\"Error listing files in bucket {bucket_name}: {e}\")\n",
    "    return all_keys\n",
    "\n",
    "bucket_name = 'centene-national-contracting-files'\n",
    "all_keys = list_all_files(bucket_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Searching v1 - using complete matching of the filenames in the s3 bucket. If ```copy``` is set to ```True```, then the files found will also be copied to the specified prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CSV processing completed.\n"
     ]
    }
   ],
   "source": [
    "# Complete matching\n",
    "def check_file_in_bucket(filename, all_keys):\n",
    "    locations = [key for key in all_keys if (key.lower().endswith(filename.lower()))]\n",
    "    if copy:\n",
    "        for l in locations:\n",
    "            s3.copy_object(\n",
    "                CopySource={'Bucket': bucket_name, 'Key': l},\n",
    "                Bucket=bucket_name,\n",
    "                Key=destination_prefix + \"/\" + filename\n",
    "            )\n",
    "    exists = bool(locations)\n",
    "    return filename, exists, locations\n",
    "\n",
    "def process_csv(input_csv, output_csv):\n",
    "    \n",
    "    with open(input_csv, mode='r', newline='', encoding='utf8') as infile, open(output_csv, mode='w', newline='', encoding='utf8') as outfile:\n",
    "        reader = csv.reader(infile)\n",
    "        writer = csv.writer(outfile)\n",
    "        writer.writerow(['File Name', 'exists', 'locations'])  # Write header\n",
    "\n",
    "        with ThreadPoolExecutor(max_workers = 50) as executor:\n",
    "            futures = [executor.submit(check_file_in_bucket, row[0][:-4] + file_extension, all_keys) for row in reader]\n",
    "            \n",
    "            for future in futures:\n",
    "                filename, exists, locations = future.result()\n",
    "                writer.writerow([filename, exists, locations])\n",
    "\n",
    "# Input csv with all file names listed without extension in the first column\n",
    "input_csv = '../CNC/batch4/batch4_missing_files_2.csv'\n",
    "# Path to output that will be generated with the following columns = File Name, exists, locations\n",
    "output_csv = '../CNC/batch4/batch4_missing_files_2_paths_txt.csv'\n",
    "# If copy = True, the files found will be copied over to this destination in s3 (User Keys needs copy permission)\n",
    "destination_prefix = 'batch_4_priority_files/hotfix_diff_121124/TXT_FILES'\n",
    "copy = True\n",
    "# File Extension that will be added to the end of the file names - to look for\n",
    "file_extension = '.txt'\n",
    "\n",
    "process_csv(input_csv, output_csv)\n",
    "print(\"CSV processing completed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Searching v2 - using partial matching of the filenames in the s3 bucket. In this specific case, we check the first 10 and last 10 characters of the filename (This logic can be changed in the 3rd line, as required). If ```copy``` is set to ```True```, then the files found will also be copied to the specified prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Partial matching\n",
    "def check_file_in_bucket(filename, all_keys):\n",
    "    locations = [key for key in all_keys if (filename[:10].lower() in key.lower() and key.lower().endswith(filename[-10:].lower()))]\n",
    "    if copy:\n",
    "        for l in locations:\n",
    "            s3.copy_object(\n",
    "                CopySource={'Bucket': bucket_name, 'Key': l},\n",
    "                Bucket=bucket_name,\n",
    "                Key=destination_prefix + \"/\" + filename\n",
    "            )\n",
    "    exists = bool(locations)\n",
    "    return filename, exists, locations\n",
    "\n",
    "def process_csv(input_csv, output_csv):\n",
    "    \n",
    "    with open(input_csv, mode='r', newline='', encoding='utf8') as infile, open(output_csv, mode='w', newline='', encoding='utf8') as outfile:\n",
    "        reader = csv.reader(infile)\n",
    "        writer = csv.writer(outfile)\n",
    "        writer.writerow(['File Name', 'exists', 'locations'])  # Write header\n",
    "\n",
    "        with ThreadPoolExecutor(max_workers = 50) as executor:\n",
    "            futures = [executor.submit(check_file_in_bucket, row[0][:-4] + file_extension, all_keys) for row in reader]\n",
    "            \n",
    "            for future in futures:\n",
    "                filename, exists, locations = future.result()\n",
    "                writer.writerow([filename, exists, locations])\n",
    "\n",
    "# Input csv with all file names listed without extension in the first column\n",
    "input_csv = '../CNC/batch4/batch4_missing_files_2.csv'\n",
    "# Path to output that will be generated with the following columns = File Name, exists, locations\n",
    "output_csv = '../CNC/batch4/batch4_missing_files_2_paths_txt.csv'\n",
    "# If copy = True, the files found will be copied over to this destination in s3 (User Keys needs copy permission)\n",
    "destination_prefix = 'batch_4_priority_files/hotfix_diff_121124/TXT_FILES'\n",
    "copy = True\n",
    "# File Extension that will be added to the end of the file names - to look for\n",
    "file_extension = '.txt'\n",
    "\n",
    "process_csv(input_csv, output_csv)\n",
    "print(\"CSV processing completed.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}