From f433da681db550d24c6f52f815b609932e3cb086 Mon Sep 17 00:00:00 2001 From: Abdullah Ceylan Date: Sat, 27 Sep 2025 19:40:36 +0300 Subject: [PATCH] Add error handling and logging --- __init__.py | 42 ++- api_routes.py | 165 +++++++--- database.py | 251 ++++++++------- error_handler.py | 467 +++++++++++++++++++++++++++ exceptions.py | 450 ++++++++++++++++++++++++++ health_check.py | 635 +++++++++++++++++++++++++++++++++++++ logging_config.py | 450 ++++++++++++++++++++++++++ performance_monitor.py | 545 +++++++++++++++++++++++++++++++ queue_service.py | 31 +- test_error_handling.py | 499 +++++++++++++++++++++++++++++ test_logging_monitoring.py | 602 +++++++++++++++++++++++++++++++++++ workflow_executor.py | 44 ++- 12 files changed, 3975 insertions(+), 206 deletions(-) create mode 100644 error_handler.py create mode 100644 exceptions.py create mode 100644 health_check.py create mode 100644 logging_config.py create mode 100644 performance_monitor.py create mode 100644 test_error_handling.py create mode 100644 test_logging_monitoring.py diff --git a/__init__.py b/__init__.py index 54a5f34..0eebc3c 100644 --- a/__init__.py +++ b/__init__.py @@ -123,8 +123,29 @@ def register_menu_extension(): def initialize_queue_manager(): """Initialize the queue manager system.""" try: + # Initialize logging and monitoring first + from logging_config import setup_logging + from performance_monitor import setup_performance_monitoring + from health_check import setup_health_checks + + # Set up logging + log_dir = current_dir / "logs" + logging_config = setup_logging( + log_dir=log_dir, + console_level=20, # INFO level + file_level=10, # DEBUG level + structured=True + ) + + logger = logging_config.get_logger("init") + logger.info(f"Queue Manager logging initialized (v{__version__})") + + # Set up performance monitoring + perf_monitor = setup_performance_monitoring(max_history_size=1000) + logger.info("Performance monitoring initialized") + # Import and initialize core services - from api_routes import setup_api_routes + from api_routes import QueueManagerAPI from database import SQLiteDatabase from queue_service import QueueService @@ -132,17 +153,34 @@ def initialize_queue_manager(): db_path = current_dir / "queue_manager.db" database = SQLiteDatabase(str(db_path)) database.initialize() + logger.info("Database initialized") # Initialize queue service queue_service = QueueService(database) + logger.info("Queue service initialized") + + # Set up health checks with actual services + health_manager = setup_health_checks(database=database, queue_service=queue_service) + health_manager.start_monitoring(interval=300.0) # 5 minutes + logger.info("Health monitoring started") # Set up API routes (will be used by the web interface) - setup_api_routes(queue_service) + api = QueueManagerAPI(queue_service) + logger.info("API routes initialized") + logger.info(f"Queue Manager initialized successfully (v{__version__})") print(f"[Queue Manager] Initialized successfully (v{__version__})") return True except Exception as e: + # Fall back to basic logging if initialization fails + import logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + ) + logger = logging.getLogger("queue_manager.init") + logger.error(f"Failed to initialize queue manager: {e}") print(f"[Queue Manager] Failed to initialize: {e}") return False diff --git a/api_routes.py b/api_routes.py index 0635c8e..47787ef 100644 --- a/api_routes.py +++ b/api_routes.py @@ -12,10 +12,14 @@ from typing import Any from flask import Flask, jsonify, request, send_from_directory from werkzeug.exceptions import BadRequest, NotFound +from error_handler import ErrorHandler, create_error_response, with_api_error_handling +from exceptions import APIError, NotFoundError, QueueManagerError, ValidationError +from logging_config import get_logger from models import QueueFilter, QueueStatus -from queue_service import QueueService, QueueServiceError +from performance_monitor import get_performance_monitor, time_function +from queue_service import QueueService -logger = logging.getLogger(__name__) +logger = get_logger("api_routes") class QueueManagerAPI: @@ -42,8 +46,10 @@ class QueueManagerAPI: def _register_routes(self) -> None: """Register all API routes.""" - # Health check + # Health check endpoints self.app.route('/health', methods=['GET'])(self.health_check) + self.app.route('/health/detailed', methods=['GET'])(self.detailed_health_check) + self.app.route('/health/metrics', methods=['GET'])(self.health_metrics) # Static files self.app.route('/', methods=['GET'])(self.serve_index) @@ -73,64 +79,115 @@ class QueueManagerAPI: def _register_error_handlers(self) -> None: """Register error handlers for the Flask app.""" - @self.app.errorhandler(QueueServiceError) - def handle_queue_service_error(error: QueueServiceError) -> tuple[dict[str, Any], int]: - """Handle queue service errors.""" - logger.error(f"Queue service error: {error}") - return { - 'error': 'Queue service error', - 'message': str(error), - 'timestamp': datetime.utcnow().isoformat() - }, 500 + @self.app.errorhandler(QueueManagerError) + def handle_queue_manager_error(error: QueueManagerError) -> tuple[dict[str, Any], int]: + """Handle queue manager errors.""" + ErrorHandler.log_error(error, context={"handler": "flask_error_handler"}) + + # Determine HTTP status code + if isinstance(error, APIError): + status_code = error.details.get('status_code', 500) + elif isinstance(error, ValidationError): + status_code = 400 + elif isinstance(error, NotFoundError): + status_code = 404 + else: + status_code = 500 + + return create_error_response(error), status_code @self.app.errorhandler(BadRequest) def handle_bad_request(error: BadRequest) -> tuple[dict[str, Any], int]: """Handle bad request errors.""" logger.warning(f"Bad request: {error}") - return { - 'error': 'Bad request', - 'message': str(error.description), - 'timestamp': datetime.utcnow().isoformat() - }, 400 + validation_error = ValidationError( + str(error.description) or "Bad request", + details={"werkzeug_error": str(error)} + ) + return create_error_response(validation_error), 400 @self.app.errorhandler(NotFound) def handle_not_found(error: NotFound) -> tuple[dict[str, Any], int]: """Handle not found errors.""" - return { - 'error': 'Not found', - 'message': 'The requested resource was not found', - 'timestamp': datetime.utcnow().isoformat() - }, 404 + not_found_error = NotFoundError( + "The requested resource was not found", + details={"werkzeug_error": str(error)} + ) + return create_error_response(not_found_error), 404 @self.app.errorhandler(Exception) def handle_generic_error(error: Exception) -> tuple[dict[str, Any], int]: """Handle generic errors.""" logger.error(f"Unexpected error: {error}", exc_info=True) - return { - 'error': 'Internal server error', - 'message': 'An unexpected error occurred', - 'timestamp': datetime.utcnow().isoformat() - }, 500 + + # Convert to structured error + api_error = APIError( + "An unexpected error occurred", + status_code=500, + details={"original_error": str(error)}, + cause=error + ) + + return create_error_response(api_error, include_traceback=False), 500 + @with_api_error_handling(endpoint="/health", method="GET") + @time_function("api.health_check") def health_check(self) -> dict[str, Any]: """Health check endpoint.""" + # Test queue service connectivity + queue_state = self.queue_service.get_queue_state() + + return { + 'status': 'healthy', + 'timestamp': datetime.utcnow().isoformat(), + 'queue_state': queue_state.value, + 'version': '1.0.0' + } + + @with_api_error_handling(endpoint="/health/detailed", method="GET") + @time_function("api.detailed_health_check") + def detailed_health_check(self) -> dict[str, Any]: + """Detailed health check endpoint with comprehensive diagnostics.""" + from health_check import get_health_check_manager + import asyncio + + # Get health check manager and run all checks + health_manager = get_health_check_manager() + + # Run health checks in a new event loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: - # Test queue service connectivity - queue_state = self.queue_service.get_queue_state() - - return { - 'status': 'healthy', - 'timestamp': datetime.utcnow().isoformat(), - 'queue_state': queue_state.value, - 'version': '1.0.0' - } - except Exception as e: - logger.error(f"Health check failed: {e}") - return { - 'status': 'unhealthy', - 'timestamp': datetime.utcnow().isoformat(), - 'error': str(e) - }, 503 + results = loop.run_until_complete(health_manager.run_all_checks()) + overall_health = health_manager.get_overall_health() + finally: + loop.close() + + return overall_health + + @with_api_error_handling(endpoint="/health/metrics", method="GET") + @time_function("api.health_metrics") + def health_metrics(self) -> dict[str, Any]: + """Health metrics endpoint with performance data.""" + perf_monitor = get_performance_monitor() + + # Get metrics summary for the last hour + since = datetime.utcnow() - timedelta(hours=1) + metrics_summary = perf_monitor.get_metrics_summary(since) + + # Get operation statistics + operation_stats = perf_monitor.get_operation_stats() + + # Get health status + health_status = perf_monitor.get_health_status() + + return { + 'timestamp': datetime.utcnow().isoformat(), + 'health_status': health_status, + 'metrics_summary': metrics_summary, + 'operation_stats': operation_stats + } def serve_index(self) -> str: """Serve the main index.html file.""" @@ -146,15 +203,15 @@ class QueueManagerAPI: def _validate_json_request(self) -> dict[str, Any]: """Validate and parse JSON request data.""" if not request.is_json: - raise BadRequest("Request must be JSON") + raise ValidationError("Request must be JSON", field="content_type") try: data = request.get_json() if data is None: - raise BadRequest("Invalid JSON data") + raise ValidationError("Invalid JSON data", field="request_body") return data except Exception as e: - raise BadRequest(f"Failed to parse JSON: {e}") + raise ValidationError(f"Failed to parse JSON: {e}", field="request_body", cause=e) def _parse_queue_filter(self, params: dict[str, Any]) -> QueueFilter: """Parse queue filter parameters from request.""" @@ -169,15 +226,21 @@ class QueueManagerAPI: try: filter_criteria.status = [QueueStatus(status) for status in status_values] except ValueError as e: - raise BadRequest(f"Invalid status value: {e}") + raise ValidationError(f"Invalid status value: {e}", field="status", value=status_values) # Parse workflow name filter if 'workflow_name' in params: - filter_criteria.workflow_name = params['workflow_name'] + workflow_name = params['workflow_name'] + if not isinstance(workflow_name, str): + raise ValidationError("Workflow name must be a string", field="workflow_name", value=workflow_name) + filter_criteria.workflow_name = workflow_name # Parse search term if 'search' in params: - filter_criteria.search_term = params['search'] + search_term = params['search'] + if not isinstance(search_term, str): + raise ValidationError("Search term must be a string", field="search", value=search_term) + filter_criteria.search_term = search_term # Parse date range if 'date_from' in params or 'date_to' in params: @@ -192,10 +255,12 @@ class QueueManagerAPI: date_to = datetime.fromisoformat(params['date_to']) if date_from and date_to: + if date_from > date_to: + raise ValidationError("date_from must be before date_to", field="date_range") filter_criteria.date_range = (date_from, date_to) except ValueError as e: - raise BadRequest(f"Invalid date format: {e}") + raise ValidationError(f"Invalid date format: {e}", field="date_range", cause=e) return filter_criteria diff --git a/database.py b/database.py index 2165e16..d3427f0 100644 --- a/database.py +++ b/database.py @@ -13,14 +13,14 @@ from datetime import datetime, timezone from pathlib import Path from typing import Any, Generator +from error_handler import ErrorHandler, with_database_error_handling +from exceptions import ConnectionError, DatabaseError, SchemaError, ValidationError from interfaces import DatabaseInterface +from logging_config import get_logger from models import QueueConfig, QueueFilter, QueueItem, QueueStatus +from performance_monitor import get_performance_monitor, time_function -logger = logging.getLogger(__name__) - - -class DatabaseError(Exception): - """Custom exception for database operations.""" +logger = get_logger("database") class SQLiteDatabase(DatabaseInterface): @@ -95,7 +95,7 @@ class SQLiteDatabase(DatabaseInterface): logger.debug(f"Created new database connection for thread {threading.current_thread().name}") except sqlite3.Error as e: logger.error(f"Failed to connect to database: {e}") - raise DatabaseError(f"Database connection failed: {e}") from e + raise ErrorHandler.handle_database_error(e, "connect", context={"db_path": str(self.db_path)}) return self._local.connection @@ -108,47 +108,53 @@ class SQLiteDatabase(DatabaseInterface): yield cursor except sqlite3.Error as e: logger.error(f"Database operation failed: {e}") - conn.rollback() - raise DatabaseError(f"Database operation failed: {e}") from e + try: + conn.rollback() + except sqlite3.Error as rollback_error: + logger.error(f"Failed to rollback transaction: {rollback_error}") + raise ErrorHandler.handle_database_error(e, "cursor_operation") finally: - cursor.close() + try: + cursor.close() + except sqlite3.Error as close_error: + logger.warning(f"Failed to close cursor: {close_error}") + @with_database_error_handling(operation="initialize") def initialize(self) -> bool: """Initialize the database and create tables if needed.""" if self._initialized: return True with self._connection_lock: - try: - with self._get_cursor() as cursor: - # Create all tables - cursor.executescript(self.CREATE_TABLES_SQL) - - # Check and update schema version - cursor.execute("SELECT version FROM schema_version LIMIT 1") - result = cursor.fetchone() - - if result is None: - # First time setup - cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) - logger.info(f"Database initialized with schema version {self.SCHEMA_VERSION}") - else: - current_version = result[0] - if current_version < self.SCHEMA_VERSION: - self._migrate_schema(cursor, current_version, self.SCHEMA_VERSION) - elif current_version > self.SCHEMA_VERSION: - logger.warning(f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}") - - # Set default configuration if not exists - self._initialize_default_config(cursor) - - self._initialized = True - logger.info("Database initialization completed successfully") - return True + with self._get_cursor() as cursor: + # Create all tables + cursor.executescript(self.CREATE_TABLES_SQL) - except Exception as e: - logger.error(f"Database initialization failed: {e}") - return False + # Check and update schema version + cursor.execute("SELECT version FROM schema_version LIMIT 1") + result = cursor.fetchone() + + if result is None: + # First time setup + cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,)) + logger.info(f"Database initialized with schema version {self.SCHEMA_VERSION}") + else: + current_version = result[0] + if current_version < self.SCHEMA_VERSION: + self._migrate_schema(cursor, current_version, self.SCHEMA_VERSION) + elif current_version > self.SCHEMA_VERSION: + logger.warning(f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}") + raise SchemaError( + f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}", + schema_version=current_version + ) + + # Set default configuration if not exists + self._initialize_default_config(cursor) + + self._initialized = True + logger.info("Database initialization completed successfully") + return True def _migrate_schema(self, cursor: sqlite3.Cursor, from_version: int, to_version: int) -> None: """Migrate database schema from one version to another.""" @@ -171,105 +177,108 @@ class SQLiteDatabase(DatabaseInterface): (key, json.dumps(value)) ) + @with_database_error_handling(operation="INSERT", table="queue_items") + @time_function("database.create_queue_item") def create_queue_item(self, item: QueueItem) -> bool: """Create a new queue item in the database.""" - try: - with self._get_cursor() as cursor: - cursor.execute(""" - INSERT INTO queue_items ( - id, workflow_name, workflow_data, status, created_at, updated_at, - started_at, completed_at, error_message, result_data - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, ( - item.id, - item.workflow_name, - json.dumps(item.workflow_data), - item.status.value, - item.created_at.isoformat(), - item.updated_at.isoformat(), - item.started_at.isoformat() if item.started_at else None, - item.completed_at.isoformat() if item.completed_at else None, - item.error_message, - json.dumps(item.result_data) if item.result_data else None - )) - logger.debug(f"Created queue item {item.id}") - return True - except Exception as e: - logger.error(f"Failed to create queue item {item.id}: {e}") - return False + # Validate item data + if not item.id: + raise ValidationError("Queue item ID cannot be empty", field="id") + if not item.workflow_name: + raise ValidationError("Workflow name cannot be empty", field="workflow_name") + + with self._get_cursor() as cursor: + cursor.execute(""" + INSERT INTO queue_items ( + id, workflow_name, workflow_data, status, created_at, updated_at, + started_at, completed_at, error_message, result_data + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, ( + item.id, + item.workflow_name, + json.dumps(item.workflow_data), + item.status.value, + item.created_at.isoformat(), + item.updated_at.isoformat(), + item.started_at.isoformat() if item.started_at else None, + item.completed_at.isoformat() if item.completed_at else None, + item.error_message, + json.dumps(item.result_data) if item.result_data else None + )) + logger.debug(f"Created queue item {item.id}") + return True + @with_database_error_handling(operation="SELECT", table="queue_items") + @time_function("database.get_queue_item") def get_queue_item(self, item_id: str) -> QueueItem | None: """Retrieve a queue item by ID.""" - try: - with self._get_cursor() as cursor: - cursor.execute("SELECT * FROM queue_items WHERE id = ?", (item_id,)) - row = cursor.fetchone() - if row: - return self._row_to_queue_item(row) - return None - except Exception as e: - logger.error(f"Failed to get queue item {item_id}: {e}") + if not item_id: + raise ValidationError("Item ID cannot be empty", field="item_id") + + with self._get_cursor() as cursor: + cursor.execute("SELECT * FROM queue_items WHERE id = ?", (item_id,)) + row = cursor.fetchone() + if row: + return self._row_to_queue_item(row) return None + @with_database_error_handling(operation="SELECT", table="queue_items") + @time_function("database.get_all_queue_items") def get_all_queue_items(self) -> list[QueueItem]: """Retrieve all queue items.""" - try: - with self._get_cursor() as cursor: - cursor.execute("SELECT * FROM queue_items ORDER BY created_at DESC") - rows = cursor.fetchall() - return [self._row_to_queue_item(row) for row in rows] - except Exception as e: - logger.error(f"Failed to get all queue items: {e}") - return [] + with self._get_cursor() as cursor: + cursor.execute("SELECT * FROM queue_items ORDER BY created_at DESC") + rows = cursor.fetchall() + return [self._row_to_queue_item(row) for row in rows] + @with_database_error_handling(operation="UPDATE", table="queue_items") def update_queue_item(self, item: QueueItem) -> bool: """Update an existing queue item.""" - try: - # Update the updated_at timestamp - item.updated_at = datetime.now(timezone.utc) + if not item.id: + raise ValidationError("Queue item ID cannot be empty", field="id") + + # Update the updated_at timestamp + item.updated_at = datetime.now(timezone.utc) + + with self._get_cursor() as cursor: + cursor.execute(""" + UPDATE queue_items SET + workflow_name = ?, workflow_data = ?, status = ?, updated_at = ?, + started_at = ?, completed_at = ?, error_message = ?, result_data = ? + WHERE id = ? + """, ( + item.workflow_name, + json.dumps(item.workflow_data), + item.status.value, + item.updated_at.isoformat(), + item.started_at.isoformat() if item.started_at else None, + item.completed_at.isoformat() if item.completed_at else None, + item.error_message, + json.dumps(item.result_data) if item.result_data else None, + item.id + )) - with self._get_cursor() as cursor: - cursor.execute(""" - UPDATE queue_items SET - workflow_name = ?, workflow_data = ?, status = ?, updated_at = ?, - started_at = ?, completed_at = ?, error_message = ?, result_data = ? - WHERE id = ? - """, ( - item.workflow_name, - json.dumps(item.workflow_data), - item.status.value, - item.updated_at.isoformat(), - item.started_at.isoformat() if item.started_at else None, - item.completed_at.isoformat() if item.completed_at else None, - item.error_message, - json.dumps(item.result_data) if item.result_data else None, - item.id - )) - - if cursor.rowcount > 0: - logger.debug(f"Updated queue item {item.id}") - return True - else: - logger.warning(f"Queue item {item.id} not found for update") - return False - except Exception as e: - logger.error(f"Failed to update queue item {item.id}: {e}") - return False + if cursor.rowcount > 0: + logger.debug(f"Updated queue item {item.id}") + return True + else: + logger.warning(f"Queue item {item.id} not found for update") + return False + @with_database_error_handling(operation="DELETE", table="queue_items") def delete_queue_item(self, item_id: str) -> bool: """Delete a queue item by ID.""" - try: - with self._get_cursor() as cursor: - cursor.execute("DELETE FROM queue_items WHERE id = ?", (item_id,)) - if cursor.rowcount > 0: - logger.debug(f"Deleted queue item {item_id}") - return True - else: - logger.warning(f"Queue item {item_id} not found for deletion") - return False - except Exception as e: - logger.error(f"Failed to delete queue item {item_id}: {e}") - return False + if not item_id: + raise ValidationError("Item ID cannot be empty", field="item_id") + + with self._get_cursor() as cursor: + cursor.execute("DELETE FROM queue_items WHERE id = ?", (item_id,)) + if cursor.rowcount > 0: + logger.debug(f"Deleted queue item {item_id}") + return True + else: + logger.warning(f"Queue item {item_id} not found for deletion") + return False def get_items_by_status(self, status: QueueStatus) -> list[QueueItem]: """Retrieve queue items by status.""" diff --git a/error_handler.py b/error_handler.py new file mode 100644 index 0000000..d2d8c16 --- /dev/null +++ b/error_handler.py @@ -0,0 +1,467 @@ +""" +Error handling utilities for the ComfyUI Queue Manager. +Provides centralized error handling, logging, and recovery mechanisms. +""" + +from __future__ import annotations + +import functools +import logging +import sqlite3 +import traceback +from datetime import datetime, timezone +from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union + +from exceptions import ( + APIError, + ConnectionError, + DatabaseError, + QueueManagerError, + QueueServiceError, + SchemaError, + TimeoutError, + ValidationError, + WorkflowExecutionError, +) + +logger = logging.getLogger(__name__) + +T = TypeVar('T') + + +class ErrorHandler: + """Centralized error handling and recovery utilities.""" + + @staticmethod + def handle_database_error( + error: Exception, + operation: str, + table: Optional[str] = None, + context: Optional[Dict[str, Any]] = None + ) -> DatabaseError: + """Handle and convert database errors to structured exceptions. + + Args: + error: Original exception + operation: Database operation that failed + table: Database table involved + context: Additional context information + + Returns: + Structured DatabaseError + """ + context = context or {} + + # Handle SQLite-specific errors + if isinstance(error, sqlite3.Error): + if isinstance(error, sqlite3.OperationalError): + error_msg = str(error).lower() + + if "database is locked" in error_msg: + return ConnectionError( + f"Database is locked during {operation}", + operation=operation, + table=table, + details=context, + cause=error + ) + elif "no such table" in error_msg: + return SchemaError( + f"Table does not exist during {operation}", + operation=operation, + table=table, + details=context, + cause=error + ) + elif "disk i/o error" in error_msg: + return DatabaseError( + f"Disk I/O error during {operation}", + error_code="DB_IO_ERROR", + operation=operation, + table=table, + details=context, + cause=error + ) + else: + return DatabaseError( + f"Database operational error during {operation}: {error}", + error_code="DB_OPERATIONAL_ERROR", + operation=operation, + table=table, + details=context, + cause=error + ) + + elif isinstance(error, sqlite3.IntegrityError): + return ValidationError( + f"Data integrity violation during {operation}: {error}", + details={**context, "operation": operation, "table": table}, + cause=error + ) + + elif isinstance(error, sqlite3.DatabaseError): + return DatabaseError( + f"Database error during {operation}: {error}", + error_code="DB_GENERAL_ERROR", + operation=operation, + table=table, + details=context, + cause=error + ) + + # Handle general exceptions + return DatabaseError( + f"Unexpected error during {operation}: {error}", + error_code="DB_UNEXPECTED_ERROR", + operation=operation, + table=table, + details=context, + cause=error + ) + + @staticmethod + def handle_api_error( + error: Exception, + endpoint: str, + method: str = "GET", + context: Optional[Dict[str, Any]] = None + ) -> APIError: + """Handle and convert API errors to structured exceptions. + + Args: + error: Original exception + endpoint: API endpoint that failed + method: HTTP method + context: Additional context information + + Returns: + Structured APIError + """ + context = context or {} + context.update({"endpoint": endpoint, "method": method}) + + if isinstance(error, QueueManagerError): + # Convert existing queue manager errors to API errors + return APIError( + error.message, + status_code=500, + endpoint=endpoint, + details=context, + cause=error + ) + + return APIError( + f"API error at {method} {endpoint}: {error}", + status_code=500, + endpoint=endpoint, + details=context, + cause=error + ) + + @staticmethod + def handle_workflow_error( + error: Exception, + workflow_id: Optional[str] = None, + stage: Optional[str] = None, + context: Optional[Dict[str, Any]] = None + ) -> WorkflowExecutionError: + """Handle and convert workflow execution errors. + + Args: + error: Original exception + workflow_id: ID of the workflow that failed + stage: Execution stage where error occurred + context: Additional context information + + Returns: + Structured WorkflowExecutionError + """ + context = context or {} + + return WorkflowExecutionError( + f"Workflow execution failed: {error}", + workflow_id=workflow_id, + execution_stage=stage, + details=context, + cause=error + ) + + @staticmethod + def log_error( + error: Exception, + context: Optional[Dict[str, Any]] = None, + level: int = logging.ERROR + ) -> None: + """Log an error with structured information. + + Args: + error: Exception to log + context: Additional context information + level: Logging level + """ + context = context or {} + + if isinstance(error, QueueManagerError): + # Log structured queue manager errors + log_data = { + "error_type": error.__class__.__name__, + "error_code": error.error_code, + "message": error.message, + "details": error.details, + **context + } + + logger.log(level, f"Queue Manager Error: {error.message}", extra=log_data) + + if error.cause: + logger.log(level, f"Caused by: {error.cause}", exc_info=error.cause) + else: + # Log general exceptions + logger.log(level, f"Unexpected error: {error}", extra=context, exc_info=error) + + +def with_error_handling( + error_type: Type[QueueManagerError] = QueueManagerError, + operation: Optional[str] = None, + log_errors: bool = True, + reraise: bool = True +): + """Decorator for automatic error handling and logging. + + Args: + error_type: Type of exception to convert errors to + operation: Name of the operation for context + log_errors: Whether to log errors automatically + reraise: Whether to reraise the converted exception + """ + def decorator(func: Callable[..., T]) -> Callable[..., T]: + @functools.wraps(func) + def wrapper(*args, **kwargs) -> T: + try: + return func(*args, **kwargs) + except QueueManagerError: + # Re-raise queue manager errors as-is + raise + except Exception as e: + # Convert other exceptions + operation_name = operation or func.__name__ + + if error_type == DatabaseError: + converted_error = ErrorHandler.handle_database_error( + e, operation_name, context={"function": func.__name__} + ) + elif error_type == APIError: + converted_error = ErrorHandler.handle_api_error( + e, operation_name, context={"function": func.__name__} + ) + elif error_type == WorkflowExecutionError: + converted_error = ErrorHandler.handle_workflow_error( + e, context={"function": func.__name__} + ) + else: + converted_error = error_type( + f"Error in {operation_name}: {e}", + details={"function": func.__name__}, + cause=e + ) + + if log_errors: + ErrorHandler.log_error(converted_error) + + if reraise: + raise converted_error + + return converted_error + + return wrapper + return decorator + + +def with_database_error_handling( + operation: Optional[str] = None, + table: Optional[str] = None, + log_errors: bool = True +): + """Decorator specifically for database operations. + + Args: + operation: Database operation name + table: Database table name + log_errors: Whether to log errors automatically + """ + def decorator(func: Callable[..., T]) -> Callable[..., T]: + @functools.wraps(func) + def wrapper(*args, **kwargs) -> T: + try: + return func(*args, **kwargs) + except QueueManagerError: + # Re-raise queue manager errors as-is + raise + except Exception as e: + operation_name = operation or func.__name__ + converted_error = ErrorHandler.handle_database_error( + e, + operation_name, + table=table, + context={"function": func.__name__} + ) + + if log_errors: + ErrorHandler.log_error(converted_error) + + raise converted_error + + return wrapper + return decorator + + +def with_api_error_handling( + endpoint: Optional[str] = None, + method: str = "GET", + log_errors: bool = True +): + """Decorator specifically for API operations. + + Args: + endpoint: API endpoint name + method: HTTP method + log_errors: Whether to log errors automatically + """ + def decorator(func: Callable[..., T]) -> Callable[..., T]: + @functools.wraps(func) + def wrapper(*args, **kwargs) -> T: + try: + return func(*args, **kwargs) + except QueueManagerError: + # Re-raise queue manager errors as-is + raise + except Exception as e: + endpoint_name = endpoint or func.__name__ + converted_error = ErrorHandler.handle_api_error( + e, + endpoint_name, + method=method, + context={"function": func.__name__} + ) + + if log_errors: + ErrorHandler.log_error(converted_error) + + raise converted_error + + return wrapper + return decorator + + +class ErrorRecovery: + """Utilities for error recovery and retry mechanisms.""" + + @staticmethod + def retry_with_backoff( + func: Callable[..., T], + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 60.0, + backoff_factor: float = 2.0, + exceptions: tuple = (Exception,) + ) -> T: + """Retry a function with exponential backoff. + + Args: + func: Function to retry + max_retries: Maximum number of retry attempts + base_delay: Initial delay between retries + max_delay: Maximum delay between retries + backoff_factor: Factor to multiply delay by each retry + exceptions: Tuple of exceptions to catch and retry on + + Returns: + Result of the function call + + Raises: + The last exception if all retries fail + """ + import time + + last_exception = None + delay = base_delay + + for attempt in range(max_retries + 1): + try: + return func() + except exceptions as e: + last_exception = e + + if attempt == max_retries: + # Last attempt failed, re-raise + break + + logger.warning( + f"Attempt {attempt + 1} failed, retrying in {delay:.2f}s: {e}" + ) + + time.sleep(delay) + delay = min(delay * backoff_factor, max_delay) + + # All retries failed + if last_exception: + raise last_exception + else: + raise RuntimeError("All retry attempts failed") + + @staticmethod + def safe_execute( + func: Callable[..., T], + default_value: Optional[T] = None, + log_errors: bool = True, + exceptions: tuple = (Exception,) + ) -> Optional[T]: + """Execute a function safely, returning a default value on error. + + Args: + func: Function to execute + default_value: Value to return on error + log_errors: Whether to log errors + exceptions: Tuple of exceptions to catch + + Returns: + Function result or default value + """ + try: + return func() + except exceptions as e: + if log_errors: + logger.error(f"Safe execution failed: {e}", exc_info=True) + return default_value + + +def create_error_response( + error: Exception, + include_traceback: bool = False +) -> Dict[str, Any]: + """Create a standardized error response dictionary. + + Args: + error: Exception to convert to response + include_traceback: Whether to include traceback in response + + Returns: + Dictionary containing error information + """ + if isinstance(error, QueueManagerError): + response = error.to_dict() + else: + response = { + "error_type": error.__class__.__name__, + "error_code": "UNEXPECTED_ERROR", + "message": str(error), + "details": {}, + "cause": None + } + + response["timestamp"] = datetime.now(timezone.utc).isoformat() + + if include_traceback: + response["traceback"] = traceback.format_exc() + + return response \ No newline at end of file diff --git a/exceptions.py b/exceptions.py new file mode 100644 index 0000000..eb36dd4 --- /dev/null +++ b/exceptions.py @@ -0,0 +1,450 @@ +""" +Custom exception classes for the ComfyUI Queue Manager. +Provides structured error handling across all components. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + + +class QueueManagerError(Exception): + """Base exception class for all queue manager errors.""" + + def __init__( + self, + message: str, + error_code: Optional[str] = None, + details: Optional[Dict[str, Any]] = None, + cause: Optional[Exception] = None + ): + """Initialize the exception. + + Args: + message: Human-readable error message + error_code: Machine-readable error code + details: Additional error details + cause: Original exception that caused this error + """ + super().__init__(message) + self.message = message + self.error_code = error_code or self.__class__.__name__ + self.details = details or {} + self.cause = cause + + def to_dict(self) -> Dict[str, Any]: + """Convert the exception to a dictionary for serialization.""" + return { + "error_type": self.__class__.__name__, + "error_code": self.error_code, + "message": self.message, + "details": self.details, + "cause": str(self.cause) if self.cause else None + } + + +class DatabaseError(QueueManagerError): + """Exception raised for database-related errors.""" + + def __init__( + self, + message: str, + operation: Optional[str] = None, + table: Optional[str] = None, + **kwargs + ): + """Initialize the database error. + + Args: + message: Error message + operation: Database operation that failed (e.g., 'INSERT', 'UPDATE') + table: Database table involved in the operation + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if operation: + details['operation'] = operation + if table: + details['table'] = table + + super().__init__(message, details=details, **kwargs) + + +class ConnectionError(DatabaseError): + """Exception raised for database connection errors.""" + + def __init__(self, message: str, **kwargs): + super().__init__( + message, + error_code="DB_CONNECTION_ERROR", + **kwargs + ) + + +class SchemaError(DatabaseError): + """Exception raised for database schema-related errors.""" + + def __init__(self, message: str, schema_version: Optional[int] = None, **kwargs): + details = kwargs.get('details', {}) + if schema_version is not None: + details['schema_version'] = schema_version + + super().__init__( + message, + error_code="DB_SCHEMA_ERROR", + details=details, + **kwargs + ) + + +class ValidationError(QueueManagerError): + """Exception raised for data validation errors.""" + + def __init__( + self, + message: str, + field: Optional[str] = None, + value: Optional[Any] = None, + **kwargs + ): + """Initialize the validation error. + + Args: + message: Error message + field: Field name that failed validation + value: Value that failed validation + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if field: + details['field'] = field + if value is not None: + details['value'] = str(value) # Convert to string for serialization + + super().__init__( + message, + error_code="VALIDATION_ERROR", + details=details, + **kwargs + ) + + +class QueueServiceError(QueueManagerError): + """Exception raised for queue service operations.""" + + def __init__( + self, + message: str, + item_id: Optional[str] = None, + operation: Optional[str] = None, + **kwargs + ): + """Initialize the queue service error. + + Args: + message: Error message + item_id: Queue item ID involved in the operation + operation: Queue operation that failed + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if item_id: + details['item_id'] = item_id + if operation: + details['operation'] = operation + + super().__init__( + message, + error_code="QUEUE_SERVICE_ERROR", + details=details, + **kwargs + ) + + +class WorkflowExecutionError(QueueManagerError): + """Exception raised for workflow execution errors.""" + + def __init__( + self, + message: str, + workflow_id: Optional[str] = None, + execution_stage: Optional[str] = None, + **kwargs + ): + """Initialize the workflow execution error. + + Args: + message: Error message + workflow_id: Workflow ID that failed + execution_stage: Stage of execution where error occurred + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if workflow_id: + details['workflow_id'] = workflow_id + if execution_stage: + details['execution_stage'] = execution_stage + + super().__init__( + message, + error_code="WORKFLOW_EXECUTION_ERROR", + details=details, + **kwargs + ) + + +class TimeoutError(QueueManagerError): + """Exception raised for timeout-related errors.""" + + def __init__( + self, + message: str, + timeout_seconds: Optional[float] = None, + operation: Optional[str] = None, + **kwargs + ): + """Initialize the timeout error. + + Args: + message: Error message + timeout_seconds: Timeout value that was exceeded + operation: Operation that timed out + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if timeout_seconds is not None: + details['timeout_seconds'] = timeout_seconds + if operation: + details['operation'] = operation + + super().__init__( + message, + error_code="TIMEOUT_ERROR", + details=details, + **kwargs + ) + + +class ConfigurationError(QueueManagerError): + """Exception raised for configuration-related errors.""" + + def __init__( + self, + message: str, + config_key: Optional[str] = None, + config_value: Optional[Any] = None, + **kwargs + ): + """Initialize the configuration error. + + Args: + message: Error message + config_key: Configuration key that caused the error + config_value: Configuration value that caused the error + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if config_key: + details['config_key'] = config_key + if config_value is not None: + details['config_value'] = str(config_value) + + super().__init__( + message, + error_code="CONFIGURATION_ERROR", + details=details, + **kwargs + ) + + +class APIError(QueueManagerError): + """Exception raised for API-related errors.""" + + def __init__( + self, + message: str, + status_code: int = 500, + endpoint: Optional[str] = None, + **kwargs + ): + """Initialize the API error. + + Args: + message: Error message + status_code: HTTP status code + endpoint: API endpoint that caused the error + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + details['status_code'] = status_code + if endpoint: + details['endpoint'] = endpoint + + super().__init__( + message, + error_code="API_ERROR", + details=details, + **kwargs + ) + + +class AuthenticationError(APIError): + """Exception raised for authentication errors.""" + + def __init__(self, message: str = "Authentication required", **kwargs): + super().__init__( + message, + status_code=401, + error_code="AUTHENTICATION_ERROR", + **kwargs + ) + + +class AuthorizationError(APIError): + """Exception raised for authorization errors.""" + + def __init__(self, message: str = "Access denied", **kwargs): + super().__init__( + message, + status_code=403, + error_code="AUTHORIZATION_ERROR", + **kwargs + ) + + +class NotFoundError(APIError): + """Exception raised when a resource is not found.""" + + def __init__( + self, + message: str, + resource_type: Optional[str] = None, + resource_id: Optional[str] = None, + **kwargs + ): + """Initialize the not found error. + + Args: + message: Error message + resource_type: Type of resource that was not found + resource_id: ID of resource that was not found + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if resource_type: + details['resource_type'] = resource_type + if resource_id: + details['resource_id'] = resource_id + + super().__init__( + message, + status_code=404, + error_code="NOT_FOUND_ERROR", + details=details, + **kwargs + ) + + +class ConflictError(APIError): + """Exception raised for resource conflicts.""" + + def __init__(self, message: str, **kwargs): + super().__init__( + message, + status_code=409, + error_code="CONFLICT_ERROR", + **kwargs + ) + + +class RateLimitError(APIError): + """Exception raised when rate limits are exceeded.""" + + def __init__( + self, + message: str = "Rate limit exceeded", + retry_after: Optional[int] = None, + **kwargs + ): + """Initialize the rate limit error. + + Args: + message: Error message + retry_after: Seconds to wait before retrying + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if retry_after is not None: + details['retry_after'] = retry_after + + super().__init__( + message, + status_code=429, + error_code="RATE_LIMIT_ERROR", + details=details, + **kwargs + ) + + +class ImportExportError(QueueManagerError): + """Exception raised for import/export operations.""" + + def __init__( + self, + message: str, + operation: Optional[str] = None, + file_path: Optional[str] = None, + **kwargs + ): + """Initialize the import/export error. + + Args: + message: Error message + operation: Operation that failed ('import' or 'export') + file_path: File path involved in the operation + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if operation: + details['operation'] = operation + if file_path: + details['file_path'] = file_path + + super().__init__( + message, + error_code="IMPORT_EXPORT_ERROR", + details=details, + **kwargs + ) + + +class ArchiveError(QueueManagerError): + """Exception raised for archive operations.""" + + def __init__( + self, + message: str, + operation: Optional[str] = None, + item_count: Optional[int] = None, + **kwargs + ): + """Initialize the archive error. + + Args: + message: Error message + operation: Archive operation that failed + item_count: Number of items involved in the operation + **kwargs: Additional arguments passed to parent + """ + details = kwargs.get('details', {}) + if operation: + details['operation'] = operation + if item_count is not None: + details['item_count'] = item_count + + super().__init__( + message, + error_code="ARCHIVE_ERROR", + details=details, + **kwargs + ) \ No newline at end of file diff --git a/health_check.py b/health_check.py new file mode 100644 index 0000000..da732ea --- /dev/null +++ b/health_check.py @@ -0,0 +1,635 @@ +""" +Health check system for the ComfyUI Queue Manager. +Provides comprehensive health monitoring and diagnostics. +""" + +from __future__ import annotations + +import asyncio +import logging +import threading +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone, timedelta +from enum import Enum +from typing import Any, Callable, Dict, List, Optional + +from logging_config import get_logger +from performance_monitor import get_performance_monitor + +logger = get_logger("health_check") + + +class HealthStatus(Enum): + """Health status enumeration.""" + + HEALTHY = "healthy" + DEGRADED = "degraded" + UNHEALTHY = "unhealthy" + UNKNOWN = "unknown" + + +@dataclass +class HealthCheckResult: + """Result of a health check.""" + + name: str + status: HealthStatus + message: str + details: Dict[str, Any] = field(default_factory=dict) + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + duration: float = 0.0 # seconds + + def to_dict(self) -> Dict[str, Any]: + """Convert result to dictionary.""" + return { + "name": self.name, + "status": self.status.value, + "message": self.message, + "details": self.details, + "timestamp": self.timestamp.isoformat(), + "duration": self.duration + } + + +class HealthCheck: + """Base class for health checks.""" + + def __init__(self, name: str, timeout: float = 30.0): + """Initialize health check. + + Args: + name: Name of the health check + timeout: Timeout in seconds + """ + self.name = name + self.timeout = timeout + + async def check(self) -> HealthCheckResult: + """Perform the health check. + + Returns: + HealthCheckResult + """ + start_time = time.time() + + try: + # Run the actual check with timeout + result = await asyncio.wait_for(self._perform_check(), timeout=self.timeout) + duration = time.time() - start_time + result.duration = duration + return result + + except asyncio.TimeoutError: + duration = time.time() - start_time + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message=f"Health check timed out after {self.timeout} seconds", + duration=duration, + details={"timeout": self.timeout} + ) + except Exception as e: + duration = time.time() - start_time + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message=f"Health check failed: {e}", + duration=duration, + details={"error": str(e), "error_type": type(e).__name__} + ) + + async def _perform_check(self) -> HealthCheckResult: + """Perform the actual health check logic. + + This method should be overridden by subclasses. + + Returns: + HealthCheckResult + """ + raise NotImplementedError("Subclasses must implement _perform_check") + + +class DatabaseHealthCheck(HealthCheck): + """Health check for database connectivity and performance.""" + + def __init__(self, database, timeout: float = 10.0): + """Initialize database health check. + + Args: + database: Database instance to check + timeout: Timeout in seconds + """ + super().__init__("database", timeout) + self.database = database + + async def _perform_check(self) -> HealthCheckResult: + """Check database health.""" + details = {} + + try: + # Check if database is initialized + if not hasattr(self.database, '_initialized') or not self.database._initialized: + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message="Database not initialized", + details={"initialized": False} + ) + + # Test basic connectivity + start_time = time.time() + test_result = await asyncio.get_event_loop().run_in_executor( + None, self._test_database_connection + ) + connection_time = time.time() - start_time + + details.update(test_result) + details["connection_time"] = connection_time + + # Determine status based on results + if not test_result.get("connected", False): + status = HealthStatus.UNHEALTHY + message = "Database connection failed" + elif connection_time > 5.0: + status = HealthStatus.DEGRADED + message = f"Database connection slow ({connection_time:.2f}s)" + else: + status = HealthStatus.HEALTHY + message = "Database is healthy" + + return HealthCheckResult( + name=self.name, + status=status, + message=message, + details=details + ) + + except Exception as e: + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message=f"Database health check failed: {e}", + details={"error": str(e)} + ) + + def _test_database_connection(self) -> Dict[str, Any]: + """Test database connection and basic operations.""" + result = { + "connected": False, + "schema_version": None, + "table_count": 0, + "last_error": None + } + + try: + # Test connection + with self.database._get_cursor() as cursor: + # Check schema version + cursor.execute("SELECT version FROM schema_version LIMIT 1") + version_row = cursor.fetchone() + if version_row: + result["schema_version"] = version_row[0] + + # Count tables + cursor.execute("SELECT COUNT(*) FROM sqlite_master WHERE type='table'") + table_count_row = cursor.fetchone() + if table_count_row: + result["table_count"] = table_count_row[0] + + # Test a simple query + cursor.execute("SELECT COUNT(*) FROM queue_items") + item_count_row = cursor.fetchone() + if item_count_row: + result["queue_item_count"] = item_count_row[0] + + result["connected"] = True + + except Exception as e: + result["last_error"] = str(e) + + return result + + +class QueueServiceHealthCheck(HealthCheck): + """Health check for queue service functionality.""" + + def __init__(self, queue_service, timeout: float = 15.0): + """Initialize queue service health check. + + Args: + queue_service: Queue service instance to check + timeout: Timeout in seconds + """ + super().__init__("queue_service", timeout) + self.queue_service = queue_service + + async def _perform_check(self) -> HealthCheckResult: + """Check queue service health.""" + details = {} + + try: + # Check if service is initialized + if not hasattr(self.queue_service, '_initialized') or not self.queue_service._initialized: + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message="Queue service not initialized", + details={"initialized": False} + ) + + # Test basic operations + start_time = time.time() + test_result = await asyncio.get_event_loop().run_in_executor( + None, self._test_queue_service_operations + ) + operation_time = time.time() - start_time + + details.update(test_result) + details["operation_time"] = operation_time + + # Determine status + if test_result.get("errors", 0) > 0: + status = HealthStatus.UNHEALTHY + message = f"Queue service has {test_result['errors']} errors" + elif operation_time > 10.0: + status = HealthStatus.DEGRADED + message = f"Queue service operations slow ({operation_time:.2f}s)" + else: + status = HealthStatus.HEALTHY + message = "Queue service is healthy" + + return HealthCheckResult( + name=self.name, + status=status, + message=message, + details=details + ) + + except Exception as e: + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message=f"Queue service health check failed: {e}", + details={"error": str(e)} + ) + + def _test_queue_service_operations(self) -> Dict[str, Any]: + """Test queue service operations.""" + result = { + "queue_state": None, + "item_count": 0, + "config_accessible": False, + "errors": 0, + "last_error": None + } + + try: + # Test getting queue state + queue_state = self.queue_service.get_queue_state() + result["queue_state"] = queue_state.value + except Exception as e: + result["errors"] += 1 + result["last_error"] = f"get_queue_state: {e}" + + try: + # Test getting queue items + items = self.queue_service.get_queue_items() + result["item_count"] = len(items) + except Exception as e: + result["errors"] += 1 + result["last_error"] = f"get_queue_items: {e}" + + try: + # Test getting configuration + config = self.queue_service.get_config() + result["config_accessible"] = config is not None + except Exception as e: + result["errors"] += 1 + result["last_error"] = f"get_config: {e}" + + return result + + +class SystemResourceHealthCheck(HealthCheck): + """Health check for system resources.""" + + def __init__(self, timeout: float = 5.0): + """Initialize system resource health check. + + Args: + timeout: Timeout in seconds + """ + super().__init__("system_resources", timeout) + + async def _perform_check(self) -> HealthCheckResult: + """Check system resource health.""" + try: + perf_monitor = get_performance_monitor() + health_status = perf_monitor.get_health_status() + + # Map performance monitor status to health status + status_mapping = { + "healthy": HealthStatus.HEALTHY, + "degraded": HealthStatus.DEGRADED, + "unhealthy": HealthStatus.UNHEALTHY + } + + status = status_mapping.get(health_status["status"], HealthStatus.UNKNOWN) + + return HealthCheckResult( + name=self.name, + status=status, + message=f"System resources are {health_status['status']}", + details=health_status + ) + + except Exception as e: + return HealthCheckResult( + name=self.name, + status=HealthStatus.UNHEALTHY, + message=f"System resource check failed: {e}", + details={"error": str(e)} + ) + + +class HealthCheckManager: + """Manager for running and coordinating health checks.""" + + def __init__(self): + """Initialize health check manager.""" + self._checks: Dict[str, HealthCheck] = {} + self._last_results: Dict[str, HealthCheckResult] = {} + self._lock = threading.RLock() + + # Monitoring settings + self._monitoring_enabled = False + self._monitoring_interval = 60.0 # seconds + self._monitoring_thread: Optional[threading.Thread] = None + self._stop_monitoring = threading.Event() + + logger.info("Health check manager initialized") + + def register_check(self, health_check: HealthCheck) -> None: + """Register a health check. + + Args: + health_check: Health check to register + """ + with self._lock: + self._checks[health_check.name] = health_check + + logger.info(f"Registered health check: {health_check.name}") + + def unregister_check(self, name: str) -> None: + """Unregister a health check. + + Args: + name: Name of the health check to remove + """ + with self._lock: + self._checks.pop(name, None) + self._last_results.pop(name, None) + + logger.info(f"Unregistered health check: {name}") + + async def run_check(self, name: str) -> HealthCheckResult: + """Run a specific health check. + + Args: + name: Name of the health check to run + + Returns: + HealthCheckResult + """ + with self._lock: + health_check = self._checks.get(name) + + if not health_check: + return HealthCheckResult( + name=name, + status=HealthStatus.UNKNOWN, + message=f"Health check '{name}' not found" + ) + + result = await health_check.check() + + with self._lock: + self._last_results[name] = result + + # Log significant health changes + if result.status != HealthStatus.HEALTHY: + logger.warning( + f"Health check failed: {name}", + extra={ + "check_name": name, + "status": result.status.value, + "message": result.message, + "duration": result.duration + } + ) + + return result + + async def run_all_checks(self) -> Dict[str, HealthCheckResult]: + """Run all registered health checks. + + Returns: + Dictionary mapping check names to results + """ + with self._lock: + check_names = list(self._checks.keys()) + + results = {} + + # Run checks concurrently + tasks = [self.run_check(name) for name in check_names] + check_results = await asyncio.gather(*tasks, return_exceptions=True) + + for name, result in zip(check_names, check_results): + if isinstance(result, Exception): + results[name] = HealthCheckResult( + name=name, + status=HealthStatus.UNHEALTHY, + message=f"Health check failed with exception: {result}", + details={"exception": str(result)} + ) + else: + results[name] = result + + return results + + def get_last_results(self) -> Dict[str, HealthCheckResult]: + """Get the last results from all health checks. + + Returns: + Dictionary mapping check names to last results + """ + with self._lock: + return self._last_results.copy() + + def get_overall_health(self) -> Dict[str, Any]: + """Get overall health status. + + Returns: + Dictionary containing overall health information + """ + with self._lock: + results = self._last_results.copy() + + if not results: + return { + "status": HealthStatus.UNKNOWN.value, + "message": "No health checks have been run", + "timestamp": datetime.now(timezone.utc).isoformat(), + "checks": {} + } + + # Determine overall status + statuses = [result.status for result in results.values()] + + if any(status == HealthStatus.UNHEALTHY for status in statuses): + overall_status = HealthStatus.UNHEALTHY + message = "One or more health checks are unhealthy" + elif any(status == HealthStatus.DEGRADED for status in statuses): + overall_status = HealthStatus.DEGRADED + message = "One or more health checks are degraded" + elif any(status == HealthStatus.UNKNOWN for status in statuses): + overall_status = HealthStatus.UNKNOWN + message = "One or more health checks have unknown status" + else: + overall_status = HealthStatus.HEALTHY + message = "All health checks are healthy" + + return { + "status": overall_status.value, + "message": message, + "timestamp": datetime.now(timezone.utc).isoformat(), + "checks": {name: result.to_dict() for name, result in results.items()} + } + + def start_monitoring(self, interval: float = 60.0) -> None: + """Start continuous health monitoring. + + Args: + interval: Monitoring interval in seconds + """ + if self._monitoring_thread and self._monitoring_thread.is_alive(): + logger.warning("Health monitoring is already running") + return + + self._monitoring_interval = interval + self._monitoring_enabled = True + self._stop_monitoring.clear() + + self._monitoring_thread = threading.Thread( + target=self._monitoring_loop, + name="HealthMonitor", + daemon=True + ) + self._monitoring_thread.start() + + logger.info(f"Health monitoring started with {interval}s interval") + + def stop_monitoring(self) -> None: + """Stop continuous health monitoring.""" + if not self._monitoring_thread or not self._monitoring_thread.is_alive(): + return + + self._monitoring_enabled = False + self._stop_monitoring.set() + self._monitoring_thread.join(timeout=5.0) + + if self._monitoring_thread.is_alive(): + logger.warning("Health monitoring thread did not stop gracefully") + else: + logger.info("Health monitoring stopped") + + def _monitoring_loop(self) -> None: + """Health monitoring loop that runs in a separate thread.""" + logger.info("Health monitoring loop started") + + while self._monitoring_enabled and not self._stop_monitoring.is_set(): + try: + # Run health checks + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + results = loop.run_until_complete(self.run_all_checks()) + + # Log overall health status + overall_health = self.get_overall_health() + logger.info( + f"Health monitoring cycle completed: {overall_health['status']}", + extra={ + "overall_status": overall_health["status"], + "check_count": len(results), + "healthy_checks": sum(1 for r in results.values() if r.status == HealthStatus.HEALTHY), + "degraded_checks": sum(1 for r in results.values() if r.status == HealthStatus.DEGRADED), + "unhealthy_checks": sum(1 for r in results.values() if r.status == HealthStatus.UNHEALTHY) + } + ) + + loop.close() + + # Wait for next cycle + if self._stop_monitoring.wait(timeout=self._monitoring_interval): + break + + except Exception as e: + logger.error(f"Error in health monitoring loop: {e}", exc_info=True) + time.sleep(5.0) # Brief pause before retrying + + logger.info("Health monitoring loop stopped") + + def close(self) -> None: + """Close the health check manager and cleanup resources.""" + self.stop_monitoring() + + with self._lock: + self._checks.clear() + self._last_results.clear() + + logger.info("Health check manager closed") + + +# Global health check manager instance +_health_check_manager: Optional[HealthCheckManager] = None + + +def get_health_check_manager() -> HealthCheckManager: + """Get the global health check manager instance. + + Returns: + HealthCheckManager instance + """ + global _health_check_manager + if _health_check_manager is None: + _health_check_manager = HealthCheckManager() + return _health_check_manager + + +def setup_health_checks(database=None, queue_service=None) -> HealthCheckManager: + """Set up health checks for the queue manager. + + Args: + database: Database instance for health checks + queue_service: Queue service instance for health checks + + Returns: + HealthCheckManager instance + """ + manager = get_health_check_manager() + + # Register system resource check + manager.register_check(SystemResourceHealthCheck()) + + # Register database check if available + if database: + manager.register_check(DatabaseHealthCheck(database)) + + # Register queue service check if available + if queue_service: + manager.register_check(QueueServiceHealthCheck(queue_service)) + + return manager \ No newline at end of file diff --git a/logging_config.py b/logging_config.py new file mode 100644 index 0000000..3f242ea --- /dev/null +++ b/logging_config.py @@ -0,0 +1,450 @@ +""" +Logging configuration for the ComfyUI Queue Manager. +Provides structured logging with rotation, filtering, and monitoring capabilities. +""" + +from __future__ import annotations + +import json +import logging +import logging.handlers +import os +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Optional + + +class StructuredFormatter(logging.Formatter): + """Custom formatter that outputs structured JSON logs.""" + + def __init__(self, include_extra: bool = True): + """Initialize the structured formatter. + + Args: + include_extra: Whether to include extra fields from log records + """ + super().__init__() + self.include_extra = include_extra + + def format(self, record: logging.LogRecord) -> str: + """Format a log record as structured JSON. + + Args: + record: Log record to format + + Returns: + JSON-formatted log string + """ + # Base log data + log_data = { + "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(), + "level": record.levelname, + "logger": record.name, + "message": record.getMessage(), + "module": record.module, + "function": record.funcName, + "line": record.lineno, + } + + # Add thread information + if record.thread: + log_data["thread_id"] = record.thread + log_data["thread_name"] = record.threadName + + # Add process information + if record.process: + log_data["process_id"] = record.process + log_data["process_name"] = record.processName + + # Add exception information + if record.exc_info: + log_data["exception"] = { + "type": record.exc_info[0].__name__ if record.exc_info[0] else None, + "message": str(record.exc_info[1]) if record.exc_info[1] else None, + "traceback": self.formatException(record.exc_info) if record.exc_info else None + } + + # Add extra fields if enabled + if self.include_extra: + # Get all extra attributes (those not in the standard LogRecord) + standard_attrs = { + 'name', 'msg', 'args', 'levelname', 'levelno', 'pathname', 'filename', + 'module', 'lineno', 'funcName', 'created', 'msecs', 'relativeCreated', + 'thread', 'threadName', 'processName', 'process', 'getMessage', + 'exc_info', 'exc_text', 'stack_info', 'message' + } + + extra_data = {} + for key, value in record.__dict__.items(): + if key not in standard_attrs and not key.startswith('_'): + # Ensure value is JSON serializable + try: + json.dumps(value) + extra_data[key] = value + except (TypeError, ValueError): + extra_data[key] = str(value) + + if extra_data: + log_data["extra"] = extra_data + + return json.dumps(log_data, ensure_ascii=False) + + +class QueueManagerFilter(logging.Filter): + """Custom filter for queue manager logs.""" + + def __init__(self, component: Optional[str] = None, min_level: int = logging.INFO): + """Initialize the filter. + + Args: + component: Specific component to filter for (e.g., 'database', 'api') + min_level: Minimum log level to allow + """ + super().__init__() + self.component = component + self.min_level = min_level + + def filter(self, record: logging.LogRecord) -> bool: + """Filter log records based on criteria. + + Args: + record: Log record to filter + + Returns: + True if record should be logged + """ + # Check minimum level + if record.levelno < self.min_level: + return False + + # Check component filter + if self.component: + # Check if the logger name contains the component + if self.component not in record.name: + return False + + return True + + +class LoggingConfig: + """Configuration manager for logging setup.""" + + DEFAULT_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S" + + def __init__(self, log_dir: Optional[Path] = None): + """Initialize logging configuration. + + Args: + log_dir: Directory for log files. If None, uses default location. + """ + self.log_dir = log_dir or Path(__file__).parent / "logs" + self.log_dir.mkdir(parents=True, exist_ok=True) + + # Configuration settings + self.console_level = logging.INFO + self.file_level = logging.DEBUG + self.structured_logging = True + self.max_file_size = 10 * 1024 * 1024 # 10MB + self.backup_count = 5 + + # Logger instances + self._loggers: Dict[str, logging.Logger] = {} + self._configured = False + + def configure_logging( + self, + console_level: int = logging.INFO, + file_level: int = logging.DEBUG, + structured: bool = True, + max_file_size: int = 10 * 1024 * 1024, + backup_count: int = 5 + ) -> None: + """Configure logging for the queue manager. + + Args: + console_level: Log level for console output + file_level: Log level for file output + structured: Whether to use structured JSON logging + max_file_size: Maximum size of log files before rotation + backup_count: Number of backup files to keep + """ + self.console_level = console_level + self.file_level = file_level + self.structured_logging = structured + self.max_file_size = max_file_size + self.backup_count = backup_count + + # Configure root logger + root_logger = logging.getLogger() + root_logger.setLevel(logging.DEBUG) + + # Remove existing handlers to avoid duplicates + for handler in root_logger.handlers[:]: + root_logger.removeHandler(handler) + + # Configure console handler + self._configure_console_handler(root_logger) + + # Configure file handlers + self._configure_file_handlers(root_logger) + + # Configure component-specific loggers + self._configure_component_loggers() + + self._configured = True + + # Log configuration completion + logger = self.get_logger("logging_config") + logger.info("Logging configuration completed", extra={ + "console_level": logging.getLevelName(console_level), + "file_level": logging.getLevelName(file_level), + "structured": structured, + "log_dir": str(self.log_dir), + "max_file_size": max_file_size, + "backup_count": backup_count + }) + + def _configure_console_handler(self, root_logger: logging.Logger) -> None: + """Configure console logging handler.""" + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(self.console_level) + + if self.structured_logging: + console_formatter = StructuredFormatter(include_extra=False) + else: + console_formatter = logging.Formatter( + self.DEFAULT_FORMAT, + datefmt=self.DEFAULT_DATE_FORMAT + ) + + console_handler.setFormatter(console_formatter) + console_handler.addFilter(QueueManagerFilter(min_level=self.console_level)) + + root_logger.addHandler(console_handler) + + def _configure_file_handlers(self, root_logger: logging.Logger) -> None: + """Configure file logging handlers.""" + # Main application log + main_log_file = self.log_dir / "queue_manager.log" + main_handler = logging.handlers.RotatingFileHandler( + main_log_file, + maxBytes=self.max_file_size, + backupCount=self.backup_count, + encoding='utf-8' + ) + main_handler.setLevel(self.file_level) + + if self.structured_logging: + main_formatter = StructuredFormatter(include_extra=True) + else: + main_formatter = logging.Formatter( + self.DEFAULT_FORMAT, + datefmt=self.DEFAULT_DATE_FORMAT + ) + + main_handler.setFormatter(main_formatter) + root_logger.addHandler(main_handler) + + # Error-only log + error_log_file = self.log_dir / "errors.log" + error_handler = logging.handlers.RotatingFileHandler( + error_log_file, + maxBytes=self.max_file_size, + backupCount=self.backup_count, + encoding='utf-8' + ) + error_handler.setLevel(logging.ERROR) + error_handler.setFormatter(main_formatter) + root_logger.addHandler(error_handler) + + # Performance log + perf_log_file = self.log_dir / "performance.log" + perf_handler = logging.handlers.RotatingFileHandler( + perf_log_file, + maxBytes=self.max_file_size, + backupCount=self.backup_count, + encoding='utf-8' + ) + perf_handler.setLevel(logging.INFO) + perf_handler.setFormatter(main_formatter) + perf_handler.addFilter(QueueManagerFilter(component="performance")) + root_logger.addHandler(perf_handler) + + def _configure_component_loggers(self) -> None: + """Configure component-specific loggers.""" + components = [ + "database", + "queue_service", + "api_routes", + "workflow_executor", + "execution_monitor", + "archive_service", + "filter_service", + "import_export_service" + ] + + for component in components: + logger = logging.getLogger(f"custom_nodes.comfyui-queue-manager.{component}") + logger.setLevel(logging.DEBUG) + # Inherit handlers from root logger + logger.propagate = True + + def get_logger(self, name: str) -> logging.Logger: + """Get a logger instance for a specific component. + + Args: + name: Logger name (component name) + + Returns: + Logger instance + """ + full_name = f"custom_nodes.comfyui-queue-manager.{name}" + + if full_name not in self._loggers: + logger = logging.getLogger(full_name) + logger.setLevel(logging.DEBUG) + self._loggers[full_name] = logger + + return self._loggers[full_name] + + def set_log_level(self, component: str, level: int) -> None: + """Set log level for a specific component. + + Args: + component: Component name + level: Log level + """ + logger = self.get_logger(component) + logger.setLevel(level) + + # Log the level change + logger.info(f"Log level changed to {logging.getLevelName(level)}", extra={ + "component": component, + "new_level": logging.getLevelName(level) + }) + + def cleanup_old_logs(self, days_to_keep: int = 30) -> int: + """Clean up old log files. + + Args: + days_to_keep: Number of days of logs to keep + + Returns: + Number of files cleaned up + """ + from datetime import timedelta + import time + + cutoff_time = time.time() - (days_to_keep * 24 * 60 * 60) + cleaned_count = 0 + + for log_file in self.log_dir.glob("*.log*"): + try: + if log_file.stat().st_mtime < cutoff_time: + log_file.unlink() + cleaned_count += 1 + except OSError as e: + logger = self.get_logger("logging_config") + logger.warning(f"Failed to clean up log file {log_file}: {e}") + + if cleaned_count > 0: + logger = self.get_logger("logging_config") + logger.info(f"Cleaned up {cleaned_count} old log files", extra={ + "cleaned_count": cleaned_count, + "days_to_keep": days_to_keep + }) + + return cleaned_count + + def get_log_stats(self) -> Dict[str, Any]: + """Get statistics about log files. + + Returns: + Dictionary containing log file statistics + """ + stats = { + "log_dir": str(self.log_dir), + "total_files": 0, + "total_size": 0, + "files": [] + } + + for log_file in self.log_dir.glob("*.log*"): + try: + file_stat = log_file.stat() + file_info = { + "name": log_file.name, + "size": file_stat.st_size, + "modified": datetime.fromtimestamp(file_stat.st_mtime, tz=timezone.utc).isoformat() + } + stats["files"].append(file_info) + stats["total_files"] += 1 + stats["total_size"] += file_stat.st_size + except OSError: + continue + + return stats + + +# Global logging configuration instance +_logging_config: Optional[LoggingConfig] = None + + +def get_logging_config() -> LoggingConfig: + """Get the global logging configuration instance. + + Returns: + LoggingConfig instance + """ + global _logging_config + if _logging_config is None: + _logging_config = LoggingConfig() + return _logging_config + + +def setup_logging( + log_dir: Optional[Path] = None, + console_level: int = logging.INFO, + file_level: int = logging.DEBUG, + structured: bool = True +) -> LoggingConfig: + """Set up logging for the queue manager. + + Args: + log_dir: Directory for log files + console_level: Log level for console output + file_level: Log level for file output + structured: Whether to use structured JSON logging + + Returns: + LoggingConfig instance + """ + global _logging_config + + if log_dir: + _logging_config = LoggingConfig(log_dir) + else: + _logging_config = get_logging_config() + + _logging_config.configure_logging( + console_level=console_level, + file_level=file_level, + structured=structured + ) + + return _logging_config + + +def get_logger(name: str) -> logging.Logger: + """Get a logger instance for a component. + + Args: + name: Component name + + Returns: + Logger instance + """ + config = get_logging_config() + return config.get_logger(name) \ No newline at end of file diff --git a/performance_monitor.py b/performance_monitor.py new file mode 100644 index 0000000..aacb801 --- /dev/null +++ b/performance_monitor.py @@ -0,0 +1,545 @@ +""" +Performance monitoring for the ComfyUI Queue Manager. +Tracks metrics, performance data, and system health. +""" + +from __future__ import annotations + +import functools +import logging +import psutil +import threading +import time +from collections import defaultdict, deque +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone, timedelta +from typing import Any, Callable, Dict, List, Optional, TypeVar + +from logging_config import get_logger + +logger = get_logger("performance_monitor") + +T = TypeVar('T') + + +@dataclass +class PerformanceMetric: + """Data class for performance metrics.""" + + name: str + value: float + unit: str + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + tags: Dict[str, str] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, Any]: + """Convert metric to dictionary.""" + return { + "name": self.name, + "value": self.value, + "unit": self.unit, + "timestamp": self.timestamp.isoformat(), + "tags": self.tags + } + + +@dataclass +class OperationStats: + """Statistics for a specific operation.""" + + operation_name: str + total_calls: int = 0 + total_time: float = 0.0 + min_time: float = float('inf') + max_time: float = 0.0 + error_count: int = 0 + last_called: Optional[datetime] = None + + @property + def average_time(self) -> float: + """Calculate average execution time.""" + return self.total_time / self.total_calls if self.total_calls > 0 else 0.0 + + @property + def success_rate(self) -> float: + """Calculate success rate as percentage.""" + if self.total_calls == 0: + return 0.0 + return ((self.total_calls - self.error_count) / self.total_calls) * 100 + + def update(self, execution_time: float, error: bool = False) -> None: + """Update statistics with new execution data.""" + self.total_calls += 1 + self.total_time += execution_time + self.min_time = min(self.min_time, execution_time) + self.max_time = max(self.max_time, execution_time) + self.last_called = datetime.now(timezone.utc) + + if error: + self.error_count += 1 + + def to_dict(self) -> Dict[str, Any]: + """Convert stats to dictionary.""" + return { + "operation_name": self.operation_name, + "total_calls": self.total_calls, + "total_time": self.total_time, + "average_time": self.average_time, + "min_time": self.min_time if self.min_time != float('inf') else 0.0, + "max_time": self.max_time, + "error_count": self.error_count, + "success_rate": self.success_rate, + "last_called": self.last_called.isoformat() if self.last_called else None + } + + +class PerformanceMonitor: + """Performance monitoring and metrics collection.""" + + def __init__(self, max_history_size: int = 1000): + """Initialize performance monitor. + + Args: + max_history_size: Maximum number of metrics to keep in history + """ + self.max_history_size = max_history_size + self._lock = threading.RLock() + + # Metrics storage + self._metrics_history: deque = deque(maxlen=max_history_size) + self._operation_stats: Dict[str, OperationStats] = {} + self._counters: Dict[str, int] = defaultdict(int) + self._gauges: Dict[str, float] = {} + + # System monitoring + self._system_metrics_enabled = True + self._system_metrics_interval = 60.0 # seconds + self._system_monitor_thread: Optional[threading.Thread] = None + self._stop_system_monitoring = threading.Event() + + # Performance thresholds + self._thresholds = { + "database_operation_time": 1.0, # seconds + "api_response_time": 2.0, # seconds + "workflow_execution_time": 300.0, # seconds + "memory_usage_percent": 80.0, # percent + "cpu_usage_percent": 90.0, # percent + } + + logger.info("Performance monitor initialized", extra={ + "max_history_size": max_history_size, + "system_metrics_enabled": self._system_metrics_enabled + }) + + def start_system_monitoring(self) -> None: + """Start system metrics monitoring thread.""" + if self._system_monitor_thread and self._system_monitor_thread.is_alive(): + logger.warning("System monitoring is already running") + return + + self._stop_system_monitoring.clear() + self._system_monitor_thread = threading.Thread( + target=self._system_monitoring_loop, + name="SystemMonitor", + daemon=True + ) + self._system_monitor_thread.start() + + logger.info("System monitoring started", extra={ + "interval": self._system_metrics_interval + }) + + def stop_system_monitoring(self) -> None: + """Stop system metrics monitoring thread.""" + if not self._system_monitor_thread or not self._system_monitor_thread.is_alive(): + return + + self._stop_system_monitoring.set() + self._system_monitor_thread.join(timeout=5.0) + + if self._system_monitor_thread.is_alive(): + logger.warning("System monitoring thread did not stop gracefully") + else: + logger.info("System monitoring stopped") + + def _system_monitoring_loop(self) -> None: + """System monitoring loop that runs in a separate thread.""" + logger.info("System monitoring loop started") + + while not self._stop_system_monitoring.is_set(): + try: + self._collect_system_metrics() + + # Wait for interval or stop signal + if self._stop_system_monitoring.wait(timeout=self._system_metrics_interval): + break + + except Exception as e: + logger.error(f"Error in system monitoring loop: {e}", exc_info=True) + time.sleep(5.0) # Brief pause before retrying + + logger.info("System monitoring loop stopped") + + def _collect_system_metrics(self) -> None: + """Collect system performance metrics.""" + try: + # CPU metrics + cpu_percent = psutil.cpu_percent(interval=1) + self.record_gauge("system.cpu.usage_percent", cpu_percent) + + # Memory metrics + memory = psutil.virtual_memory() + self.record_gauge("system.memory.usage_percent", memory.percent) + self.record_gauge("system.memory.available_mb", memory.available / 1024 / 1024) + self.record_gauge("system.memory.used_mb", memory.used / 1024 / 1024) + + # Disk metrics + disk = psutil.disk_usage('/') + self.record_gauge("system.disk.usage_percent", (disk.used / disk.total) * 100) + self.record_gauge("system.disk.free_gb", disk.free / 1024 / 1024 / 1024) + + # Process metrics + process = psutil.Process() + process_memory = process.memory_info() + self.record_gauge("process.memory.rss_mb", process_memory.rss / 1024 / 1024) + self.record_gauge("process.memory.vms_mb", process_memory.vms / 1024 / 1024) + self.record_gauge("process.cpu.percent", process.cpu_percent()) + self.record_gauge("process.threads.count", process.num_threads()) + + # Check thresholds + self._check_performance_thresholds({ + "cpu_usage_percent": cpu_percent, + "memory_usage_percent": memory.percent, + }) + + except Exception as e: + logger.error(f"Failed to collect system metrics: {e}") + + def _check_performance_thresholds(self, metrics: Dict[str, float]) -> None: + """Check if metrics exceed performance thresholds.""" + for metric_name, value in metrics.items(): + threshold = self._thresholds.get(metric_name) + if threshold and value > threshold: + logger.warning( + f"Performance threshold exceeded: {metric_name}", + extra={ + "metric": metric_name, + "value": value, + "threshold": threshold, + "severity": "high" if value > threshold * 1.2 else "medium" + } + ) + + def record_metric(self, metric: PerformanceMetric) -> None: + """Record a performance metric. + + Args: + metric: Performance metric to record + """ + with self._lock: + self._metrics_history.append(metric) + + # Log metric if it's significant + if metric.name.endswith("_time") and metric.value > 1.0: + logger.info(f"Performance metric recorded: {metric.name}", extra={ + "metric_name": metric.name, + "value": metric.value, + "unit": metric.unit, + "tags": metric.tags + }) + + def record_counter(self, name: str, increment: int = 1, tags: Optional[Dict[str, str]] = None) -> None: + """Record a counter metric. + + Args: + name: Counter name + increment: Amount to increment by + tags: Optional tags for the metric + """ + with self._lock: + self._counters[name] += increment + + metric = PerformanceMetric( + name=name, + value=increment, + unit="count", + tags=tags or {} + ) + self.record_metric(metric) + + def record_gauge(self, name: str, value: float, tags: Optional[Dict[str, str]] = None) -> None: + """Record a gauge metric. + + Args: + name: Gauge name + value: Gauge value + tags: Optional tags for the metric + """ + with self._lock: + self._gauges[name] = value + + metric = PerformanceMetric( + name=name, + value=value, + unit="gauge", + tags=tags or {} + ) + self.record_metric(metric) + + def record_timing(self, name: str, duration: float, tags: Optional[Dict[str, str]] = None) -> None: + """Record a timing metric. + + Args: + name: Timing metric name + duration: Duration in seconds + tags: Optional tags for the metric + """ + metric = PerformanceMetric( + name=name, + value=duration, + unit="seconds", + tags=tags or {} + ) + self.record_metric(metric) + + # Check timing thresholds + threshold_key = name.replace(".", "_") + threshold = self._thresholds.get(threshold_key) + if threshold and duration > threshold: + logger.warning( + f"Timing threshold exceeded: {name}", + extra={ + "metric": name, + "duration": duration, + "threshold": threshold, + "tags": tags or {} + } + ) + + def update_operation_stats(self, operation: str, execution_time: float, error: bool = False) -> None: + """Update statistics for an operation. + + Args: + operation: Operation name + execution_time: Execution time in seconds + error: Whether the operation resulted in an error + """ + with self._lock: + if operation not in self._operation_stats: + self._operation_stats[operation] = OperationStats(operation) + + self._operation_stats[operation].update(execution_time, error) + + # Record timing metric + self.record_timing(f"operation.{operation}.time", execution_time, { + "operation": operation, + "status": "error" if error else "success" + }) + + @contextmanager + def time_operation(self, operation: str, tags: Optional[Dict[str, str]] = None): + """Context manager for timing operations. + + Args: + operation: Operation name + tags: Optional tags for the metric + """ + start_time = time.time() + error_occurred = False + + try: + yield + except Exception as e: + error_occurred = True + raise + finally: + execution_time = time.time() - start_time + self.update_operation_stats(operation, execution_time, error_occurred) + + # Add error status to tags + final_tags = tags or {} + final_tags["status"] = "error" if error_occurred else "success" + + self.record_timing(f"operation.{operation}.time", execution_time, final_tags) + + def get_operation_stats(self, operation: Optional[str] = None) -> Dict[str, Any]: + """Get operation statistics. + + Args: + operation: Specific operation name, or None for all operations + + Returns: + Dictionary containing operation statistics + """ + with self._lock: + if operation: + stats = self._operation_stats.get(operation) + return stats.to_dict() if stats else {} + else: + return { + op_name: stats.to_dict() + for op_name, stats in self._operation_stats.items() + } + + def get_metrics_summary(self, since: Optional[datetime] = None) -> Dict[str, Any]: + """Get summary of recorded metrics. + + Args: + since: Only include metrics since this timestamp + + Returns: + Dictionary containing metrics summary + """ + with self._lock: + # Filter metrics by timestamp if specified + if since: + filtered_metrics = [ + metric for metric in self._metrics_history + if metric.timestamp >= since + ] + else: + filtered_metrics = list(self._metrics_history) + + # Group metrics by name + metrics_by_name = defaultdict(list) + for metric in filtered_metrics: + metrics_by_name[metric.name].append(metric.value) + + # Calculate summary statistics + summary = {} + for name, values in metrics_by_name.items(): + if values: + summary[name] = { + "count": len(values), + "min": min(values), + "max": max(values), + "avg": sum(values) / len(values), + "total": sum(values) if name.endswith("_count") else None + } + + return { + "period": { + "start": since.isoformat() if since else None, + "end": datetime.now(timezone.utc).isoformat(), + "total_metrics": len(filtered_metrics) + }, + "metrics": summary, + "counters": dict(self._counters), + "gauges": dict(self._gauges) + } + + def get_health_status(self) -> Dict[str, Any]: + """Get overall system health status. + + Returns: + Dictionary containing health status information + """ + with self._lock: + # Check recent errors + recent_time = datetime.now(timezone.utc) - timedelta(minutes=5) + recent_errors = sum( + 1 for metric in self._metrics_history + if (metric.timestamp >= recent_time and + metric.tags.get("status") == "error") + ) + + # Calculate error rates + total_operations = sum(stats.total_calls for stats in self._operation_stats.values()) + total_errors = sum(stats.error_count for stats in self._operation_stats.values()) + overall_error_rate = (total_errors / total_operations * 100) if total_operations > 0 else 0 + + # Determine health status + if recent_errors > 10 or overall_error_rate > 10: + health_status = "unhealthy" + elif recent_errors > 5 or overall_error_rate > 5: + health_status = "degraded" + else: + health_status = "healthy" + + return { + "status": health_status, + "timestamp": datetime.now(timezone.utc).isoformat(), + "metrics": { + "recent_errors": recent_errors, + "total_operations": total_operations, + "total_errors": total_errors, + "error_rate_percent": overall_error_rate + }, + "system": { + "cpu_percent": self._gauges.get("system.cpu.usage_percent", 0), + "memory_percent": self._gauges.get("system.memory.usage_percent", 0), + "disk_percent": self._gauges.get("system.disk.usage_percent", 0) + } + } + + def reset_stats(self) -> None: + """Reset all performance statistics.""" + with self._lock: + self._metrics_history.clear() + self._operation_stats.clear() + self._counters.clear() + self._gauges.clear() + + logger.info("Performance statistics reset") + + def close(self) -> None: + """Close the performance monitor and cleanup resources.""" + self.stop_system_monitoring() + logger.info("Performance monitor closed") + + +def time_function(operation_name: Optional[str] = None, tags: Optional[Dict[str, str]] = None): + """Decorator for timing function execution. + + Args: + operation_name: Name for the operation (defaults to function name) + tags: Optional tags for the metric + """ + def decorator(func: Callable[..., T]) -> Callable[..., T]: + @functools.wraps(func) + def wrapper(*args, **kwargs) -> T: + op_name = operation_name or func.__name__ + + with get_performance_monitor().time_operation(op_name, tags): + return func(*args, **kwargs) + + return wrapper + return decorator + + +# Global performance monitor instance +_performance_monitor: Optional[PerformanceMonitor] = None + + +def get_performance_monitor() -> PerformanceMonitor: + """Get the global performance monitor instance. + + Returns: + PerformanceMonitor instance + """ + global _performance_monitor + if _performance_monitor is None: + _performance_monitor = PerformanceMonitor() + _performance_monitor.start_system_monitoring() + return _performance_monitor + + +def setup_performance_monitoring(max_history_size: int = 1000) -> PerformanceMonitor: + """Set up performance monitoring. + + Args: + max_history_size: Maximum number of metrics to keep in history + + Returns: + PerformanceMonitor instance + """ + global _performance_monitor + + if _performance_monitor: + _performance_monitor.close() + + _performance_monitor = PerformanceMonitor(max_history_size) + _performance_monitor.start_system_monitoring() + + return _performance_monitor \ No newline at end of file diff --git a/queue_service.py b/queue_service.py index 7b6cee2..95d96ef 100644 --- a/queue_service.py +++ b/queue_service.py @@ -11,6 +11,8 @@ from typing import Any from archive_service import ArchiveService from database import SQLiteDatabase +from error_handler import ErrorHandler, with_error_handling +from exceptions import QueueServiceError, ValidationError from filter_service import FilterService from interfaces import DatabaseInterface, QueueServiceInterface from models import QueueConfig, QueueFilter, QueueItem, QueueState, QueueStatus @@ -18,10 +20,6 @@ from models import QueueConfig, QueueFilter, QueueItem, QueueState, QueueStatus logger = logging.getLogger(__name__) -class QueueServiceError(Exception): - """Custom exception for queue service operations.""" - - class QueueService(QueueServiceInterface): """Core queue service implementation.""" @@ -44,6 +42,7 @@ class QueueService(QueueServiceInterface): self._initialized = True logger.info("Queue service initialized successfully") + @with_error_handling(error_type=QueueServiceError, operation="add_workflow") def add_workflow( self, workflow_data: dict[str, Any], workflow_name: str = "" ) -> str: @@ -60,10 +59,13 @@ class QueueService(QueueServiceInterface): QueueServiceError: If the workflow cannot be added """ if not self._initialized: - raise QueueServiceError("Queue service not initialized") + raise QueueServiceError("Queue service not initialized", operation="add_workflow") if not workflow_data: - raise QueueServiceError("Workflow data cannot be empty") + raise ValidationError("Workflow data cannot be empty", field="workflow_data") + + if not isinstance(workflow_data, dict): + raise ValidationError("Workflow data must be a dictionary", field="workflow_data", value=type(workflow_data).__name__) # Create a new queue item item = QueueItem( @@ -74,11 +76,12 @@ class QueueService(QueueServiceInterface): # Save to database if not self.database.create_queue_item(item): - raise QueueServiceError(f"Failed to create queue item {item.id}") + raise QueueServiceError(f"Failed to create queue item {item.id}", item_id=item.id, operation="create") logger.info(f"Added workflow to queue: {item.id} - {item.workflow_name}") return item.id + @with_error_handling(error_type=QueueServiceError, operation="get_queue_items") def get_queue_items(self, status: QueueStatus | None = None) -> list[QueueItem]: """Get queue items, optionally filtered by status. @@ -89,16 +92,12 @@ class QueueService(QueueServiceInterface): List of queue items """ if not self._initialized: - raise QueueServiceError("Queue service not initialized") + raise QueueServiceError("Queue service not initialized", operation="get_queue_items") - try: - if status is not None: - return self.database.get_items_by_status(status) - else: - return self.database.get_all_queue_items() - except Exception as e: - logger.error(f"Failed to get queue items: {e}") - raise QueueServiceError(f"Failed to get queue items: {e}") from e + if status is not None: + return self.database.get_items_by_status(status) + else: + return self.database.get_all_queue_items() def get_queue_item(self, item_id: str) -> QueueItem | None: """Get a specific queue item by ID. diff --git a/test_error_handling.py b/test_error_handling.py new file mode 100644 index 0000000..08f403b --- /dev/null +++ b/test_error_handling.py @@ -0,0 +1,499 @@ +""" +Tests for error handling functionality in the ComfyUI Queue Manager. +""" + +import json +import pytest +import sqlite3 +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +from database import SQLiteDatabase +from error_handler import ErrorHandler, ErrorRecovery, with_database_error_handling +from exceptions import ( + APIError, + ConnectionError, + DatabaseError, + NotFoundError, + QueueManagerError, + QueueServiceError, + SchemaError, + ValidationError, + WorkflowExecutionError, +) +from models import QueueItem, QueueStatus +from queue_service import QueueService + + +class TestCustomExceptions: + """Test custom exception classes.""" + + def test_queue_manager_error_basic(self): + """Test basic QueueManagerError functionality.""" + error = QueueManagerError("Test error") + assert error.message == "Test error" + assert error.error_code == "QueueManagerError" + assert error.details == {} + assert error.cause is None + + def test_queue_manager_error_with_details(self): + """Test QueueManagerError with additional details.""" + details = {"key": "value", "number": 42} + cause = ValueError("Original error") + + error = QueueManagerError( + "Test error", + error_code="CUSTOM_ERROR", + details=details, + cause=cause + ) + + assert error.message == "Test error" + assert error.error_code == "CUSTOM_ERROR" + assert error.details == details + assert error.cause == cause + + def test_queue_manager_error_to_dict(self): + """Test QueueManagerError serialization.""" + error = QueueManagerError( + "Test error", + error_code="TEST_ERROR", + details={"field": "value"}, + cause=ValueError("Cause") + ) + + result = error.to_dict() + + assert result["error_type"] == "QueueManagerError" + assert result["error_code"] == "TEST_ERROR" + assert result["message"] == "Test error" + assert result["details"] == {"field": "value"} + assert result["cause"] == "Cause" + + def test_database_error_with_operation(self): + """Test DatabaseError with operation details.""" + error = DatabaseError( + "Database failed", + operation="INSERT", + table="queue_items" + ) + + assert error.message == "Database failed" + assert error.details["operation"] == "INSERT" + assert error.details["table"] == "queue_items" + + def test_validation_error_with_field(self): + """Test ValidationError with field details.""" + error = ValidationError( + "Invalid value", + field="workflow_name", + value="invalid" + ) + + assert error.message == "Invalid value" + assert error.details["field"] == "workflow_name" + assert error.details["value"] == "invalid" + + def test_api_error_with_status_code(self): + """Test APIError with status code.""" + error = APIError( + "API failed", + status_code=404, + endpoint="/api/items" + ) + + assert error.message == "API failed" + assert error.details["status_code"] == 404 + assert error.details["endpoint"] == "/api/items" + + def test_workflow_execution_error(self): + """Test WorkflowExecutionError with workflow details.""" + error = WorkflowExecutionError( + "Execution failed", + workflow_id="test-123", + execution_stage="validation" + ) + + assert error.message == "Execution failed" + assert error.details["workflow_id"] == "test-123" + assert error.details["execution_stage"] == "validation" + + +class TestErrorHandler: + """Test ErrorHandler utility functions.""" + + def test_handle_database_error_sqlite_locked(self): + """Test handling of SQLite database locked error.""" + original_error = sqlite3.OperationalError("database is locked") + + result = ErrorHandler.handle_database_error( + original_error, + "INSERT", + "queue_items" + ) + + assert isinstance(result, ConnectionError) + assert "Database is locked during INSERT" in result.message + assert result.details["operation"] == "INSERT" + assert result.details["table"] == "queue_items" + assert result.cause == original_error + + def test_handle_database_error_no_such_table(self): + """Test handling of SQLite no such table error.""" + original_error = sqlite3.OperationalError("no such table: missing_table") + + result = ErrorHandler.handle_database_error( + original_error, + "SELECT", + "missing_table" + ) + + assert isinstance(result, SchemaError) + assert "Table does not exist during SELECT" in result.message + assert result.details["operation"] == "SELECT" + assert result.details["table"] == "missing_table" + + def test_handle_database_error_integrity_error(self): + """Test handling of SQLite integrity error.""" + original_error = sqlite3.IntegrityError("UNIQUE constraint failed") + + result = ErrorHandler.handle_database_error( + original_error, + "INSERT", + "queue_items" + ) + + assert isinstance(result, ValidationError) + assert "Data integrity violation during INSERT" in result.message + assert result.details["operation"] == "INSERT" + assert result.details["table"] == "queue_items" + + def test_handle_database_error_generic(self): + """Test handling of generic database error.""" + original_error = Exception("Generic error") + + result = ErrorHandler.handle_database_error( + original_error, + "UPDATE", + "queue_items" + ) + + assert isinstance(result, DatabaseError) + assert "Unexpected error during UPDATE" in result.message + assert result.details["operation"] == "UPDATE" + assert result.details["table"] == "queue_items" + + def test_handle_api_error(self): + """Test handling of API errors.""" + original_error = ValueError("Invalid input") + + result = ErrorHandler.handle_api_error( + original_error, + "/api/items", + "POST" + ) + + assert isinstance(result, APIError) + assert "API error at POST /api/items" in result.message + assert result.details["endpoint"] == "/api/items" + assert result.details["method"] == "POST" + assert result.cause == original_error + + def test_handle_workflow_error(self): + """Test handling of workflow execution errors.""" + original_error = RuntimeError("Execution failed") + + result = ErrorHandler.handle_workflow_error( + original_error, + workflow_id="test-123", + stage="validation" + ) + + assert isinstance(result, WorkflowExecutionError) + assert "Workflow execution failed" in result.message + assert result.details["workflow_id"] == "test-123" + assert result.details["execution_stage"] == "validation" + assert result.cause == original_error + + +class TestErrorHandlingDecorators: + """Test error handling decorators.""" + + def test_with_database_error_handling_success(self): + """Test database error handling decorator with successful operation.""" + @with_database_error_handling(operation="test_op", table="test_table") + def successful_function(): + return "success" + + result = successful_function() + assert result == "success" + + def test_with_database_error_handling_sqlite_error(self): + """Test database error handling decorator with SQLite error.""" + @with_database_error_handling(operation="test_op", table="test_table") + def failing_function(): + raise sqlite3.OperationalError("database is locked") + + with pytest.raises(ConnectionError) as exc_info: + failing_function() + + error = exc_info.value + assert "Database is locked during test_op" in error.message + assert error.details["operation"] == "test_op" + assert error.details["table"] == "test_table" + + def test_with_database_error_handling_generic_error(self): + """Test database error handling decorator with generic error.""" + @with_database_error_handling(operation="test_op") + def failing_function(): + raise ValueError("Generic error") + + with pytest.raises(DatabaseError) as exc_info: + failing_function() + + error = exc_info.value + assert "Unexpected error during test_op" in error.message + assert error.details["operation"] == "test_op" + + def test_with_database_error_handling_reraise_queue_manager_error(self): + """Test that QueueManagerError is re-raised as-is.""" + @with_database_error_handling(operation="test_op") + def failing_function(): + raise ValidationError("Already a queue manager error") + + with pytest.raises(ValidationError) as exc_info: + failing_function() + + error = exc_info.value + assert error.message == "Already a queue manager error" + + +class TestErrorRecovery: + """Test error recovery utilities.""" + + def test_retry_with_backoff_success_first_try(self): + """Test retry mechanism with success on first try.""" + call_count = 0 + + def successful_function(): + nonlocal call_count + call_count += 1 + return "success" + + result = ErrorRecovery.retry_with_backoff(successful_function, max_retries=3) + + assert result == "success" + assert call_count == 1 + + def test_retry_with_backoff_success_after_retries(self): + """Test retry mechanism with success after retries.""" + call_count = 0 + + def eventually_successful_function(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ValueError("Temporary error") + return "success" + + with patch('time.sleep'): # Mock sleep to speed up test + result = ErrorRecovery.retry_with_backoff( + eventually_successful_function, + max_retries=3, + base_delay=0.1 + ) + + assert result == "success" + assert call_count == 3 + + def test_retry_with_backoff_all_retries_fail(self): + """Test retry mechanism when all retries fail.""" + call_count = 0 + + def always_failing_function(): + nonlocal call_count + call_count += 1 + raise ValueError(f"Error {call_count}") + + with patch('time.sleep'): # Mock sleep to speed up test + with pytest.raises(ValueError) as exc_info: + ErrorRecovery.retry_with_backoff( + always_failing_function, + max_retries=2, + base_delay=0.1 + ) + + assert "Error 3" in str(exc_info.value) # Last error + assert call_count == 3 # Initial + 2 retries + + def test_safe_execute_success(self): + """Test safe execution with successful function.""" + def successful_function(): + return "success" + + result = ErrorRecovery.safe_execute(successful_function) + assert result == "success" + + def test_safe_execute_with_error(self): + """Test safe execution with error and default value.""" + def failing_function(): + raise ValueError("Error") + + result = ErrorRecovery.safe_execute(failing_function, default_value="default") + assert result == "default" + + def test_safe_execute_with_error_no_default(self): + """Test safe execution with error and no default value.""" + def failing_function(): + raise ValueError("Error") + + result = ErrorRecovery.safe_execute(failing_function) + assert result is None + + +class TestDatabaseErrorHandling: + """Test error handling in database operations.""" + + def setup_method(self): + """Set up test database.""" + self.temp_dir = tempfile.mkdtemp() + self.db_path = Path(self.temp_dir) / "test.db" + self.database = SQLiteDatabase(self.db_path) + + def teardown_method(self): + """Clean up test database.""" + if hasattr(self, 'database'): + self.database.close() + + def test_database_initialization_error_handling(self): + """Test error handling during database initialization.""" + # Create a database with invalid path (directory instead of file) + invalid_path = Path(self.temp_dir) / "invalid_dir" + invalid_path.mkdir() + + database = SQLiteDatabase(invalid_path) + + # This should handle the error gracefully + with pytest.raises(DatabaseError): + database.initialize() + + def test_create_queue_item_validation_error(self): + """Test validation error when creating queue item.""" + self.database.initialize() + + # Create item with empty ID + item = QueueItem(id="", workflow_name="test") + + with pytest.raises(ValidationError) as exc_info: + self.database.create_queue_item(item) + + error = exc_info.value + assert "Queue item ID cannot be empty" in error.message + assert error.details["field"] == "id" + + def test_get_queue_item_validation_error(self): + """Test validation error when getting queue item.""" + self.database.initialize() + + with pytest.raises(ValidationError) as exc_info: + self.database.get_queue_item("") + + error = exc_info.value + assert "Item ID cannot be empty" in error.message + assert error.details["field"] == "item_id" + + def test_database_connection_error_handling(self): + """Test database connection error handling.""" + # Create database with invalid permissions + invalid_db = SQLiteDatabase("/root/invalid.db") # Should fail on most systems + + with pytest.raises(DatabaseError): + invalid_db._get_connection() + + +class TestQueueServiceErrorHandling: + """Test error handling in queue service operations.""" + + def setup_method(self): + """Set up test queue service.""" + self.temp_dir = tempfile.mkdtemp() + self.db_path = Path(self.temp_dir) / "test.db" + self.database = SQLiteDatabase(self.db_path) + self.database.initialize() + self.queue_service = QueueService(self.database) + + def teardown_method(self): + """Clean up test resources.""" + if hasattr(self, 'queue_service'): + self.queue_service.close() + + def test_add_workflow_validation_error_empty_data(self): + """Test validation error when adding workflow with empty data.""" + with pytest.raises(ValidationError) as exc_info: + self.queue_service.add_workflow({}) + + error = exc_info.value + assert "Workflow data cannot be empty" in error.message + assert error.details["field"] == "workflow_data" + + def test_add_workflow_validation_error_invalid_type(self): + """Test validation error when adding workflow with invalid data type.""" + with pytest.raises(ValidationError) as exc_info: + self.queue_service.add_workflow("not a dict") + + error = exc_info.value + assert "Workflow data must be a dictionary" in error.message + assert error.details["field"] == "workflow_data" + assert error.details["value"] == "str" + + def test_queue_service_not_initialized_error(self): + """Test error when using uninitialized queue service.""" + # Create service without initializing + uninit_service = QueueService.__new__(QueueService) + uninit_service._initialized = False + + with pytest.raises(QueueServiceError) as exc_info: + uninit_service.add_workflow({"test": "data"}) + + error = exc_info.value + assert "Queue service not initialized" in error.message + assert error.details["operation"] == "add_workflow" + + +class TestWorkflowExecutorErrorHandling: + """Test error handling in workflow executor.""" + + def setup_method(self): + """Set up test workflow executor.""" + from workflow_executor import WorkflowExecutor + self.executor = WorkflowExecutor() + + def teardown_method(self): + """Clean up test resources.""" + if hasattr(self, 'executor'): + self.executor.close() + + def test_execute_workflow_validation_error_empty_data(self): + """Test validation error when executing workflow with empty data.""" + with pytest.raises(ValidationError) as exc_info: + self.executor.execute_workflow({}) + + error = exc_info.value + assert "Workflow data cannot be empty" in error.message + assert error.details["field"] == "workflow_data" + + def test_execute_workflow_validation_error_invalid_type(self): + """Test validation error when executing workflow with invalid data type.""" + with pytest.raises(ValidationError) as exc_info: + self.executor.execute_workflow("not a dict") + + error = exc_info.value + assert "Workflow data must be a dictionary" in error.message + assert error.details["field"] == "workflow_data" + assert error.details["value"] == "str" + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/test_logging_monitoring.py b/test_logging_monitoring.py new file mode 100644 index 0000000..89a09a0 --- /dev/null +++ b/test_logging_monitoring.py @@ -0,0 +1,602 @@ +""" +Tests for logging and monitoring functionality in the ComfyUI Queue Manager. +""" + +import asyncio +import json +import logging +import tempfile +import threading +import time +from datetime import datetime, timezone, timedelta +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + +from health_check import ( + DatabaseHealthCheck, + HealthCheck, + HealthCheckManager, + HealthCheckResult, + HealthStatus, + QueueServiceHealthCheck, + SystemResourceHealthCheck, +) +from logging_config import LoggingConfig, QueueManagerFilter, StructuredFormatter +from performance_monitor import ( + OperationStats, + PerformanceMetric, + PerformanceMonitor, + time_function, +) + + +class TestStructuredFormatter: + """Test structured logging formatter.""" + + def test_basic_formatting(self): + """Test basic log record formatting.""" + formatter = StructuredFormatter() + + # Create a log record + record = logging.LogRecord( + name="test.logger", + level=logging.INFO, + pathname="/test/path.py", + lineno=42, + msg="Test message", + args=(), + exc_info=None + ) + + # Format the record + formatted = formatter.format(record) + + # Parse the JSON + log_data = json.loads(formatted) + + assert log_data["level"] == "INFO" + assert log_data["logger"] == "test.logger" + assert log_data["message"] == "Test message" + assert log_data["line"] == 42 + assert "timestamp" in log_data + + def test_formatting_with_extra_fields(self): + """Test formatting with extra fields.""" + formatter = StructuredFormatter(include_extra=True) + + record = logging.LogRecord( + name="test.logger", + level=logging.ERROR, + pathname="/test/path.py", + lineno=42, + msg="Test error", + args=(), + exc_info=None + ) + + # Add extra fields + record.user_id = "test-user" + record.operation = "test_operation" + record.custom_data = {"key": "value"} + + formatted = formatter.format(record) + log_data = json.loads(formatted) + + assert "extra" in log_data + assert log_data["extra"]["user_id"] == "test-user" + assert log_data["extra"]["operation"] == "test_operation" + assert log_data["extra"]["custom_data"] == {"key": "value"} + + def test_formatting_with_exception(self): + """Test formatting with exception information.""" + formatter = StructuredFormatter() + + try: + raise ValueError("Test exception") + except ValueError: + record = logging.LogRecord( + name="test.logger", + level=logging.ERROR, + pathname="/test/path.py", + lineno=42, + msg="Error occurred", + args=(), + exc_info=True + ) + + formatted = formatter.format(record) + log_data = json.loads(formatted) + + assert "exception" in log_data + assert log_data["exception"]["type"] == "ValueError" + assert log_data["exception"]["message"] == "Test exception" + assert "traceback" in log_data["exception"] + + +class TestQueueManagerFilter: + """Test logging filter.""" + + def test_level_filtering(self): + """Test filtering by log level.""" + filter_obj = QueueManagerFilter(min_level=logging.WARNING) + + # Create records with different levels + info_record = logging.LogRecord( + name="test", level=logging.INFO, pathname="", lineno=0, + msg="Info message", args=(), exc_info=None + ) + warning_record = logging.LogRecord( + name="test", level=logging.WARNING, pathname="", lineno=0, + msg="Warning message", args=(), exc_info=None + ) + + assert not filter_obj.filter(info_record) + assert filter_obj.filter(warning_record) + + def test_component_filtering(self): + """Test filtering by component.""" + filter_obj = QueueManagerFilter(component="database") + + # Create records with different logger names + db_record = logging.LogRecord( + name="queue_manager.database", level=logging.INFO, pathname="", lineno=0, + msg="Database message", args=(), exc_info=None + ) + api_record = logging.LogRecord( + name="queue_manager.api", level=logging.INFO, pathname="", lineno=0, + msg="API message", args=(), exc_info=None + ) + + assert filter_obj.filter(db_record) + assert not filter_obj.filter(api_record) + + +class TestLoggingConfig: + """Test logging configuration.""" + + def setup_method(self): + """Set up test logging configuration.""" + self.temp_dir = Path(tempfile.mkdtemp()) + self.log_config = LoggingConfig(log_dir=self.temp_dir) + + def test_configure_logging(self): + """Test logging configuration setup.""" + self.log_config.configure_logging( + console_level=logging.INFO, + file_level=logging.DEBUG, + structured=True + ) + + # Check that log files are created + assert (self.temp_dir / "queue_manager.log").exists() + assert (self.temp_dir / "errors.log").exists() + assert (self.temp_dir / "performance.log").exists() + + def test_get_logger(self): + """Test getting component loggers.""" + self.log_config.configure_logging() + + logger = self.log_config.get_logger("test_component") + + assert logger.name == "custom_nodes.comfyui-queue-manager.test_component" + assert logger.level == logging.DEBUG + + def test_set_log_level(self): + """Test setting log level for components.""" + self.log_config.configure_logging() + + self.log_config.set_log_level("test_component", logging.WARNING) + logger = self.log_config.get_logger("test_component") + + assert logger.level == logging.WARNING + + def test_cleanup_old_logs(self): + """Test cleaning up old log files.""" + self.log_config.configure_logging() + + # Create some old log files + old_log = self.temp_dir / "old.log" + old_log.touch() + + # Set modification time to 31 days ago + old_time = time.time() - (31 * 24 * 60 * 60) + old_log.stat().st_mtime = old_time + + cleaned_count = self.log_config.cleanup_old_logs(days_to_keep=30) + + # Note: This test might not work on all systems due to file system limitations + # The cleanup_old_logs method should handle this gracefully + assert cleaned_count >= 0 + + def test_get_log_stats(self): + """Test getting log file statistics.""" + self.log_config.configure_logging() + + stats = self.log_config.get_log_stats() + + assert "log_dir" in stats + assert "total_files" in stats + assert "total_size" in stats + assert "files" in stats + assert stats["total_files"] >= 0 + + +class TestPerformanceMetric: + """Test performance metric data class.""" + + def test_metric_creation(self): + """Test creating performance metrics.""" + metric = PerformanceMetric( + name="test.metric", + value=1.5, + unit="seconds", + tags={"operation": "test"} + ) + + assert metric.name == "test.metric" + assert metric.value == 1.5 + assert metric.unit == "seconds" + assert metric.tags["operation"] == "test" + assert isinstance(metric.timestamp, datetime) + + def test_metric_to_dict(self): + """Test converting metric to dictionary.""" + metric = PerformanceMetric( + name="test.metric", + value=2.0, + unit="count" + ) + + result = metric.to_dict() + + assert result["name"] == "test.metric" + assert result["value"] == 2.0 + assert result["unit"] == "count" + assert "timestamp" in result + + +class TestOperationStats: + """Test operation statistics.""" + + def test_stats_initialization(self): + """Test initializing operation stats.""" + stats = OperationStats("test_operation") + + assert stats.operation_name == "test_operation" + assert stats.total_calls == 0 + assert stats.total_time == 0.0 + assert stats.error_count == 0 + assert stats.average_time == 0.0 + assert stats.success_rate == 0.0 + + def test_stats_update(self): + """Test updating operation stats.""" + stats = OperationStats("test_operation") + + # Update with successful operation + stats.update(1.5, error=False) + + assert stats.total_calls == 1 + assert stats.total_time == 1.5 + assert stats.average_time == 1.5 + assert stats.min_time == 1.5 + assert stats.max_time == 1.5 + assert stats.error_count == 0 + assert stats.success_rate == 100.0 + + # Update with failed operation + stats.update(2.0, error=True) + + assert stats.total_calls == 2 + assert stats.total_time == 3.5 + assert stats.average_time == 1.75 + assert stats.min_time == 1.5 + assert stats.max_time == 2.0 + assert stats.error_count == 1 + assert stats.success_rate == 50.0 + + def test_stats_to_dict(self): + """Test converting stats to dictionary.""" + stats = OperationStats("test_operation") + stats.update(1.0, error=False) + + result = stats.to_dict() + + assert result["operation_name"] == "test_operation" + assert result["total_calls"] == 1 + assert result["average_time"] == 1.0 + assert result["success_rate"] == 100.0 + + +class TestPerformanceMonitor: + """Test performance monitoring.""" + + def setup_method(self): + """Set up test performance monitor.""" + self.monitor = PerformanceMonitor(max_history_size=100) + + def teardown_method(self): + """Clean up test resources.""" + if hasattr(self, 'monitor'): + self.monitor.close() + + def test_record_metric(self): + """Test recording performance metrics.""" + metric = PerformanceMetric("test.metric", 1.0, "seconds") + + self.monitor.record_metric(metric) + + # Check that metric was recorded + assert len(self.monitor._metrics_history) == 1 + assert self.monitor._metrics_history[0] == metric + + def test_record_counter(self): + """Test recording counter metrics.""" + self.monitor.record_counter("test.counter", increment=5) + + assert self.monitor._counters["test.counter"] == 5 + assert len(self.monitor._metrics_history) == 1 + + def test_record_gauge(self): + """Test recording gauge metrics.""" + self.monitor.record_gauge("test.gauge", 42.0) + + assert self.monitor._gauges["test.gauge"] == 42.0 + assert len(self.monitor._metrics_history) == 1 + + def test_record_timing(self): + """Test recording timing metrics.""" + self.monitor.record_timing("test.timing", 1.5) + + assert len(self.monitor._metrics_history) == 1 + metric = self.monitor._metrics_history[0] + assert metric.name == "test.timing" + assert metric.value == 1.5 + assert metric.unit == "seconds" + + def test_update_operation_stats(self): + """Test updating operation statistics.""" + self.monitor.update_operation_stats("test_op", 1.0, error=False) + + stats = self.monitor.get_operation_stats("test_op") + assert stats["operation_name"] == "test_op" + assert stats["total_calls"] == 1 + assert stats["average_time"] == 1.0 + + def test_time_operation_context_manager(self): + """Test timing operation context manager.""" + with self.monitor.time_operation("test_context"): + time.sleep(0.01) # Small delay + + stats = self.monitor.get_operation_stats("test_context") + assert stats["total_calls"] == 1 + assert stats["average_time"] > 0 + + def test_time_operation_with_error(self): + """Test timing operation with error.""" + with pytest.raises(ValueError): + with self.monitor.time_operation("test_error"): + raise ValueError("Test error") + + stats = self.monitor.get_operation_stats("test_error") + assert stats["total_calls"] == 1 + assert stats["error_count"] == 1 + assert stats["success_rate"] == 0.0 + + def test_get_metrics_summary(self): + """Test getting metrics summary.""" + # Record some metrics + self.monitor.record_timing("test.timing", 1.0) + self.monitor.record_timing("test.timing", 2.0) + self.monitor.record_counter("test.counter", 5) + + summary = self.monitor.get_metrics_summary() + + assert "metrics" in summary + assert "test.timing" in summary["metrics"] + assert summary["metrics"]["test.timing"]["count"] == 2 + assert summary["metrics"]["test.timing"]["avg"] == 1.5 + + def test_get_health_status(self): + """Test getting health status.""" + # Record some operations + self.monitor.update_operation_stats("test_op", 1.0, error=False) + self.monitor.update_operation_stats("test_op", 1.5, error=False) + + health = self.monitor.get_health_status() + + assert "status" in health + assert "metrics" in health + assert health["metrics"]["total_operations"] == 2 + assert health["metrics"]["total_errors"] == 0 + + @patch('psutil.cpu_percent') + @patch('psutil.virtual_memory') + @patch('psutil.disk_usage') + @patch('psutil.Process') + def test_collect_system_metrics(self, mock_process, mock_disk, mock_memory, mock_cpu): + """Test collecting system metrics.""" + # Mock system metrics + mock_cpu.return_value = 50.0 + mock_memory.return_value = Mock(percent=60.0, available=1024*1024*1024, used=512*1024*1024) + mock_disk.return_value = Mock(total=1024*1024*1024*100, used=1024*1024*1024*50, free=1024*1024*1024*50) + mock_process.return_value = Mock( + memory_info=Mock(rss=1024*1024*100, vms=1024*1024*200), + cpu_percent=Mock(return_value=25.0), + num_threads=Mock(return_value=10) + ) + + self.monitor._collect_system_metrics() + + # Check that system metrics were recorded + assert "system.cpu.usage_percent" in self.monitor._gauges + assert "system.memory.usage_percent" in self.monitor._gauges + assert self.monitor._gauges["system.cpu.usage_percent"] == 50.0 + assert self.monitor._gauges["system.memory.usage_percent"] == 60.0 + + +class TestTimeFunctionDecorator: + """Test time_function decorator.""" + + def setup_method(self): + """Set up test performance monitor.""" + self.monitor = PerformanceMonitor() + + def teardown_method(self): + """Clean up test resources.""" + if hasattr(self, 'monitor'): + self.monitor.close() + + def test_time_function_decorator(self): + """Test timing function with decorator.""" + @time_function("test_function") + def test_func(): + time.sleep(0.01) + return "result" + + result = test_func() + + assert result == "result" + + # Check that timing was recorded + from performance_monitor import get_performance_monitor + monitor = get_performance_monitor() + stats = monitor.get_operation_stats("test_function") + assert stats["total_calls"] >= 1 + + def test_time_function_with_error(self): + """Test timing function with error.""" + @time_function("test_error_function") + def error_func(): + raise ValueError("Test error") + + with pytest.raises(ValueError): + error_func() + + # Check that error was recorded + from performance_monitor import get_performance_monitor + monitor = get_performance_monitor() + stats = monitor.get_operation_stats("test_error_function") + assert stats["total_calls"] >= 1 + assert stats["error_count"] >= 1 + + +class TestHealthCheck: + """Test health check base class.""" + + def test_health_check_timeout(self): + """Test health check timeout handling.""" + class SlowHealthCheck(HealthCheck): + async def _perform_check(self): + await asyncio.sleep(2.0) # Longer than timeout + return HealthCheckResult( + name=self.name, + status=HealthStatus.HEALTHY, + message="Should not reach here" + ) + + check = SlowHealthCheck("slow_check", timeout=0.1) + + # Run the check + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(check.check()) + finally: + loop.close() + + assert result.status == HealthStatus.UNHEALTHY + assert "timed out" in result.message.lower() + + def test_health_check_exception(self): + """Test health check exception handling.""" + class FailingHealthCheck(HealthCheck): + async def _perform_check(self): + raise ValueError("Test error") + + check = FailingHealthCheck("failing_check") + + # Run the check + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(check.check()) + finally: + loop.close() + + assert result.status == HealthStatus.UNHEALTHY + assert "Test error" in result.message + + +class TestDatabaseHealthCheck: + """Test database health check.""" + + def test_database_health_check_uninitialized(self): + """Test health check with uninitialized database.""" + mock_db = Mock() + mock_db._initialized = False + + check = DatabaseHealthCheck(mock_db) + + # Run the check + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(check.check()) + finally: + loop.close() + + assert result.status == HealthStatus.UNHEALTHY + assert "not initialized" in result.message.lower() + + +class TestHealthCheckManager: + """Test health check manager.""" + + def setup_method(self): + """Set up test health check manager.""" + self.manager = HealthCheckManager() + + def teardown_method(self): + """Clean up test resources.""" + if hasattr(self, 'manager'): + self.manager.close() + + def test_register_and_unregister_check(self): + """Test registering and unregistering health checks.""" + check = SystemResourceHealthCheck() + + self.manager.register_check(check) + assert "system_resources" in self.manager._checks + + self.manager.unregister_check("system_resources") + assert "system_resources" not in self.manager._checks + + def test_run_nonexistent_check(self): + """Test running a non-existent health check.""" + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(self.manager.run_check("nonexistent")) + finally: + loop.close() + + assert result.status == HealthStatus.UNKNOWN + assert "not found" in result.message.lower() + + def test_get_overall_health_no_checks(self): + """Test getting overall health with no checks.""" + health = self.manager.get_overall_health() + + assert health["status"] == HealthStatus.UNKNOWN.value + assert "no health checks" in health["message"].lower() + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/workflow_executor.py b/workflow_executor.py index dbf87c0..6c05c0f 100644 --- a/workflow_executor.py +++ b/workflow_executor.py @@ -13,16 +13,14 @@ import uuid from datetime import datetime, timezone from typing import Any, Callable, Dict, Optional +from error_handler import ErrorHandler, with_error_handling +from exceptions import ValidationError, WorkflowExecutionError from interfaces import WorkflowExecutorInterface from models import QueueStatus logger = logging.getLogger(__name__) -class WorkflowExecutionError(Exception): - """Custom exception for workflow execution errors.""" - - class WorkflowExecutor(WorkflowExecutorInterface): """ Handles workflow execution integration with ComfyUI. @@ -166,6 +164,7 @@ class WorkflowExecutor(WorkflowExecutorInterface): return workflow_id + @with_error_handling(error_type=WorkflowExecutionError, operation="execute_workflow") def execute_workflow(self, workflow_data: Dict[str, Any]) -> Dict[str, Any]: """Execute a workflow and return the results. @@ -178,21 +177,27 @@ class WorkflowExecutor(WorkflowExecutorInterface): Raises: WorkflowExecutionError: If execution fails """ + if not workflow_data: + raise ValidationError("Workflow data cannot be empty", field="workflow_data") + + if not isinstance(workflow_data, dict): + raise ValidationError("Workflow data must be a dictionary", field="workflow_data", value=type(workflow_data).__name__) + workflow_id = str(uuid.uuid4()) + # Create workflow info for tracking + with self._lock: + self._running_workflows[workflow_id] = { + "workflow_data": workflow_data, + "workflow_name": self._extract_workflow_name(workflow_data), + "status": QueueStatus.RUNNING, + "started_at": datetime.now(timezone.utc), + "completed_at": None, + "error_message": None, + "result_data": None, + } + try: - # Create workflow info for tracking - with self._lock: - self._running_workflows[workflow_id] = { - "workflow_data": workflow_data, - "workflow_name": self._extract_workflow_name(workflow_data), - "status": QueueStatus.RUNNING, - "started_at": datetime.now(timezone.utc), - "completed_at": None, - "error_message": None, - "result_data": None, - } - # Execute the workflow execution_id = self._execute_workflow_directly(workflow_data) @@ -221,7 +226,12 @@ class WorkflowExecutor(WorkflowExecutorInterface): QueueStatus.FAILED, error_message=str(e) ) - raise WorkflowExecutionError(f"Workflow execution failed: {e}") from e + raise WorkflowExecutionError( + f"Workflow execution failed: {e}", + workflow_id=workflow_id, + execution_stage="execution", + cause=e + ) def _execute_workflow_directly(self, workflow_data: Dict[str, Any]) -> str: """Execute a workflow directly through ComfyUI.