mockupAWS/src/tasks/cleanup.py

"""Background cleanup tasks."""

import asyncio
from datetime import datetime, timedelta
from uuid import UUID
import time

from celery import shared_task

from src.core.celery_app import celery_app
from src.core.database import AsyncSessionLocal
from src.core.cache import cache_manager
from src.core.logging_config import get_logger, set_correlation_id
from src.core.monitoring import metrics
from src.repositories.report import report_repository
from src.services.report_service import report_service


logger = get_logger(__name__)


@celery_app.task(
    bind=True,
    time_limit=1800,  # 30 minutes
    rate_limit="1/h",  # Run once per hour
)
def cleanup_old_reports(self, max_age_days: int = 30):
    """Clean up old report files and database entries.

    Args:
        max_age_days: Maximum age of reports in days
    """
    correlation_id = set_correlation_id()
    start_time = datetime.utcnow()

    logger.info(
        "Starting old reports cleanup",
        extra={"max_age_days": max_age_days, "correlation_id": correlation_id},
    )

    try:
        # Run cleanup
        deleted_count = asyncio.run(_cleanup_reports_async(max_age_days))

        duration = (datetime.utcnow() - start_time).total_seconds()

        logger.info(
            "Reports cleanup completed",
            extra={
                "deleted_count": deleted_count,
                "duration_seconds": duration,
            },
        )

        return {
            "status": "completed",
            "deleted_count": deleted_count,
            "duration_seconds": duration,
        }

    except Exception as exc:
        logger.exception("Reports cleanup failed")
        raise self.retry(exc=exc, countdown=3600)


async def _cleanup_reports_async(max_age_days: int) -> int:
    """Async helper for report cleanup."""
    async with AsyncSessionLocal() as db:
        try:
            # Cleanup files
            deleted_count = await report_service.cleanup_old_reports(max_age_days)

            # Cleanup database entries
            cutoff_date = datetime.now() - timedelta(days=max_age_days)
            db_deleted = await report_repository.delete_old_reports(db, cutoff_date)

            await db.commit()

            return deleted_count + db_deleted

        except Exception as e:
            await db.rollback()
            raise


@celery_app.task(
    bind=True,
    time_limit=600,  # 10 minutes
)
def cleanup_expired_sessions(self):
    """Clean up expired user sessions from cache."""
    correlation_id = set_correlation_id()

    logger.info(
        "Starting expired sessions cleanup", extra={"correlation_id": correlation_id}
    )

    try:
        # Initialize cache manager
        asyncio.run(cache_manager.initialize())

        # Delete session pattern
        deleted = asyncio.run(cache_manager.delete_pattern("session:*"))

        logger.info(
            "Expired sessions cleanup completed",
            extra={"deleted_sessions": deleted},
        )

        return {"status": "completed", "deleted_sessions": deleted}

    except Exception as exc:
        logger.exception("Sessions cleanup failed")
        raise self.retry(exc=exc, countdown=1800)


@celery_app.task(
    bind=True,
    time_limit=300,  # 5 minutes
)
def cleanup_stale_cache(self, pattern: str = "*"):
    """Clean up stale cache entries.

    Args:
        pattern: Cache key pattern to clean up
    """
    correlation_id = set_correlation_id()

    logger.info(
        "Starting stale cache cleanup",
        extra={"pattern": pattern, "correlation_id": correlation_id},
    )

    try:
        asyncio.run(cache_manager.initialize())

        # Get cache stats before cleanup
        stats_before = asyncio.run(cache_manager.get_stats())

        # Clean up expired keys (Redis does this automatically, but we can force it)
        # This is mostly for checking cache health
        stats_after = asyncio.run(cache_manager.get_stats())

        logger.info(
            "Cache cleanup completed",
            extra={
                "stats_before": stats_before,
                "stats_after": stats_after,
            },
        )

        return {
            "status": "completed",
            "stats": stats_after,
        }

    except Exception as exc:
        logger.exception("Cache cleanup failed")
        raise self.retry(exc=exc, countdown=3600)


@celery_app.task(
    bind=True,
    time_limit=60,
)
def health_check_task(self):
    """Periodic health check task.

    This task runs frequently to verify system health.
    """
    correlation_id = set_correlation_id()

    health_status = {
        "timestamp": datetime.utcnow().isoformat(),
        "status": "healthy",
        "checks": {},
    }

    # Check database connectivity
    try:
        asyncio.run(_check_database())
        health_status["checks"]["database"] = "healthy"
    except Exception as e:
        health_status["checks"]["database"] = f"unhealthy: {str(e)}"
        health_status["status"] = "degraded"
        logger.error(f"Database health check failed: {e}")

    # Check cache connectivity
    try:
        asyncio.run(cache_manager.initialize())
        stats = asyncio.run(cache_manager.get_stats())
        health_status["checks"]["cache"] = "healthy"
        health_status["checks"]["cache_stats"] = stats
    except Exception as e:
        health_status["checks"]["cache"] = f"unhealthy: {str(e)}"
        health_status["status"] = "degraded"
        logger.error(f"Cache health check failed: {e}")

    # Log health status
    if health_status["status"] == "healthy":
        logger.debug("Health check passed", extra=health_status)
    else:
        logger.warning("Health check detected issues", extra=health_status)

    return health_status


async def _check_database():
    """Check database connectivity."""
    async with AsyncSessionLocal() as db:
        from sqlalchemy import text

        result = await db.execute(text("SELECT 1"))
        result.scalar()