homelab_automation/app/crud/docker_alert.py

"""CRUD operations for Docker alerts."""
from datetime import datetime
from typing import List, Optional

from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.docker_alert import DockerAlert


class DockerAlertRepository:
    """Repository for Docker alert CRUD operations."""

    def __init__(self, session: AsyncSession):
        self.session = session

    async def get(self, alert_id: int) -> Optional[DockerAlert]:
        """Get an alert by its ID."""
        result = await self.session.execute(
            select(DockerAlert).where(DockerAlert.id == alert_id)
        )
        return result.scalar_one_or_none()

    async def get_open_alert(self, host_id: str, container_name: str) -> Optional[DockerAlert]:
        """Get an open alert for a specific container."""
        result = await self.session.execute(
            select(DockerAlert).where(
                DockerAlert.host_id == host_id,
                DockerAlert.container_name == container_name,
                DockerAlert.state == "open"
            )
        )
        return result.scalar_one_or_none()

    async def list_alerts(
        self,
        host_id: Optional[str] = None,
        state: Optional[str] = None,
        severity: Optional[str] = None,
        limit: int = 100,
        offset: int = 0
    ) -> List[DockerAlert]:
        """List alerts with optional filters."""
        query = select(DockerAlert)
        
        if host_id:
            query = query.where(DockerAlert.host_id == host_id)
        if state:
            query = query.where(DockerAlert.state == state)
        if severity:
            query = query.where(DockerAlert.severity == severity)
        
        query = query.order_by(DockerAlert.opened_at.desc()).limit(limit).offset(offset)
        result = await self.session.execute(query)
        return list(result.scalars().all())

    async def count_alerts(
        self,
        host_id: Optional[str] = None,
        state: Optional[str] = None
    ) -> dict:
        """Count alerts by state."""
        base_query = select(DockerAlert)
        if host_id:
            base_query = base_query.where(DockerAlert.host_id == host_id)
        
        # Count open
        open_query = select(func.count(DockerAlert.id)).where(DockerAlert.state == "open")
        if host_id:
            open_query = open_query.where(DockerAlert.host_id == host_id)
        
        # Count acknowledged
        ack_query = select(func.count(DockerAlert.id)).where(DockerAlert.state == "acknowledged")
        if host_id:
            ack_query = ack_query.where(DockerAlert.host_id == host_id)
        
        # Count total
        total_query = select(func.count(DockerAlert.id))
        if host_id:
            total_query = total_query.where(DockerAlert.host_id == host_id)
        
        open_result = await self.session.execute(open_query)
        ack_result = await self.session.execute(ack_query)
        total_result = await self.session.execute(total_query)
        
        return {
            "total": total_result.scalar() or 0,
            "open": open_result.scalar() or 0,
            "acknowledged": ack_result.scalar() or 0
        }

    async def count_open_by_host(self, host_id: str) -> int:
        """Count open alerts for a host."""
        result = await self.session.execute(
            select(func.count(DockerAlert.id)).where(
                DockerAlert.host_id == host_id,
                DockerAlert.state == "open"
            )
        )
        return result.scalar() or 0

    async def create(
        self,
        host_id: str,
        container_name: str,
        severity: str = "warning",
        message: Optional[str] = None
    ) -> DockerAlert:
        """Create a new alert."""
        alert = DockerAlert(
            host_id=host_id,
            container_name=container_name,
            severity=severity,
            state="open",
            message=message,
            opened_at=datetime.utcnow()
        )
        self.session.add(alert)
        return alert

    async def acknowledge(
        self,
        alert_id: int,
        acknowledged_by: str
    ) -> Optional[DockerAlert]:
        """Acknowledge an alert."""
        alert = await self.get(alert_id)
        if alert and alert.state == "open":
            alert.state = "acknowledged"
            alert.acknowledged_at = datetime.utcnow()
            alert.acknowledged_by = acknowledged_by
            return alert
        return None

    async def close(self, alert_id: int) -> Optional[DockerAlert]:
        """Close an alert."""
        alert = await self.get(alert_id)
        if alert and alert.state in ("open", "acknowledged"):
            alert.state = "closed"
            alert.closed_at = datetime.utcnow()
            return alert
        return None

    async def close_for_container(self, host_id: str, container_name: str) -> int:
        """Close all open alerts for a container."""
        result = await self.session.execute(
            select(DockerAlert).where(
                DockerAlert.host_id == host_id,
                DockerAlert.container_name == container_name,
                DockerAlert.state.in_(["open", "acknowledged"])
            )
        )
        alerts = result.scalars().all()
        count = 0
        for alert in alerts:
            alert.state = "closed"
            alert.closed_at = datetime.utcnow()
            count += 1
        return count

    async def update_last_notified(self, alert_id: int) -> None:
        """Update the last notified timestamp."""
        alert = await self.get(alert_id)
        if alert:
            alert.last_notified_at = datetime.utcnow()