""" Retry handling utilities for the owned clip analysis backfill job. This module provides centralized retry logic with exponential backoff, jitter, and specific handling for different types of errors. """ import time import random from typing import Callable, TypeVar, Any, Optional from functools import wraps from config.job_config import JobConfig, RetryConfig from lib.exceptions import ( DeadlockError, ConnectionError, ForeignKeyConstraintError, RetryExhaustedException, classify_database_error, is_retryable_error ) from lib.metrics_collector import MetricsCollector T = TypeVar('T') class RetryHandler: """ Centralized retry handling with exponential backoff and jitter. This class provides various retry strategies for different types of operations, with configurable policies and comprehensive logging. """ def __init__(self, config: JobConfig, metrics_collector: Optional[MetricsCollector] = None): """ Initialize the retry handler. Args: config: Job configuration containing retry settings metrics_collector: Optional metrics collector for logging """ self.config = config self.retry_config = config.retry self.metrics_collector = metrics_collector def _calculate_delay(self, retry_count: int, base_delay: Optional[float] = None) -> float: """ Calculate delay with exponential backoff and jitter. Args: retry_count: Current retry attempt (0-based) base_delay: Optional override for base delay Returns: Delay in seconds """ if base_delay is None: base_delay = self.retry_config.base_delay # Exponential backoff delay = base_delay * (self.retry_config.backoff_multiplier ** retry_count) # Cap at max_delay delay = min(delay, self.retry_config.max_delay) # Add jitter to prevent thundering herd jitter = random.uniform(-self.retry_config.jitter_range, self.retry_config.jitter_range) delay = max(0, delay + jitter) return delay def _log_retry_attempt(self, operation_name: str, retry_count: int, max_retries: int, error: Exception, delay: float, **context): """ Log retry attempt with context. Args: operation_name: Name of the operation being retried retry_count: Current retry attempt (1-based) max_retries: Maximum number of retries error: The error that triggered the retry delay: Delay before next attempt **context: Additional context for logging """ if self.metrics_collector: self.metrics_collector.log_warning( f"Retrying {operation_name}", retry_attempt=retry_count, max_retries=max_retries, delay_seconds=round(delay, 2), error_type=type(error).__name__, error_snippet=str(error)[:200], **context ) def with_retry(self, operation: Callable[[], T], operation_name: str = "operation", custom_max_retries: Optional[int] = None, custom_is_retryable: Optional[Callable[[Exception], bool]] = None, **context) -> T: """ Execute an operation with retry logic. Args: operation: Function to execute operation_name: Human-readable name for logging custom_max_retries: Override default max retries custom_is_retryable: Custom function to determine if error is retryable **context: Additional context for logging Returns: Result of the operation Raises: RetryExhaustedException: If all retry attempts are exhausted """ max_retries = custom_max_retries or self.retry_config.max_retries is_retryable_fn = custom_is_retryable or is_retryable_error last_error = None for attempt in range(max_retries + 1): # +1 for initial attempt try: return operation() except Exception as e: last_error = e # Classify database errors for better handling if hasattr(e, '__class__') and 'database' in str(e.__class__.__module__).lower(): e = classify_database_error(e, context) # Don't retry on the last attempt if attempt == max_retries: break # Check if error is retryable if not is_retryable_fn(e): break # Calculate delay and wait delay = self._calculate_delay(attempt) self._log_retry_attempt(operation_name, attempt + 1, max_retries, e, delay, **context) if delay > 0: time.sleep(delay) # All retries exhausted raise RetryExhaustedException( f"Operation '{operation_name}' failed after {max_retries} retries", max_retries=max_retries, last_error=last_error, **context ) def with_database_retry(self, operation: Callable[[], T], operation_name: str = "database_operation", partition_id: Optional[str] = None, **context) -> T: """ Execute a database operation with specialized retry logic. This method provides enhanced handling for database-specific errors like deadlocks and connection issues. Args: operation: Database operation to execute operation_name: Human-readable name for logging partition_id: Optional partition ID for context **context: Additional context for logging Returns: Result of the operation """ def is_database_retryable(error: Exception) -> bool: """Check if a database error is retryable.""" if isinstance(error, (DeadlockError, ConnectionError)): return True elif isinstance(error, ForeignKeyConstraintError): # FK errors might be retryable in case of race conditions return True return False context_with_partition = context.copy() if partition_id: context_with_partition["partition_id"] = partition_id return self.with_retry( operation=operation, operation_name=operation_name, custom_is_retryable=is_database_retryable, **context_with_partition ) def retry_decorator(self, operation_name: Optional[str] = None, max_retries: Optional[int] = None): """ Decorator for automatic retry handling. Args: operation_name: Optional name for the operation (uses function name if not provided) max_retries: Optional override for max retries Returns: Decorated function with retry capability """ def decorator(func: Callable[..., T]) -> Callable[..., T]: @wraps(func) def wrapper(*args, **kwargs) -> T: nonlocal operation_name if operation_name is None: operation_name = func.__name__ def operation(): return func(*args, **kwargs) return self.with_retry( operation=operation, operation_name=operation_name, custom_max_retries=max_retries ) return wrapper return decorator def handle_connection_refresh(self, refresh_callback: Callable[[], None], operation_name: str = "connection_refresh", **context): """ Handle connection refresh with retry logic. Args: refresh_callback: Function to call for refreshing connections operation_name: Name for logging **context: Additional context """ def refresh_operation(): refresh_callback() return None try: self.with_retry( operation=refresh_operation, operation_name=operation_name, custom_max_retries=2, # Shorter retry for connection refresh **context ) except RetryExhaustedException as e: if self.metrics_collector: self.metrics_collector.log_error( f"Failed to refresh connections after retries", operation_name=operation_name, **context ) raise def handle_deadlock(self, retry_count: int, **context): """ Handle deadlock-specific logic. Args: retry_count: Current retry attempt **context: Additional context for logging """ delay = self._calculate_delay(retry_count) if self.metrics_collector: self.metrics_collector.log_warning( "Database deadlock detected - will retry", retry_count=retry_count + 1, max_retries=self.retry_config.max_retries, delay_seconds=round(delay, 2), **context ) if delay > 0: time.sleep(delay) def handle_race_condition(self, revalidation_callback: Callable[[], T], operation_name: str = "race_condition_recovery", **context) -> T: """ Handle race condition recovery with revalidation. Args: revalidation_callback: Function to call for revalidation and retry operation_name: Name for logging **context: Additional context Returns: Result of the revalidation operation """ if self.metrics_collector: self.metrics_collector.log_warning( "Race condition detected - performing revalidation", operation_name=operation_name, **context ) return revalidation_callback() def create_retry_handler(config: JobConfig, metrics_collector: Optional[MetricsCollector] = None) -> RetryHandler: """ Factory function to create a RetryHandler instance. Args: config: Job configuration metrics_collector: Optional metrics collector Returns: Configured RetryHandler instance """ return RetryHandler(config, metrics_collector)