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