Nexus/server/infrastructure/database/server_repo.py

"""Nexus — Server Repository (Async SQLAlchemy implementation)
Concrete implementation of ServerRepository protocol.
"""

from datetime import datetime, timezone
from typing import Optional, List, Tuple
from sqlalchemy import select, update, func, asc, desc
from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server

# Whitelist of sortable columns (prevents SQL injection via sort_by)
_SORT_COLUMNS = {
    "id": Server.id,
    "name": Server.name,
    "domain": Server.domain,
    "status": Server.is_online,
    "heartbeat": Server.last_heartbeat,
    "agent_version": Server.agent_version,
    "category": Server.category,
}


class ServerRepositoryImpl:
    """Async SQLAlchemy implementation of Server data access"""

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

    async def get_by_id(self, id: int) -> Optional[Server]:
        result = await self.session.execute(select(Server).where(Server.id == id))
        return result.scalar_one_or_none()

    async def get_all(self) -> List[Server]:
        result = await self.session.execute(select(Server).order_by(Server.id))
        return list(result.scalars().all())

    async def get_paginated(
        self,
        category: Optional[str] = None,
        platform_id: Optional[int] = None,
        node_id: Optional[int] = None,
        offset: int = 0,
        limit: int = 50,
        search: Optional[str] = None,
        sort_by: Optional[str] = None,
        sort_order: Optional[str] = "asc",
        is_online: Optional[bool] = None,
    ) -> Tuple[List[Server], int]:
        """Get paginated server list with total count.

        Supports text search (name/domain), sorting, and online status filter.
        Returns (servers, total_count) — DB-level pagination for 2000+ servers.
        When is_online filter is active, we over-fetch (limit + 10) to compensate
        for Redis overlay post-filter corrections in the API layer.
        """
        filters = []
        if category:
            filters.append(Server.category == category)
        if platform_id:
            filters.append(Server.platform_id == platform_id)
        if node_id:
            filters.append(Server.node_id == node_id)
        if search:
            term = f"%{search}%"
            filters.append(Server.name.ilike(term) | Server.domain.ilike(term))
        if is_online is not None:
            filters.append(Server.is_online == is_online)

        # Count query
        count_query = select(func.count(Server.id))
        for f in filters:
            count_query = count_query.where(f)
        count_result = await self.session.execute(count_query)
        total = count_result.scalar_one() or 0

        # Over-fetch when is_online filter is active (Redis overlay may change status)
        fetch_limit = limit + 10 if is_online is not None else limit

        # Paginated data query with sort
        data_query = select(Server)
        for f in filters:
            data_query = data_query.where(f)

        # Sorting (whitelist-validated)
        sort_col = _SORT_COLUMNS.get(sort_by, Server.id)
        order_func = desc if sort_order == "desc" else asc
        data_query = data_query.order_by(order_func(sort_col), Server.id)
        data_query = data_query.offset(offset).limit(fetch_limit)
        result = await self.session.execute(data_query)
        servers = list(result.scalars().all())

        return servers, total

    async def get_by_category(self, category: str) -> List[Server]:
        result = await self.session.execute(
            select(Server).where(Server.category == category).order_by(Server.id)
        )
        return list(result.scalars().all())

    async def get_online(self) -> List[Server]:
        result = await self.session.execute(
            select(Server).where(Server.is_online == True).order_by(Server.id)
        )
        return list(result.scalars().all())

    async def create(self, server: Server) -> Server:
        self.session.add(server)
        await self.session.commit()
        return server

    async def update(self, server: Server) -> Server:
        await self.session.commit()
        return server

    async def delete(self, id: int) -> bool:
        server = await self.get_by_id(id)
        if server:
            await self.session.delete(server)
            await self.session.commit()
            return True
        return False

    async def update_heartbeat(self, id: int, is_online: bool, system_info: str, agent_version: str) -> None:
        """Update server online status from Agent heartbeat"""
        await self.session.execute(
            update(Server)
            .where(Server.id == id)
            .values(
                is_online=is_online,
                last_heartbeat=datetime.now(timezone.utc),
                system_info=system_info,
                agent_version=agent_version,
            )
        )
        await self.session.commit()