131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
"""PostgreSQL persistence for Agent command transport."""
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import psycopg
|
|
from psycopg.rows import dict_row
|
|
from psycopg.types.json import Jsonb
|
|
|
|
from app.agent_commands.exceptions import AgentCommandConflict, AgentCommandNotFound
|
|
from app.config import settings
|
|
|
|
|
|
ConnectionFactory = Callable[[], psycopg.Connection]
|
|
|
|
CLAIM_COMMANDS_SQL = """
|
|
WITH candidates AS (
|
|
SELECT id
|
|
FROM minerops_agent_command_queue
|
|
WHERE agent_id = %s
|
|
AND command_status = 'pending'
|
|
ORDER BY created_at, id
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT %s
|
|
)
|
|
UPDATE minerops_agent_command_queue AS queue
|
|
SET command_status = 'claimed', claimed_at = now()
|
|
FROM candidates
|
|
WHERE queue.id = candidates.id
|
|
AND queue.command_status = 'pending'
|
|
RETURNING queue.id AS command_id, queue.command_type,
|
|
queue.request_payload, queue.created_at
|
|
"""
|
|
|
|
|
|
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 AgentCommandRepository:
|
|
"""Own queue locking and immutable terminal transitions."""
|
|
|
|
def __init__(self, connection_factory: ConnectionFactory | None = None):
|
|
self._connection_factory = connection_factory or default_connection_factory
|
|
|
|
def claim_pending(self, agent_id: str, limit: int = 5) -> list[dict]:
|
|
connection = self._connection_factory()
|
|
try:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(CLAIM_COMMANDS_SQL, (agent_id, limit))
|
|
rows = [dict(row) for row in cur.fetchall()]
|
|
connection.commit()
|
|
return rows
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
def get_command(self, command_id: str) -> dict | None:
|
|
connection = self._connection_factory()
|
|
try:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT id, agent_id, command_type, command_status, request_payload
|
|
FROM minerops_agent_command_queue WHERE id = %s
|
|
""",
|
|
(command_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
return dict(row) if row else None
|
|
finally:
|
|
connection.close()
|
|
|
|
def finish_command(
|
|
self,
|
|
command_id: str,
|
|
agent_id: str,
|
|
status: str,
|
|
result_payload: dict[str, Any],
|
|
error_message: str | None,
|
|
) -> None:
|
|
connection = self._connection_factory()
|
|
try:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE minerops_agent_command_queue
|
|
SET command_status = %s, result_payload = %s,
|
|
error_message = %s, completed_at = now()
|
|
WHERE id = %s AND agent_id = %s AND command_status = 'claimed'
|
|
RETURNING id
|
|
""",
|
|
(status, Jsonb(result_payload), error_message, command_id, agent_id),
|
|
)
|
|
if cur.fetchone():
|
|
connection.commit()
|
|
return
|
|
|
|
cur.execute(
|
|
"""
|
|
SELECT agent_id, command_status
|
|
FROM minerops_agent_command_queue WHERE id = %s
|
|
""",
|
|
(command_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
connection.rollback()
|
|
if not row:
|
|
raise AgentCommandNotFound("Agent command not found")
|
|
if row["agent_id"] != agent_id:
|
|
raise AgentCommandConflict("Agent command belongs to another Agent")
|
|
raise AgentCommandConflict(
|
|
f"Agent command is {row['command_status']}, not claimed"
|
|
)
|
|
except (AgentCommandNotFound, AgentCommandConflict):
|
|
raise
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|