import io import pandas as pd from typing import Dict, List, Tuple, Any from datetime import datetime, date from utils import get_cloudwatch_logger, execute_query, engine import numpy as np import uuid import threading from scripts.mysql_queries import ( insert_file_upload_record_query, get_user_name_by_id_query, insert_pincode_query, soft_delete_pincode_query, cleanup_old_pincode_batches_query ) STREAM_NAME = 'pincode_master_data' logger = get_cloudwatch_logger(stream_name=STREAM_NAME) class PincodeMasterProcessor: EXPECTED_COLUMNS = [ "pincode", "type", "state", "city", "comment", "agency_name" ] def __init__(self, user_id: str, file_name: str): self.user_id = user_id self.df = None self.file_name = file_name def log_error(self, message: str): full_msg = f"[User: {self.user_id}] Error: {message}" print(full_msg) try: logger.send_log_event(full_msg) except Exception as log_err: print(f"Failed to send log event: {log_err}") def load_file(self, file_content: bytes) -> Tuple[bool, str]: try: if not file_content: return False, "File content is empty." self.df = pd.read_excel(io.BytesIO(file_content), engine="calamine") return True, "File loaded successfully." except Exception as e: return False, f"Failed to parse Excel file: {str(e)}" def validate_columns(self) -> Tuple[bool, str]: if self.df is None: return False, "DataFrame not loaded." self.df.columns = [str(col).strip().lower() for col in self.df.columns] self.df.rename(columns={"agencyname": "agency_name"}, inplace=True) # Check for mandatory columns mandatory = ["pincode", "type"] uploaded_columns = set(self.df.columns) missing = set(mandatory) - uploaded_columns if missing: return False, f"Missing mandatory columns: {', '.join(missing)}" return True, "Column structure valid." def get_uploader_name(self) -> str: try: query = get_user_name_by_id_query.replace("%s", str(self.user_id)) df_result = execute_query(query) if df_result is not None and not df_result.empty: return df_result.iloc[0]["name"] return "Unknown" except Exception as e: self.log_error(f"Could not fetch uploader name: {e}") return "System" def validate_data_integrity(self) -> Tuple[bool, Any]: errors = [] # 1. Standardize Data Types self.df["pincode_clean"] = self.df["pincode"].astype(str).str.strip().replace('nan', '') self.df["type_clean"] = self.df["type"].astype(str).str.strip().str.upper() # Mandatory missing mask_pincode_missing = self.df["pincode_clean"] == "" mask_type_missing = self.df["type_clean"] == "NAN" # Valid Types valid_types = {"NEGATIVE": 0, "POSITIVE": 1, "RISKY": 2, "0": 0, "1": 1, "2": 2} mask_invalid_type = (~mask_type_missing) & (~self.df["type_clean"].isin(valid_types.keys())) def add_errors(mask, message): if mask.any(): bad_rows = mask[mask].index + 2 for r in bad_rows: errors.append(f"Row {r}: {message}") add_errors(mask_pincode_missing, "Missing 'pincode'.") add_errors(mask_type_missing, "Missing 'type'.") add_errors(mask_invalid_type, "Invalid 'type'. Must be POSITIVE, NEGATIVE, or RISKY.") # Duplicate pincode validation duplicate_mask = self.df.duplicated(subset=['pincode_clean'], keep=False) & (~mask_pincode_missing) add_errors(duplicate_mask, "Duplicate 'pincode' found in the file.") # Map type strings to integers self.df["type_mapped"] = self.df["type_clean"].map(valid_types) if errors: return False, errors return True, "Data integrity checks passed." def save_to_database(self) -> int: records_count = len(self.df) print(f"--- Inserting {records_count} records into DB ---") uploader_name = self.get_uploader_name() current_date = date.today() current_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") db_df = pd.DataFrame() db_df["pincode"] = self.df["pincode_clean"] db_df["type"] = self.df["type_mapped"] # Optional fields def get_optional(col): if col in self.df.columns: return self.df[col].where(self.df[col].notna(), None) return None db_df["state"] = get_optional("state") db_df["city"] = get_optional("city") db_df["area"] = get_optional("area") db_df["location"] = get_optional("location") db_df["comment"] = get_optional("comment") db_df["agency_name"] = get_optional("agency_name") # Level defaults to 1 for NEGATIVE types, else None if "level" in self.df.columns: db_df["level"] = pd.to_numeric(self.df["level"], errors='coerce') db_df["level"] = np.where(db_df["type"] == 0, db_df["level"].fillna(1), db_df["level"]) db_df["level"] = db_df["level"].replace({np.nan: None}) else: db_df["level"] = np.where(db_df["type"] == 0, 1, None) batch_id = str(uuid.uuid4()) db_df["upload_batch_id"] = batch_id db_df["user_id"] = self.user_id db_df["created_at"] = current_timestamp db_df["updated_at"] = current_timestamp # First truncate/delete old records since we are doing full replacement execute_query(soft_delete_pincode_query) if not db_df.empty: db_df.to_sql( name="pincode_classifications", con=engine, if_exists="append", index=False, chunksize=1000, method="multi" ) file_upload_params = [ { "file_name": self.file_name, "uploaded_by": uploader_name, "upload_date": current_date, "total_rows": records_count, "created_at": current_timestamp, } ] execute_query(insert_file_upload_record_query, file_upload_params) # Run hard delete of older batches asynchronously def async_cleanup(): try: print(f"--- Running background cleanup for old pincode batches ---") execute_query(cleanup_old_pincode_batches_query) print(f"--- Background cleanup complete ---") except Exception as e: self.log_error(f"Failed to cleanup old batches: {str(e)}") threading.Thread(target=async_cleanup, daemon=True).start() return records_count def run(self, file_content: bytes) -> Dict[str, Any]: print(f"--- Processing Started for User: {self.user_id} ---") success, msg = self.load_file(file_content) if not success: self.log_error(msg) return {"status": "error", "message": msg} success, msg = self.validate_columns() if not success: self.log_error(msg) return {"status": "error", "message": msg} success, msg = self.validate_data_integrity() if not success: self.log_error(msg) return {"status": "error", "message": msg} try: count = self.save_to_database() return { "status": "success", "message": "File processed and data inserted successfully.", "processed_records": count } except Exception as e: err_msg = f"Database Insertion Failed: {str(e)}" self.log_error(err_msg) return {"status": "error", "message": err_msg} def process_pincode_master_excel(file_content: bytes, user_id: str, file_name: str) -> Dict[str, Any]: processor = PincodeMasterProcessor(user_id, file_name) return processor.run(file_content)