Admiral/admiral-worker/app/repos/sql/WorkerStatusSqlRepo.py
2025-06-24 14:22:50 +02:00

70 lines
2.0 KiB
Python

import uuid
from dataclasses import dataclass
from typing import Optional
from typing_extensions import Self
from sqlalchemy.engine import Engine
from sqlmodel import SQLModel, Field, Session, select
from typing_extensions import override
from app.repos.sql import dbRetry
from core.domain.worker.Worker import Worker
from core.domain.worker.WorkerStatus import WorkerStatus
from core.repos.WorkerStatusRepo import WorkerStatusRepo
from core.types.Id import Id
@dataclass
class WorkerStatusSqlRepo(WorkerStatusRepo):
db: Engine
class Table(SQLModel, table=True):
__tablename__ = "worker_status"
id: str = Field(primary_key=True)
worker_id: str = Field(foreign_key="worker.id")
ram_available: float
cpu_utilization: float
created_at: int
def toDomain(self) -> WorkerStatus:
return WorkerStatus(
ownerId=Id(value=uuid.UUID(self.worker_id)),
ramAvailable=self.ram_available,
cpuUtilization=self.cpu_utilization,
createdAt=self.created_at,
id=Id(value=uuid.UUID(self.id))
)
@classmethod
def toRow(cls, obj: WorkerStatus) -> Self:
return cls(
worker_id=obj.ownerId.value,
ram_available=obj.ramAvailable,
cpu_utilization=obj.cpuUtilization,
created_at=obj.createdAt,
id=obj.id.value
)
@override
def getAll(self) -> list[WorkerStatus]:
with Session(self.db) as conn:
query = select(self.Table)
return [row.toDomain() for row in conn.exec(query).all()]
@override
def get(self, id: Id[WorkerStatus]) -> Optional[WorkerStatus]:
with Session(self.db) as conn:
query = select(self.Table).filter_by(id=id.value)
row = conn.exec(query).one_or_none()
return row.toDomain() if row is not None else None
@dbRetry
@override
def post(self, workerId: Id[Worker], ramAvailable: float, cpuUtilization: float) -> WorkerStatus:
workerStatus = WorkerStatus(ownerId=workerId, ramAvailable=ramAvailable, cpuUtilization=cpuUtilization)
with Session(self.db) as conn:
conn.add(self.Table.toRow(workerStatus))
conn.commit()
return workerStatus