feat: add agent zones
This commit is contained in:
lussierj committed
1 parent
d8cd7ad058
commit
82309db79a
29 files changed
+2242
-255
No files matched your search
@@ -0,0 +1,300 @@
|
||||
"""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 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()
|
||||
Reference in new issue
Block a user