minerops/backend/app/agent_zones/repository.py

"""PostgreSQL persistence for Agent Zones and ownership assignments."""

from collections.abc import Callable, Iterable
from typing import Any
from uuid import UUID

import psycopg
from psycopg.errors import ForeignKeyViolation, UniqueViolation
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from app.agent_zones.exceptions import (
    AgentZoneConflict,
    AgentZoneNotFound,
)
from app.config import settings


ConnectionFactory = Callable[[], psycopg.Connection]


def default_connection_factory() -> psycopg.Connection:
    return psycopg.connect(
        host=settings.postgres_host,
        port=settings.postgres_port,
        dbname=settings.postgres_db,
        user=settings.postgres_user,
        password=settings.postgres_password,
        connect_timeout=5,
    )


class AgentZoneRepository:
    def __init__(self, connection_factory: ConnectionFactory | None = None):
        self._connection_factory = connection_factory or default_connection_factory

    def get_agent(self, agent_id: str) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT agent_id, site_id, agent_zone_id, site_name, hostname,
                           version, ip_address, last_checkin_at, last_payload
                    FROM minerops_site_agents
                    WHERE agent_id = %s
                    """,
                    (agent_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def get_zone(self, zone_id: UUID) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT id, site_id, name, description, is_enabled,
                           created_at, updated_at
                    FROM agent_zones
                    WHERE id = %s
                    """,
                    (zone_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def list_usable_zones(self, site_id: UUID) -> list[dict]:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT id, site_id, name, description, is_enabled
                    FROM agent_zones
                    WHERE site_id = %s AND is_enabled
                    ORDER BY lower(name), id
                    """,
                    (site_id,),
                )
                return [dict(row) for row in cur.fetchall()]
        finally:
            connection.close()

    def list_zones(self, site_id: UUID | None = None) -> list[dict]:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT zone.id, zone.site_id, zone.name, zone.description,
                           zone.is_enabled, agent.agent_id AS assigned_agent_id,
                           range.id AS range_id, range.network::text AS network
                    FROM agent_zones AS zone
                    LEFT JOIN minerops_site_agents AS agent
                      ON agent.agent_zone_id = zone.id
                    LEFT JOIN agent_zone_network_ranges AS range
                      ON range.agent_zone_id = zone.id
                    WHERE (%s::uuid IS NULL OR zone.site_id = %s::uuid)
                    ORDER BY lower(zone.name), zone.id, range.network, range.id
                    """,
                    (site_id, site_id),
                )
                rows = cur.fetchall()
        finally:
            connection.close()

        zones: dict[UUID, dict] = {}
        for row in rows:
            zone = zones.setdefault(
                row["id"],
                {
                    "id": row["id"],
                    "site_id": row["site_id"],
                    "name": row["name"],
                    "description": row["description"],
                    "is_enabled": row["is_enabled"],
                    "assigned_agent_id": row["assigned_agent_id"],
                    "networks": [],
                },
            )
            if row["range_id"] is not None:
                zone["networks"].append(
                    {"id": row["range_id"], "network": row["network"]}
                )
        return list(zones.values())

    def create_zone(
        self,
        *,
        site_id: UUID,
        name: str,
        description: str | None,
        networks: list[str],
    ) -> UUID:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    INSERT INTO agent_zones (site_id, name, description)
                    VALUES (%s, %s, %s)
                    RETURNING id
                    """,
                    (site_id, name, description),
                )
                zone_id = cur.fetchone()["id"]
                for network in networks:
                    cur.execute(
                        """
                        INSERT INTO agent_zone_network_ranges (agent_zone_id, network)
                        VALUES (%s, %s::cidr)
                        """,
                        (zone_id, network),
                    )
            connection.commit()
            return zone_id
        except ForeignKeyViolation as exc:
            connection.rollback()
            raise AgentZoneNotFound("Site not found") from exc
        except UniqueViolation as exc:
            connection.rollback()
            raise AgentZoneConflict(
                "An enabled Agent Zone with this name already exists at the Site"
            ) from exc
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def save_checkin(
        self,
        *,
        agent_id: str,
        site_id: UUID | None,
        agent_zone_id: UUID | None,
        site_name: str,
        hostname: str | None,
        version: str | None,
        ip_address: str | None,
        payload: dict[str, Any],
    ) -> None:
        connection = self._connection_factory()
        try:
            with connection.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO minerops_site_agents (
                        agent_id, site_id, agent_zone_id, site_name, hostname,
                        version, ip_address, last_checkin_at, last_payload
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, now(), %s)
                    ON CONFLICT (agent_id) DO UPDATE SET
                        site_id = EXCLUDED.site_id,
                        agent_zone_id = EXCLUDED.agent_zone_id,
                        site_name = EXCLUDED.site_name,
                        hostname = EXCLUDED.hostname,
                        version = EXCLUDED.version,
                        ip_address = EXCLUDED.ip_address,
                        last_checkin_at = now(),
                        last_payload = EXCLUDED.last_payload
                    WHERE minerops_site_agents.site_id IS NOT DISTINCT FROM EXCLUDED.site_id
                      AND minerops_site_agents.agent_zone_id IS NOT DISTINCT FROM EXCLUDED.agent_zone_id
                    """,
                    (
                        agent_id,
                        site_id,
                        agent_zone_id,
                        site_name,
                        hostname,
                        version,
                        ip_address,
                        Jsonb(payload),
                    ),
                )
                if cur.rowcount != 1:
                    raise AgentZoneConflict("Agent assignment changed concurrently")
            connection.commit()
        except UniqueViolation as exc:
            connection.rollback()
            raise AgentZoneConflict(
                "The Agent Zone already has an assigned Agent"
            ) from exc
        except (ForeignKeyViolation, psycopg.errors.CheckViolation) as exc:
            connection.rollback()
            raise AgentZoneConflict("Agent Zone and Site assignment conflict") from exc
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()

    def get_agent_assignment(self, agent_id: str) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT agent.agent_id, agent.site_id, agent.agent_zone_id,
                           zone.is_enabled AS zone_enabled
                    FROM minerops_site_agents AS agent
                    LEFT JOIN agent_zones AS zone ON zone.id = agent.agent_zone_id
                    WHERE agent.agent_id = %s
                    """,
                    (agent_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def get_agent_configuration(self, agent_id: str) -> dict | None:
        connection = self._connection_factory()
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    """
                    SELECT
                        agent.agent_id,
                        agent.site_id,
                        agent.agent_zone_id,
                        coalesce(site.name, agent.site_name) AS site_name,
                        zone.name AS agent_zone_name,
                        zone.is_enabled AS agent_zone_enabled,
                        coalesce(
                            array_agg(
                                network.network::text
                                ORDER BY network.network::text, network.id
                            ) FILTER (WHERE network.id IS NOT NULL),
                            ARRAY[]::text[]
                        ) AS networks
                    FROM minerops_site_agents AS agent
                    LEFT JOIN sites AS site
                      ON site.id = agent.site_id
                    LEFT JOIN agent_zones AS zone
                      ON zone.id = agent.agent_zone_id
                    LEFT JOIN agent_zone_network_ranges AS network
                      ON network.agent_zone_id = zone.id
                    WHERE agent.agent_id = %s
                    GROUP BY
                        agent.agent_id,
                        agent.site_id,
                        agent.agent_zone_id,
                        site.name,
                        agent.site_name,
                        zone.name,
                        zone.is_enabled
                    """,
                    (agent_id,),
                )
                row = cur.fetchone()
                return dict(row) if row else None
        finally:
            connection.close()

    def find_miner_in_zone(self, *, agent_id: str, agent_zone_id: UUID,
                           mac_address: str | None, ip_address: str | None) -> dict | None:
        """Retained helper uses the same atomic authority/collision boundary as HTTP."""
        from app.discovery.repository import DiscoveryRepository, reject
        repository = DiscoveryRepository(self._connection_factory)
        with repository.transaction() as cur:
            context = repository.context(cur, agent_id)
            if context['agent_zone_id'] != agent_zone_id:
                reject('discovery_stale_authority')
            if not mac_address and not ip_address:
                return None
            row = repository.match(cur, context, {'mac_address': mac_address, 'ip_address': ip_address})
            return {'id': row['id'], 'matched_by': 'mac' if mac_address else 'ip'} if row else None

    def assign_miners_to_zone(
        self, miner_ids: Iterable[UUID], agent_zone_id: UUID
    ) -> dict[UUID, str]:
        unique_ids = list(dict.fromkeys(miner_ids))
        if not unique_ids:
            return {}
        connection = self._connection_factory()
        outcomes: dict[UUID, str] = {}
        try:
            with connection.cursor(row_factory=dict_row) as cur:
                cur.execute(
                    "SELECT id, site_id FROM agent_zones WHERE id = %s",
                    (agent_zone_id,),
                )
                zone = cur.fetchone()
                if not zone:
                    raise AgentZoneNotFound("Agent Zone not found")
                for miner_id in unique_ids:
                    cur.execute(
                        "SELECT id, site_id, agent_zone_id FROM miners WHERE id = %s FOR UPDATE",
                        (miner_id,),
                    )
                    miner = cur.fetchone()
                    if not miner:
                        outcomes[miner_id] = "miner_not_found"
                    elif miner["agent_zone_id"] == agent_zone_id:
                        outcomes[miner_id] = "preserved"
                    elif miner["agent_zone_id"] is not None:
                        outcomes[miner_id] = "agent_zone_conflict"
                    elif miner["site_id"] not in {None, zone["site_id"]}:
                        outcomes[miner_id] = "site_conflict"
                    else:
                        outcomes[miner_id] = "miner_zone_unassigned"
            connection.commit()
            return outcomes
        except Exception:
            connection.rollback()
            raise
        finally:
            connection.close()