diff --git a/decnet/web/db/sqlmodel_repo/__init__.py b/decnet/web/db/sqlmodel_repo/__init__.py index 5563fbb7..cf1580b8 100644 --- a/decnet/web/db/sqlmodel_repo/__init__.py +++ b/decnet/web/db/sqlmodel_repo/__init__.py @@ -40,7 +40,6 @@ from decnet.web.db.models import ( Campaign, SessionProfile, SmtpTarget, - SwarmHost, DeckyShard, FleetDecky, LOCAL_HOST_SENTINEL, @@ -67,10 +66,12 @@ from decnet.web.db.sqlmodel_repo._helpers import ( # noqa: F401 (re-exported f _cleanup_tasks, ) from decnet.web.db.sqlmodel_repo.attacker_intel import AttackerIntelMixin +from decnet.web.db.sqlmodel_repo.swarm import SwarmMixin class SQLModelRepository( AttackerIntelMixin, + SwarmMixin, BaseRepository, ): """Concrete SQLModel/SQLAlchemy-async repository. @@ -1775,65 +1776,6 @@ class SQLModelRepository( # ------------------------------------------------------------- swarm - async def add_swarm_host(self, data: dict[str, Any]) -> None: - async with self._session() as session: - session.add(SwarmHost(**data)) - await session.commit() - - async def get_swarm_host_by_name(self, name: str) -> Optional[dict[str, Any]]: - async with self._session() as session: - result = await session.execute(select(SwarmHost).where(SwarmHost.name == name)) - row = result.scalar_one_or_none() - return row.model_dump(mode="json") if row else None - - async def get_swarm_host_by_uuid(self, uuid: str) -> Optional[dict[str, Any]]: - async with self._session() as session: - result = await session.execute(select(SwarmHost).where(SwarmHost.uuid == uuid)) - row = result.scalar_one_or_none() - return row.model_dump(mode="json") if row else None - - async def get_swarm_host_by_fingerprint(self, fingerprint: str) -> Optional[dict[str, Any]]: - async with self._session() as session: - result = await session.execute( - select(SwarmHost).where(SwarmHost.client_cert_fingerprint == fingerprint) - ) - row = result.scalar_one_or_none() - return row.model_dump(mode="json") if row else None - - async def list_swarm_hosts(self, status: Optional[str] = None) -> list[dict[str, Any]]: - statement = select(SwarmHost).order_by(asc(SwarmHost.name)) - if status: - statement = statement.where(SwarmHost.status == status) - async with self._session() as session: - result = await session.execute(statement) - return [r.model_dump(mode="json") for r in result.scalars().all()] - - async def update_swarm_host(self, uuid: str, fields: dict[str, Any]) -> None: - if not fields: - return - async with self._session() as session: - await session.execute( - update(SwarmHost).where(SwarmHost.uuid == uuid).values(**fields) - ) - await session.commit() - - async def delete_swarm_host(self, uuid: str) -> bool: - async with self._session() as session: - # Clean up child shards first (no ON DELETE CASCADE portable across dialects). - await session.execute( - text("DELETE FROM decky_shards WHERE host_uuid = :u"), {"u": uuid} - ) - result = await session.execute( - select(SwarmHost).where(SwarmHost.uuid == uuid) - ) - host = result.scalar_one_or_none() - if not host: - await session.commit() - return False - await session.delete(host) - await session.commit() - return True - async def upsert_decky_shard(self, data: dict[str, Any]) -> None: payload = {**data, "updated_at": datetime.now(timezone.utc)} if isinstance(payload.get("services"), list): diff --git a/decnet/web/db/sqlmodel_repo/swarm.py b/decnet/web/db/sqlmodel_repo/swarm.py new file mode 100644 index 00000000..9fea2d9d --- /dev/null +++ b/decnet/web/db/sqlmodel_repo/swarm.py @@ -0,0 +1,71 @@ +"""Swarm host CRUD.""" +from __future__ import annotations + +from typing import Any, Optional + +from sqlalchemy import asc, select, text, update + +from decnet.web.db.models import SwarmHost + + +class SwarmMixin: + """Mixin: composed onto ``SQLModelRepository``. Expects ``self._session()``.""" + + async def add_swarm_host(self, data: dict[str, Any]) -> None: + async with self._session() as session: + session.add(SwarmHost(**data)) + await session.commit() + + async def get_swarm_host_by_name(self, name: str) -> Optional[dict[str, Any]]: + async with self._session() as session: + result = await session.execute(select(SwarmHost).where(SwarmHost.name == name)) + row = result.scalar_one_or_none() + return row.model_dump(mode="json") if row else None + + async def get_swarm_host_by_uuid(self, uuid: str) -> Optional[dict[str, Any]]: + async with self._session() as session: + result = await session.execute(select(SwarmHost).where(SwarmHost.uuid == uuid)) + row = result.scalar_one_or_none() + return row.model_dump(mode="json") if row else None + + async def get_swarm_host_by_fingerprint(self, fingerprint: str) -> Optional[dict[str, Any]]: + async with self._session() as session: + result = await session.execute( + select(SwarmHost).where(SwarmHost.client_cert_fingerprint == fingerprint) + ) + row = result.scalar_one_or_none() + return row.model_dump(mode="json") if row else None + + async def list_swarm_hosts(self, status: Optional[str] = None) -> list[dict[str, Any]]: + statement = select(SwarmHost).order_by(asc(SwarmHost.name)) + if status: + statement = statement.where(SwarmHost.status == status) + async with self._session() as session: + result = await session.execute(statement) + return [r.model_dump(mode="json") for r in result.scalars().all()] + + async def update_swarm_host(self, uuid: str, fields: dict[str, Any]) -> None: + if not fields: + return + async with self._session() as session: + await session.execute( + update(SwarmHost).where(SwarmHost.uuid == uuid).values(**fields) + ) + await session.commit() + + async def delete_swarm_host(self, uuid: str) -> bool: + async with self._session() as session: + # Clean up child shards first (no ON DELETE CASCADE portable across dialects). + await session.execute( + text("DELETE FROM decky_shards WHERE host_uuid = :u"), {"u": uuid} + ) + result = await session.execute( + select(SwarmHost).where(SwarmHost.uuid == uuid) + ) + host = result.scalar_one_or_none() + if not host: + await session.commit() + return False + await session.delete(host) + await session.commit() + return True