import os import json from supabase import create_client, Client import boto3 import pandas as pd from io import BytesIO from datetime import datetime, timedelta from typing import Dict, List, Any, Optional import logging # Configure logging logger = logging.getLogger() logger.setLevel(logging.INFO) # Initialize AWS SSM client ssm = boto3.client('ssm') def get_parameter(name: str) -> str: """Get a parameter from AWS SSM Parameter Store""" try: response = ssm.get_parameter( Name=name, WithDecryption=True ) return response['Parameter']['Value'] except Exception as e: logger.error(f"Error getting parameter {name}: {str(e)}") raise def get_supabase_client() -> Client: """Initialize and return Supabase client""" try: url: str = get_parameter('/supabase/url') key: str = get_parameter('/supabase/anon') supabase: Client = create_client(url, key) supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"}) return supabase except Exception as e: logger.error(f"Error initializing Supabase client: {str(e)}") raise def get_unprocessed_upload(supabase: Client) -> Optional[Dict]: """Get the next unprocessed upload record""" try: upload_response = supabase.table("sjreport_uploads").select("*").eq("processed", "FALSE").limit(1).execute() return upload_response.data[0] if upload_response.data else None except Exception as e: logger.error(f"Error getting unprocessed upload: {str(e)}") raise def download_file(supabase: Client, bucket_name: str, file_name: str) -> bytes: """Download file from Supabase storage""" try: response = supabase.storage.from_(bucket_name).download(file_name) return response except Exception as e: logger.error(f"Error downloading file {file_name}: {str(e)}") raise def process_excel_file(file_content: bytes) -> Dict[str, pd.DataFrame]: """Process Excel file and return dictionary of DataFrames""" try: # Load only visible sheets; ignore hidden/veryHidden sheets try: from openpyxl import load_workbook # Lazy import to reduce cold start wb = load_workbook(BytesIO(file_content), read_only=True, data_only=True) visible_sheet_names = [ ws.title for ws in wb.worksheets if getattr(ws, "sheet_state", "visible") == "visible" ] if visible_sheet_names: sheets = pd.read_excel(BytesIO(file_content), sheet_name=visible_sheet_names, dtype=str) else: sheets = {} except Exception: # Fallback: load all sheets if workbook inspection fails sheets = pd.read_excel(BytesIO(file_content), sheet_name=None, dtype=str) for sheet_name in list(sheets.keys()): sheets[sheet_name].columns = sheets[sheet_name].columns.str.lower().str.replace(' ', '_') return sheets except Exception as e: logger.error(f"Error processing Excel file: {str(e)}") raise def process_row(row: pd.Series, url: str, sheet_name: str) -> Dict: """Process a single row of data""" try: row_data = row.to_dict() modified_row_data = {} # Handle null values for col, value in row_data.items(): modified_row_data[col] = "" if pd.isna(value) else value # Process SPAR store codes if "SPAR" in modified_row_data.get("channel", ""): modified_row_data["store_code"] = "S" + modified_row_data["store_code"] # Process date date_string = modified_row_data["date"][:10] date = datetime.strptime(date_string, '%Y-%m-%d').date() week_start_date = date + timedelta(days=(7 - date.weekday())) # Add metadata modified_row_data.update({ 'row': row.name + 1, 'week_start_date': week_start_date.strftime('%Y-%m-%d'), 'file': url, 'sheetname': sheet_name }) return modified_row_data except Exception as e: logger.error(f"Error processing row: {str(e)}") raise def insert_batch(supabase: Client, table_name: str, batch: List[Dict]) -> None: """Insert a batch of records into the database""" try: response = supabase.table(table_name).upsert(batch, on_conflict='row, week_start_date, file, sheetname').execute() if hasattr(response, 'error') and response.error: raise Exception(response.error.message) logger.info(f"Successfully inserted {len(batch)} records") except Exception as e: logger.error(f"Error inserting batch: {str(e)}") raise def update_upload_status(supabase: Client, record_id: str, error: str = '') -> None: """Update the upload record status""" try: supabase.table("sjreport_uploads").update({ "processed": True, "error": error }).eq("id", record_id).execute() except Exception as e: logger.error(f"Error updating upload status: {str(e)}") raise def invoke_create_tasks(lambda_client: boto3.client, week_start_date: str) -> None: """Invoke the create tasks Lambda function""" try: payload = {"week_start_date": week_start_date} lambda_client.invoke( FunctionName='datafyNew-CreateTasks-ccMR48YAjXLP', InvocationType='Event', Payload=json.dumps(payload) ) except Exception as e: logger.error(f"Error invoking create tasks: {str(e)}") raise def lambda_handler(event, context): """Main Lambda handler function""" try: # Initialize clients lambda_client = boto3.client('lambda') supabase = get_supabase_client() # Get unprocessed upload record = get_unprocessed_upload(supabase) if not record: return { 'statusCode': 400, 'body': json.dumps({'error': 'No unprocessed uploads found'}) } record_id = record['id'] bucket_name = "public/" + record['payload']['bucket_name'] file_name = record['payload']['file_name'] table_name = record['payload']['path_tokens'][0] # Download and process file file_content = download_file(supabase, bucket_name, file_name) sheets = process_excel_file(file_content) # Process each sheet last_week_start_date = None BATCH_SIZE = 500 # Reduced batch size for better memory management for sheet_name, sheet_df in sheets.items(): logger.info(f"Processing sheet: {sheet_name}") data_to_insert = [] for _, row in sheet_df.iterrows(): try: processed_row = process_row(row, file_name, sheet_name) data_to_insert.append(processed_row) last_week_start_date = processed_row['week_start_date'] # Insert in smaller batches if len(data_to_insert) >= BATCH_SIZE: insert_batch(supabase, table_name, data_to_insert) data_to_insert = [] except Exception as e: logger.error(f"Error processing row in sheet {sheet_name}: {str(e)}") update_upload_status(supabase, record_id, f"Error processing row: {str(e)}") return { 'statusCode': 400, 'body': json.dumps({'error': str(e)}) } # Insert remaining records if data_to_insert: insert_batch(supabase, table_name, data_to_insert) supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"}) supabase.rpc("update_product_table").execute() # Update product table and create tasks if last_week_start_date: invoke_create_tasks(lambda_client, last_week_start_date) # Mark upload as processed update_upload_status(supabase, record_id) return { 'statusCode': 200, 'body': json.dumps({'message': 'Data processed successfully'}) } except Exception as e: logger.error(f"Error in lambda_handler: {str(e)}") return { 'statusCode': 500, 'body': json.dumps({'error': str(e)}) }