345 lines
13 KiB
Python
345 lines
13 KiB
Python
"""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
|
|
""",
|
|
(
|
|
agent_id,
|
|
site_id,
|
|
agent_zone_id,
|
|
site_name,
|
|
hostname,
|
|
version,
|
|
ip_address,
|
|
Jsonb(payload),
|
|
),
|
|
)
|
|
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 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:
|
|
cur.execute(
|
|
"""
|
|
UPDATE miners
|
|
SET site_id = %s, agent_zone_id = %s, updated_at = now()
|
|
WHERE id = %s AND agent_zone_id IS NULL
|
|
""",
|
|
(zone["site_id"], agent_zone_id, miner_id),
|
|
)
|
|
outcomes[miner_id] = "assigned"
|
|
connection.commit()
|
|
return outcomes
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|