226 lines
6.9 KiB
Python
226 lines
6.9 KiB
Python
"""PostgreSQL persistence for repair workflows."""
|
|
|
|
from collections.abc import Callable, Iterator, Mapping
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
import psycopg
|
|
from psycopg.errors import UniqueViolation
|
|
from psycopg.rows import dict_row
|
|
from psycopg.types.json import Jsonb
|
|
|
|
from app.config import settings
|
|
from app.repairs.exceptions import (
|
|
ActiveRepairConflict,
|
|
RepairPersistenceError,
|
|
)
|
|
|
|
|
|
ConnectionFactory = Callable[[], psycopg.Connection]
|
|
ACTIVE_REPAIR_STATUSES = (
|
|
"OPEN",
|
|
"IN_PROGRESS",
|
|
"WAITING_FOR_PARTS",
|
|
"READY_FOR_TEST",
|
|
)
|
|
ACTIVE_REPAIR_CONSTRAINT = "uq_repairs_one_active_per_miner"
|
|
|
|
|
|
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 RepairRepository:
|
|
"""Persistence boundary for the atomic Begin Repair workflow."""
|
|
|
|
def __init__(
|
|
self,
|
|
connection_factory: ConnectionFactory | None = None,
|
|
*,
|
|
connection: psycopg.Connection | None = None,
|
|
):
|
|
self._connection_factory = connection_factory or default_connection_factory
|
|
self._connection = connection
|
|
|
|
@contextmanager
|
|
def transaction(self) -> Iterator["RepairRepository"]:
|
|
if self._connection is not None:
|
|
yield self
|
|
return
|
|
|
|
connection = self._connection_factory()
|
|
transaction_repository = RepairRepository(
|
|
self._connection_factory,
|
|
connection=connection,
|
|
)
|
|
try:
|
|
yield transaction_repository
|
|
connection.commit()
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
@contextmanager
|
|
def _connection_scope(
|
|
self,
|
|
*,
|
|
write: bool = False,
|
|
) -> Iterator[psycopg.Connection]:
|
|
if self._connection is not None:
|
|
yield self._connection
|
|
return
|
|
|
|
connection = self._connection_factory()
|
|
try:
|
|
yield connection
|
|
if write:
|
|
connection.commit()
|
|
except Exception:
|
|
if write:
|
|
connection.rollback()
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
def lock_miner_for_update(self, miner_id: UUID) -> dict | None:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT id, lifecycle_status
|
|
FROM miners
|
|
WHERE id = %s
|
|
FOR UPDATE
|
|
""",
|
|
(miner_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_active_repair(self, miner_id: UUID) -> dict | None:
|
|
with self._connection_scope() as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT
|
|
id,
|
|
miner_id,
|
|
status,
|
|
reason,
|
|
opened_at,
|
|
opened_by,
|
|
completed_at,
|
|
completed_by
|
|
FROM repairs
|
|
WHERE miner_id = %s
|
|
AND status = ANY(%s)
|
|
ORDER BY opened_at, id
|
|
LIMIT 1
|
|
""",
|
|
(miner_id, list(ACTIVE_REPAIR_STATUSES)),
|
|
)
|
|
row = cur.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def create_repair(
|
|
self,
|
|
miner_id: UUID,
|
|
reason: str,
|
|
opened_by: UUID,
|
|
) -> dict:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
try:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO repairs (miner_id, reason, opened_by)
|
|
VALUES (%s, %s, %s)
|
|
RETURNING
|
|
id,
|
|
miner_id,
|
|
status,
|
|
reason,
|
|
opened_at,
|
|
opened_by,
|
|
completed_at,
|
|
completed_by
|
|
""",
|
|
(miner_id, reason, opened_by),
|
|
)
|
|
except UniqueViolation as exc:
|
|
if exc.diag.constraint_name == ACTIVE_REPAIR_CONSTRAINT:
|
|
raise ActiveRepairConflict(miner_id) from exc
|
|
raise
|
|
return dict(cur.fetchone())
|
|
|
|
def create_event(
|
|
self,
|
|
repair_id: UUID,
|
|
*,
|
|
sequence: int,
|
|
event_type: str,
|
|
created_by: UUID,
|
|
payload: Mapping[str, Any],
|
|
) -> dict:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO repair_events (
|
|
repair_id,
|
|
sequence,
|
|
event_type,
|
|
created_by,
|
|
payload
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s)
|
|
RETURNING
|
|
id,
|
|
repair_id,
|
|
sequence,
|
|
event_type,
|
|
created_at,
|
|
created_by,
|
|
payload
|
|
""",
|
|
(
|
|
repair_id,
|
|
sequence,
|
|
event_type,
|
|
created_by,
|
|
Jsonb(dict(payload)),
|
|
),
|
|
)
|
|
return dict(cur.fetchone())
|
|
|
|
def set_miner_lifecycle(
|
|
self,
|
|
miner_id: UUID,
|
|
lifecycle_status: str,
|
|
) -> None:
|
|
with self._connection_scope(write=True) as connection:
|
|
with connection.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE miners
|
|
SET lifecycle_status = %s,
|
|
updated_at = now()
|
|
WHERE id = %s
|
|
""",
|
|
(lifecycle_status, miner_id),
|
|
)
|
|
if cur.rowcount != 1:
|
|
raise RepairPersistenceError(
|
|
"Miner lifecycle transition did not update one record"
|
|
)
|