sientia-dataops-model-manager/model_manager/schedules/cleanup_schedule.py

"""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,
        )