"""Schedule configuration for cleanup workflow.""" import os from datetime import timedelta from typing import Any from sientia_do.observability.logger import Logger as SientiaLogger from temporalio.client import ( Client, Schedule, ScheduleActionStartWorkflow, ScheduleSpec, ) from model_manager.worker.prepare_worker import build_queue_name # Schedule configuration from environment variables CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 0 * * *') # Default: midnight UTC CLEANUP_TIMEZONE = os.getenv('CLEANUP_TIMEZONE', 'UTC') CLEANUP_EXECUTION_TIMEOUT_HOURS = int(os.getenv('CLEANUP_EXECUTION_TIMEOUT_HOURS', '1')) def build_cleanup_schedule_id(runtime: str | None) -> str: """ Build cleanup schedule ID using runtime-derived naming. Args: runtime: Runtime suffix used by workers Return: str: Cleanup schedule ID """ normalized_runtime = runtime.strip() if runtime else '' return f'cleanup-files-{normalized_runtime or "single"}-daily' async def _needs_schedule_reconcile( schedule_handle: Any, cleanup_task_queue: str, logger: SientiaLogger, metadata: dict[str, str | None], ) -> bool: """ Compare configured cleanup schedule against current expected values. Args: schedule_handle: Temporal schedule handle for current schedule ID cleanup_task_queue: Expected cleanup task queue logger: Logger instance for errors metadata: Metadata dictionary for logging context Return: bool: True when schedule should be recreated to apply current config """ try: schedule_description = await schedule_handle.describe() schedule = getattr(schedule_description, 'schedule', None) action = getattr(schedule, 'action', None) spec = getattr(schedule, 'spec', None) current_task_queue = getattr(action, 'task_queue', None) current_execution_timeout = getattr(action, 'execution_timeout', None) current_cron = getattr(spec, 'cron_expressions', None) current_timezone = getattr(spec, 'time_zone_name', None) return ( current_task_queue != cleanup_task_queue or current_execution_timeout != timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS) or current_cron != [CLEANUP_CRON] or current_timezone != CLEANUP_TIMEZONE ) except Exception as e: # noqa: BLE001 logger.custom_error(f'Error describing cleanup schedule for reconcile: {e}', metadata) return True async def schedule_exists( client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None] ) -> bool: """ Check if a schedule already exists. Args: client: Temporal client instance schedule_id: ID of the schedule to check logger: Logger instance for error logging metadata: Metadata dictionary for logging context Returns: True if schedule exists, False otherwise """ try: async for schedule in await client.list_schedules(): if schedule.id == schedule_id: return True return False except Exception as e: # noqa: BLE001 logger.custom_error(f'Error checking if schedule exists: {e}', metadata) return False async def create_cleanup_schedule( client: Client, logger: SientiaLogger, metadata: dict[str, str | None] ) -> None: """ Create or update the cleanup files schedule. This function is idempotent and can be called multiple times safely. It will only create the schedule if it doesn't already exist. Args: client: Temporal client instance logger: Logger instance for logging schedule operations metadata: Metadata dictionary for logging context """ runtime = (os.getenv('RUNTIME') or 'single').strip() cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single') schedule_id = build_cleanup_schedule_id(runtime) updated = False if await schedule_exists(client, schedule_id, logger, metadata): handle = client.get_schedule_handle(schedule_id) if await _needs_schedule_reconcile(handle, cleanup_task_queue, logger, metadata): await handle.delete() updated = True else: logger.custom_info( f"Schedule '{schedule_id}' is already up to date, no-op reconcile", metadata, ) return await client.create_schedule( schedule_id, Schedule( action=ScheduleActionStartWorkflow( 'cleanup_files', {}, # Empty input, will use default bucket from environment id=f'cleanup-files-scheduled-{schedule_id}', task_queue=cleanup_task_queue, execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS), ), spec=ScheduleSpec( cron_expressions=[CLEANUP_CRON], time_zone_name=CLEANUP_TIMEZONE, ), ), ) if updated: logger.custom_info( f"Schedule '{schedule_id}' reconciled successfully. " f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', metadata, ) else: logger.custom_info( f"Schedule '{schedule_id}' created successfully. " f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})', metadata, )