66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.edge_node import EdgeNode
|
|
from app.models.edge_task import EdgeTask
|
|
|
|
|
|
class EdgeNodeRepository:
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
|
|
def add_or_update(self, node: EdgeNode) -> EdgeNode:
|
|
self.db.add(node)
|
|
self.db.commit()
|
|
self.db.refresh(node)
|
|
return node
|
|
|
|
def get_by_edge_id(self, edge_id: str) -> EdgeNode | None:
|
|
statement = select(EdgeNode).where(EdgeNode.edge_id == edge_id)
|
|
return self.db.execute(statement).scalar_one_or_none()
|
|
|
|
|
|
class EdgeTaskRepository:
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
|
|
def add(self, edge_task: EdgeTask) -> EdgeTask:
|
|
self.db.add(edge_task)
|
|
self.db.commit()
|
|
self.db.refresh(edge_task)
|
|
return edge_task
|
|
|
|
def update(self, edge_task: EdgeTask) -> EdgeTask:
|
|
self.db.add(edge_task)
|
|
self.db.commit()
|
|
self.db.refresh(edge_task)
|
|
return edge_task
|
|
|
|
def get_by_step_id(self, step_id: str) -> EdgeTask | None:
|
|
statement = select(EdgeTask).where(EdgeTask.step_id == step_id)
|
|
return self.db.execute(statement).scalar_one_or_none()
|
|
|
|
def list_by_task_id(self, task_id: str) -> list[EdgeTask]:
|
|
statement = select(EdgeTask).where(EdgeTask.task_id == task_id).order_by(EdgeTask.created_at.desc())
|
|
return list(self.db.execute(statement).scalars())
|
|
|
|
def list_pending_by_edge_id(self, edge_id: str) -> list[EdgeTask]:
|
|
statement = (
|
|
select(EdgeTask)
|
|
.where(EdgeTask.edge_id == edge_id)
|
|
.where(EdgeTask.step_status == "PENDING")
|
|
.order_by(EdgeTask.created_at.asc())
|
|
)
|
|
return list(self.db.execute(statement).scalars())
|
|
|
|
def list_active_by_task_id(self, task_id: str) -> list[EdgeTask]:
|
|
statement = (
|
|
select(EdgeTask)
|
|
.where(EdgeTask.task_id == task_id)
|
|
.where(EdgeTask.step_status.in_(["PENDING", "RUNNING"]))
|
|
.order_by(EdgeTask.created_at.asc())
|
|
)
|
|
return list(self.db.execute(statement).scalars())
|