from fastapi import APIRouter, HTTPException from app.salary_analytics.services.main import SalaryAnalyticsPipeline from app.salary_analytics.helpers.response_helpers import AnalysisResponse, BatchResponse from app.salary_analytics.helpers.data_checks import check_data_loaded from app.salary_analytics.services.data_loader import DataLoader from app.salary_analytics.core.state import state from app.models.db_operations import DatabaseOperations from app.config import OUTPUT_PATHS, TABLE_NAME from app.utils.logger import logger from typing import Optional, List, Union from sqlalchemy import text from datetime import datetime import pandas as pd, os, tempfile, time from typing import Optional, Union from fastapi import UploadFile, File router = APIRouter() @router.post("/run/pipeline", response_model=AnalysisResponse) async def run_full_pipeline(): """Run the complete salary analytics pipeline.""" start_time = time.time() try: check_data_loaded() logger.info("Starting full pipeline...") success = state.pipeline.run_full_pipeline() if not success: logger.error("Pipeline failed") logger.info(f"Full pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=500, detail="Pipeline failed") logger.info("Pipeline completed successfully") response = AnalysisResponse( message="Pipeline completed successfully" ) logger.info(f"Full pipeline endpoint completed in {time.time() - start_time:.2f} seconds") return response except Exception as e: logger.error(f"Error in pipeline: {str(e)}") logger.info(f"Full pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=500, detail=str(e)) @router.post("/run/streaming-pipeline", response_model=List[BatchResponse]) async def run_streaming_pipeline( source: str = "db", batch_size: int = 10000, file: Optional[Union[UploadFile, str]] = File(None) ): """ Run the complete salary analytics pipeline in batches. Args: source (str): Source of data ('db' or 'csv') batch_size (int): Number of rows to process in each batch file (UploadFile, optional): CSV file to load (required if source is 'csv') Returns: List[BatchResponse]: List of responses for each batch processed """ start_time = time.time() try: if source not in ['db', 'csv']: logger.error(f"Invalid source: {source}") logger.info(f"Streaming pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=400, detail="Source must be either 'db' or 'csv'") if source == 'csv' and not file: logger.error("No file provided for CSV source") logger.info(f"Streaming pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=400, detail="File must be provided when loading from CSV") # Initialize data loader state.data_loader = DataLoader() state.data_loader.chunk_size = batch_size # Create output directory for batch results timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") batch_output_dir = os.path.join(os.path.dirname(OUTPUT_PATHS['final_table']), f"batch_results_{timestamp}") os.makedirs(batch_output_dir, exist_ok=True) # Initialize database operations if not state.data_loader.connect(): logger.error("Failed to connect to database") logger.info(f"Streaming pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=500, detail="Failed to connect to database") db_ops = DatabaseOperations(state.data_loader.engine) if not db_ops.create_batch_results_table(): logger.error("Failed to create batch results table") logger.info(f"Streaming pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=500, detail="Failed to create batch results table") responses = [] batch_number = 0 batch_start_time = time.time() def preprocess_chunk(chunk): """Preprocess a chunk of data with the same logic as DataLoader.""" # Convert dates chunk['trx_start_date'] = pd.to_datetime(chunk['trx_start_date']) chunk['trx_end_date'] = pd.to_datetime(chunk['trx_end_date']) # Rename columns chunk = chunk.rename(columns={ 'd1': 'trx_type', 'd2': 'trx_subtype', 'd3': 'initiated_by', 'd4': 'customer_id' }) chunk = chunk.dropna() return chunk if source == 'csv': # Save uploaded file temporarily with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as temp_file: content = await file.read() temp_file.write(content) temp_file_path = temp_file.name try: # Process CSV in chunks for chunk in pd.read_csv(temp_file_path, chunksize=batch_size): batch_number += 1 logger.info(f"Processing batch {batch_number}") # Preprocess chunk chunk = preprocess_chunk(chunk) # Run pipeline on chunk state.pipeline = SalaryAnalyticsPipeline() state.pipeline.df = chunk try: batch_start_time = time.time() # Run analyses state.pipeline.run_keyword_analysis() state.pipeline.run_consistent_amount_analysis() state.pipeline.run_transaction_type_analysis() # Generate reports reports = state.pipeline.generate_salary_earner_reports() # Add batch metadata to results results_df = reports['final_table'].copy() results_df['batch_number'] = batch_number results_df['total_batches'] = -1 # Unknown for CSV results_df['processed_at'] = datetime.now() # Save batch results to CSV batch_results_path = os.path.join(batch_output_dir, f"batch_{batch_number}_results.csv") results_df.to_csv(batch_results_path, index=False) # Save to database db_ops.save_batch_to_db( batch_number=batch_number, total_batches=-1, # Unknown for CSV results_df=results_df, status="success" ) logger.info(f"Batch {batch_number} processed in {time.time() - batch_start_time:.2f} seconds") responses.append(BatchResponse( batch_number=batch_number, total_batches=-1, # Unknown for CSV processed_rows=len(chunk), results_path=batch_results_path, message=f"Successfully processed batch {batch_number}" )) except Exception as e: error_message = str(e) logger.error(f"Error processing batch {batch_number}: {error_message}") # Save error to database db_ops.save_batch_to_db( batch_number=batch_number, total_batches=-1, results_df=pd.DataFrame(), # Empty DataFrame for error case status="error" ) responses.append(BatchResponse( batch_number=batch_number, total_batches=-1, processed_rows=len(chunk), results_path="", message=f"Error processing batch {batch_number}: {error_message}" )) finally: # Clean up temporary file os.unlink(temp_file_path) else: # Process database in chunks if not state.data_loader.connect(): raise HTTPException(status_code=500, detail="Failed to connect to database") # Get total row count with state.data_loader.engine.connect() as conn: count_query = text(f"SELECT COUNT(*) FROM {TABLE_NAME}") total_rows = conn.execute(count_query).scalar() total_batches = (total_rows + batch_size - 1) // batch_size offset = 0 while offset < total_rows: batch_number += 1 logger.info(f"Processing batch {batch_number} of {total_batches}") # Load chunk from database query = f"SELECT * FROM {TABLE_NAME} LIMIT {batch_size} OFFSET {offset}" chunk = pd.read_sql(query, state.data_loader.engine) if chunk.empty: break # Preprocess chunk chunk = preprocess_chunk(chunk) # Run pipeline on chunk pipeline = SalaryAnalyticsPipeline() state.pipeline.df = chunk try: batch_start_time = time.time() # Run analyses state.pipeline.run_keyword_analysis() state.pipeline.run_consistent_amount_analysis() state.pipeline.run_transaction_type_analysis() # Generate reports reports = state.pipeline.generate_salary_earner_reports() # Add batch metadata to results results_df = reports['final_table'].copy() results_df['batch_number'] = batch_number results_df['total_batches'] = total_batches results_df['processed_at'] = datetime.now() # Save batch results to CSV batch_results_path = os.path.join(batch_output_dir, f"batch_{batch_number}_results.csv") results_df.to_csv(batch_results_path, index=False) # Save to database db_ops.save_batch_to_db( batch_number=batch_number, total_batches=total_batches, results_df=results_df, status="success" ) logger.info(f"Batch {batch_number} of {total_batches} processed in {time.time() - batch_start_time:.2f} seconds") responses.append(BatchResponse( batch_number=batch_number, total_batches=total_batches, processed_rows=len(chunk), results_path=batch_results_path, message=f"Successfully processed batch {batch_number} of {total_batches}" )) except Exception as e: error_message = str(e) logger.error(f"Error processing batch {batch_number}: {error_message}") # Save error to database db_ops.save_batch_to_db( batch_number=batch_number, total_batches=total_batches, results_df=pd.DataFrame(), # Empty DataFrame for error case status="error" ) responses.append(BatchResponse( batch_number=batch_number, total_batches=total_batches, processed_rows=len(chunk), results_path="", message=f"Error processing batch {batch_number}: {error_message}" )) offset += batch_size logger.info(f"Streaming pipeline endpoint completed in {time.time() - start_time:.2f} seconds") return responses except Exception as e: logger.error(f"Error in streaming pipeline: {str(e)}") logger.info(f"Streaming pipeline endpoint failed after {time.time() - start_time:.2f} seconds") raise HTTPException(status_code=500, detail=str(e))