#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Base Automation Class Provides common workflow patterns, standardized result handling, and progress reporting """ import os import sys import csv import time import logging import datetime import io from typing import Dict, Any, Optional, List, Callable from pathlib import Path from abc import ABC, abstractmethod # Import shared utilities try: from .electron_browser import ElectronBrowser, ElectronWait, By, EC from .config import get_config from .error_handler import RetryHandler, ErrorClassifier, get_error_suggestion from .session_manager import SessionManager from .validators import CSVValidator, ValidationError except ImportError: # Fallback for direct execution import sys import os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from automation_scripts.shared.electron_browser import ElectronBrowser, ElectronWait, By, EC from automation_scripts.shared.config import get_config from automation_scripts.shared.error_handler import RetryHandler, ErrorClassifier, get_error_suggestion from automation_scripts.shared.session_manager import SessionManager from automation_scripts.shared.validators import CSVValidator, ValidationError class BaseAutomation(ABC): """Base class for all automation scripts.""" def __init__( self, service_name: str, csv_file: str, output_dir: str, job_id: str, config_file: Optional[str] = None, api_base_url: Optional[str] = None, user_id: Optional[str] = None, credits_used: float = 0.0 ): """Initialize base automation. Args: service_name: Name of the service/automation csv_file: Path to input CSV file output_dir: Base output directory job_id: Unique job identifier config_file: Optional path to config file api_base_url: Optional API base URL for server communication user_id: Optional user ID for server communication credits_used: Credits used for this job """ self.service_name = service_name self.csv_file = csv_file self.base_output_dir = output_dir self.job_id = job_id self.api_base_url = api_base_url self.user_id = user_id self.credits_used = credits_used # Configuration self.config = get_config(config_file) self.service_config = self.config.get_service_config(service_name) self.browser_config = self.config.get_browser_config() self.download_config = self.config.get_download_config() self.session_config = self.config.get_session_config() self.error_config = self.config.get_error_config() # Setup directories self.output_dir = os.path.join(output_dir, f"{service_name}_{job_id}") os.makedirs(self.output_dir, exist_ok=True) # Local log directory disabled; no log files will be written self.session_dir = os.path.join(self.output_dir, 'sessions', job_id) os.makedirs(self.session_dir, exist_ok=True) self.log_buffer = io.StringIO() # Setup logging self._setup_logging() # Initialize components self.retry_handler = RetryHandler(self.error_config) self.session_manager = SessionManager(self.session_config) self.csv_validator = CSVValidator(self.logger) # State self.driver: Optional[ElectronBrowser] = None self.results: List[Dict[str, Any]] = [] self.success_count = 0 self.failed_count = 0 self.start_time: Optional[datetime.datetime] = None self.end_time: Optional[datetime.datetime] = None # Server client (if needed) self.server_client = None if api_base_url and user_id: try: try: from .server_client import ServerClient except ImportError: from automation_scripts.shared.server_client import ServerClient self.server_client = ServerClient(api_base_url, user_id, job_id, self.logger) except ImportError: if self.logger: self.logger.warning("ServerClient not available, server communication disabled") def _setup_logging(self): """Setup logging configuration.""" log_config = self.config.get_logging_config() log_level = getattr(logging, log_config.get('level', 'INFO')) log_format = log_config.get('format', '%(asctime)s [%(levelname)s] %(message)s') self.logger = logging.getLogger(self.service_name) self.logger.setLevel(log_level) # Clear existing handlers self.logger.handlers.clear() buffer_handler = logging.StreamHandler(self.log_buffer) buffer_handler.setLevel(log_level) buffer_handler.setFormatter(logging.Formatter(log_format)) self.logger.addHandler(buffer_handler) # Console handler if log_config.get('console_output', True): console_handler = logging.StreamHandler(sys.stdout) console_handler.setLevel(log_level) console_handler.setFormatter(logging.Formatter(log_format)) self.logger.addHandler(console_handler) def validate_input(self, required_columns: Optional[List[str]] = None, required_fields: Optional[List[str]] = None) -> Dict[str, Any]: """Validate input CSV file. Args: required_columns: List of required column names required_fields: List of required field names per row Returns: Validation result dictionary """ try: result = self.csv_validator.validate_complete( self.csv_file, required_columns=required_columns, required_fields=required_fields ) self.logger.info(f"✓ CSV validation passed: {result['row_count']} rows, {len(result['columns'])} columns") return result except ValidationError as e: self.logger.error(f"❌ CSV validation failed: {e}") raise def setup_driver(self, headless: Optional[bool] = None, use_session: bool = True) -> ElectronBrowser: """Setup Electron browser. Args: headless: Override headless mode (uses config if None) use_session: Whether to use session persistence (unused for Electron but kept for API parity) Returns: WebDriver instance """ setup_start_time = time.time() self.logger.info("🚀 Starting Electron browser setup...") if headless is None: headless = self.browser_config.get('headless', True) self.logger.info(f"🔧 Configuration:") self.logger.info(f" - Headless mode: {'ENABLED' if headless else 'DISABLED'}") self.logger.info("⚡ Creating Electron browser...") self.driver = ElectronBrowser.create( job_id=self.job_id, headless=headless, logger=self.logger if hasattr(self, 'logger') else None # compat ) setup_duration = time.time() - setup_start_time self.logger.info(f"✅ Electron browser setup complete in {setup_duration:.2f} seconds") return self.driver def wait_for_element(self, by, value, timeout: Optional[int] = None, retry_count: int = 3): """Wait for an element to be present with retry. Args: by: Selenium By locator value: Locator value timeout: Optional timeout override retry_count: Number of retry attempts Returns: WebElement or None """ if timeout is None: timeout = self.browser_config.get('explicit_wait', 15) wait = ElectronWait(self.driver, timeout=timeout) for attempt in range(retry_count): try: return wait.until(EC.presence_of_element_located((by, value))) except Exception: if attempt == retry_count - 1: self.logger.warning(f"Timeout waiting for element: {by}={value}") return None time.sleep(0.5) return None def wait_for_clickable(self, by, value, timeout: Optional[int] = None, retry_count: int = 3): """Wait for an element to be clickable with retry. Args: by: Selenium By locator value: Locator value timeout: Optional timeout override retry_count: Number of retry attempts Returns: WebElement or None """ if timeout is None: timeout = self.browser_config.get('explicit_wait', 15) wait = ElectronWait(self.driver, timeout=timeout) for attempt in range(retry_count): try: return wait.until(EC.element_to_be_clickable((by, value))) except Exception: if attempt == retry_count - 1: self.logger.warning(f"Timeout waiting for clickable element: {by}={value}") return None time.sleep(0.5) return None def wait_for_page_load(self, timeout: Optional[int] = None): """Wait for page to be ready. Args: timeout: Optional timeout override """ if timeout is None: timeout = self.browser_config.get('page_load_timeout', 30) try: wait = ElectronWait(self.driver, timeout=timeout) wait.until(lambda d: d.execute_script("return document.readyState") == "complete") except Exception: pass def check_stop_signal(self) -> bool: """Check if stop signal exists. Returns: True if stop signal found, False otherwise """ stop_file = os.path.join(self.output_dir, f"stop_{self.job_id}.txt") if os.path.exists(stop_file): self.logger.info("🛑 Stop signal detected - initiating cleanup...") self._cleanup_on_stop() return True return False def _cleanup_on_stop(self): """Clean up all generated files and directories when automation is stopped.""" try: import shutil self.logger.info("🗑️ Cleaning up generated files and directories...") # List of files/directories to clean up cleanup_items = [] # 1. Result CSV file result_csv = os.path.join(self.output_dir, f"{self.service_name}_results_{self.job_id}.csv") if os.path.exists(result_csv): cleanup_items.append(result_csv) # 2. Log files (if any were created locally) log_file = os.path.join(self.output_dir, f"{self.service_name}_{self.job_id}.log") if os.path.exists(log_file): cleanup_items.append(log_file) # 3. Session directory if os.path.exists(self.session_dir): cleanup_items.append(self.session_dir) # 4. Any downloaded files or PDFs in output directory if os.path.exists(self.output_dir): for item in os.listdir(self.output_dir): item_path = os.path.join(self.output_dir, item) # Skip the stop signal file itself if not item.startswith(f"stop_{self.job_id}"): cleanup_items.append(item_path) # Delete all items deleted_count = 0 for item in cleanup_items: try: if os.path.isfile(item): os.remove(item) deleted_count += 1 self.logger.info(f" ✓ Deleted file: {os.path.basename(item)}") elif os.path.isdir(item): shutil.rmtree(item) deleted_count += 1 self.logger.info(f" ✓ Deleted directory: {os.path.basename(item)}") except Exception as e: self.logger.warning(f" ⚠️ Failed to delete {item}: {e}") # Finally, remove the entire output directory try: if os.path.exists(self.output_dir): shutil.rmtree(self.output_dir) self.logger.info(f"✅ Cleanup complete: Removed output directory and {deleted_count} items") except Exception as e: self.logger.warning(f"⚠️ Failed to remove output directory: {e}") except Exception as e: self.logger.error(f"❌ Error during cleanup: {e}") def write_result(self, row_data: Dict[str, str], success: bool, message: str): """Write result to CSV file. Args: row_data: Dictionary containing row data success: Boolean indicating success or failure message: Status message to write """ if self.check_stop_signal(): return result_csv = os.path.join(self.output_dir, f"{self.service_name}_results_{self.job_id}.csv") file_exists = os.path.exists(result_csv) with open(result_csv, "a", newline='', encoding='utf-8') as f: writer = csv.writer(f) if not file_exists: # Write header - use row_data keys plus status columns header = list(row_data.keys()) + ["Status", "Message", "Timestamp"] writer.writerow(header) row = [row_data.get(key, "N/A") for key in row_data.keys()] row.extend([ "✅ Success" if success else "❌ Failed", message, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") ]) writer.writerow(row) def record_result(self, row_data: Dict[str, str], success: bool, message: str, result_data: Optional[Dict[str, Any]] = None): """Record a processing result. Args: row_data: Original row data success: Whether processing was successful message: Status message result_data: Optional additional result data """ if success: self.success_count += 1 else: self.failed_count += 1 result = { 'row_data': row_data, 'success': success, 'message': message, 'timestamp': datetime.datetime.now().isoformat() } if result_data: result.update(result_data) self.results.append(result) self.write_result(row_data, success, message) def log_progress(self, current: int, total: int, item_name: str = "item"): """Log progress information. Args: current: Current item number total: Total number of items item_name: Name of item being processed """ items_remaining = total - current + 1 percentage = (current / total * 100) if total > 0 else 0 self.logger.info(f"📊 Progress: {current}/{total} ({percentage:.1f}%) - {items_remaining} {item_name}(s) remaining") self.logger.info(f" Success: {self.success_count}, Failed: {self.failed_count}") @abstractmethod def login(self, driver: ElectronBrowser) -> bool: """Perform login. Must be implemented by subclasses. Args: driver: WebDriver instance Returns: True if login successful, False otherwise """ pass @abstractmethod def process_row(self, driver: ElectronBrowser, row: Dict[str, str], index: int, total: int) -> bool: """Process a single CSV row. Must be implemented by subclasses. Args: driver: WebDriver instance row: CSV row data index: Row index (1-based) total: Total number of rows Returns: True if processing successful, False otherwise """ pass def run(self) -> bool: """Run the automation. Returns: True if automation completed successfully, False otherwise """ self.start_time = datetime.datetime.now() try: self.logger.info("=" * 80) self.logger.info(f"🚀 {self.service_name} Automation Started") self.logger.info("=" * 80) self.logger.info(f"📁 Input CSV: {self.csv_file}") self.logger.info(f"📂 Output Directory: {self.output_dir}") self.logger.info(f"🆔 Job ID: {self.job_id}") self.logger.info("=" * 80) # Validate input validation_result = self.validate_input() total_rows = validation_result['row_count'] # Setup driver self.setup_driver() # Login if not self.login(self.driver): self.logger.error("❌ Login failed") return False # Process CSV rows with open(self.csv_file, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) rows = [row for row in reader if any(row.values())] for idx, row in enumerate(rows, 1): # Check for stop signal if self.check_stop_signal(): self.logger.info("Stop signal received, stopping automation...") if self.driver: try: self.driver.quit() except: pass sys.exit(130) # Log progress self.log_progress(idx, len(rows)) # Process row try: success = self.process_row(self.driver, row, idx, len(rows)) if success: self.record_result(row, True, "Processing completed successfully") else: self.record_result(row, False, "Processing failed") except Exception as e: error_type = ErrorClassifier.classify(e) suggestion = get_error_suggestion(e) self.logger.error(f"❌ Error processing row {idx}: {e}") self.logger.error(f" Suggestion: {suggestion}") self.record_result(row, False, f"Error: {str(e)}") # Print summary self.end_time = datetime.datetime.now() duration = (self.end_time - self.start_time).total_seconds() self.logger.info("\n" + "=" * 80) self.logger.info("PROCESSING SUMMARY") self.logger.info("=" * 80) self.logger.info(f"Total entries: {total_rows}") self.logger.info(f"✓ Successful: {self.success_count}") self.logger.info(f"✗ Failed: {self.failed_count}") self.logger.info(f"⏱️ Duration: {duration:.2f} seconds") self.logger.info("=" * 80) # Upload results if server client available if self.server_client: self._upload_results() return self.failed_count == 0 except Exception as e: self.logger.error(f"❌ Fatal error in automation: {e}") import traceback traceback.print_exc() return False finally: self.cleanup() def _upload_results(self): """Upload results to server if server client is available.""" if not self.server_client: return try: job_data = { 'userId': self.user_id, 'serviceId': self.service_name, 'serviceName': self.service_name, 'inputFileName': os.path.basename(self.csv_file), 'inputFilePath': self.csv_file, 'creditsUsed': self.credits_used } uploaded_log_path = None uploaded_csv_path = None if self.log_buffer: log_text = self.log_buffer.getvalue() if log_text: uploaded_log_path = self.server_client.upload_log_buffer( log_text, f'{self.service_name}_{self.job_id}.log', job_data ) uploaded_csv_path = None if self.results: import io as _io import csv as _csv csv_io = _io.StringIO() writer = _csv.writer(csv_io) header_keys = list(self.results[0].get('row_data', {}).keys()) writer.writerow(header_keys + ['Status','Message','Timestamp']) for r in self.results: row_data = r.get('row_data', {}) row = [row_data.get(k, 'N/A') for k in header_keys] row.extend(['✅ Success' if r.get('success') else '❌ Failed', r.get('message',''), r.get('timestamp','')]) writer.writerow(row) uploaded_csv_path = self.server_client.upload_results_buffer( csv_io.getvalue().encode('utf-8'), f'{self.service_name}_results_{self.job_id}.csv', job_data ) result_files = [] if uploaded_csv_path: result_files.append(uploaded_csv_path) self.server_client.report_job_result( result_files, self.output_dir, successful_count=self.success_count ) except Exception as e: self.logger.warning(f"⚠️ Error uploading results to server: {e}") def cleanup(self): """Clean up resources.""" if self.driver: try: self.logger.info("Closing browser...") self.driver.quit() except Exception as e: self.logger.warning(f"Error closing browser: {e}") self.session_manager.cleanup_all()