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, or_, not_
from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server

UNCATEGORIZED_CATEGORY = "uncategorized"


def _escape_like(s: str) -> str:
    return s.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")


def _search_filter(search: str):
    """Substring match on name, domain, category (case-insensitive)."""
    term = f"%{_escape_like(search.strip())}%"
    return or_(
        Server.name.ilike(term, escape="\\"),
        Server.domain.ilike(term, escape="\\"),
        Server.category.ilike(term, escape="\\"),
    )


def _category_filter(category: str):
    """Map API category filter to SQLAlchemy criterion."""
    if category == UNCATEGORIZED_CATEGORY:
        return or_(Server.category.is_(None), Server.category == "")
    return Server.category == category


def _target_path_unset_filter(unset: bool):
    """Filter servers by whether target_path is unset (empty or BT default).

    Mirrors :func:`server.utils.posix_paths.is_unset_target_path` using SQL TRIM.
    """
    tp = func.trim(Server.target_path)
    unset_expr = or_(
        Server.target_path.is_(None),
        tp == "",
        tp == "/www/wwwroot",
    )
    return unset_expr if unset else not_(unset_expr)


def _agent_not_installed_filter(not_installed: bool):
    """Filter by DB agent_version empty (mirrors agent_is_installed without heartbeat)."""
    av = func.trim(Server.agent_version)
    expr = or_(Server.agent_version.is_(None), av == "")
    return expr if not_installed else not_(expr)


# 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,
    "target_path": Server.target_path,
    "created_at": Server.created_at,
}


def _server_list_order(sort_by: Optional[str], sort_order: Optional[str]):
    """Default: newest created first."""
    if not sort_by:
        return desc(Server.created_at), desc(Server.id)
    sort_col = _SORT_COLUMNS.get(sort_by)
    if sort_col is None:
        return desc(Server.created_at), desc(Server.id)
    reverse = sort_order == "desc"
    primary = desc(sort_col) if reverse else asc(sort_col)
    secondary = desc(Server.id) if reverse else asc(Server.id)
    return primary, secondary


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_by_ids(self, ids: List[int]) -> List[Server]:
        """Fetch multiple servers in one query; order matches *ids* (skips missing)."""
        if not ids:
            return []
        result = await self.session.execute(select(Server).where(Server.id.in_(ids)))
        by_id = {s.id: s for s in result.scalars().all()}
        return [by_id[i] for i in ids if i in by_id]

    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,
        target_path_unset: Optional[bool] = None,
        agent_not_installed: 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(_category_filter(category))
        if platform_id:
            filters.append(Server.platform_id == platform_id)
        if node_id:
            filters.append(Server.node_id == node_id)
        if search:
            filters.append(_search_filter(search))
        if is_online is not None:
            filters.append(Server.is_online == is_online)
        if target_path_unset is not None:
            filters.append(_target_path_unset_filter(target_path_unset))
        if agent_not_installed is not None:
            filters.append(_agent_not_installed_filter(agent_not_installed))

        # 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)

        primary, secondary = _server_list_order(sort_by, sort_order)
        data_query = data_query.order_by(primary, secondary)
        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_filtered(
        self,
        category: Optional[str] = None,
        platform_id: Optional[int] = None,
        node_id: Optional[int] = None,
        search: Optional[str] = None,
        is_online: Optional[bool] = None,
        target_path_unset: Optional[bool] = None,
        agent_not_installed: Optional[bool] = None,
        max_rows: int = 5000,
    ) -> Tuple[List[Server], int]:
        """All servers matching filters (no sort/pagination), capped at *max_rows*."""
        filters = []
        if category:
            filters.append(_category_filter(category))
        if platform_id:
            filters.append(Server.platform_id == platform_id)
        if node_id:
            filters.append(Server.node_id == node_id)
        if search:
            filters.append(_search_filter(search))
        if is_online is not None:
            filters.append(Server.is_online == is_online)
        if target_path_unset is not None:
            filters.append(_target_path_unset_filter(target_path_unset))
        if agent_not_installed is not None:
            filters.append(_agent_not_installed_filter(agent_not_installed))

        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

        data_query = select(Server)
        for f in filters:
            data_query = data_query.where(f)
        data_query = data_query.order_by(desc(Server.created_at), desc(Server.id)).limit(max_rows)
        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 update_category_batch(self, ids: List[int], category: Optional[str]) -> int:
        """Set category for multiple servers in one UPDATE."""
        if not ids:
            return 0
        result = await self.session.execute(
            update(Server).where(Server.id.in_(ids)).values(category=category)
        )
        await self.session.commit()
        return int(result.rowcount or 0)

    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()