Make queries.get_severity_counts() return a dict

This commit is contained in:
Alexis Métaireau 2023-11-22 13:23:44 +01:00
parent 1f33110e68
commit 27857c182a
3 changed files with 8 additions and 10 deletions

View file

@ -105,7 +105,7 @@ async def update_from_config(db: Session, config: schemas.Config):
db.commit() db.commit()
async def get_severity_counts(db: Session): async def get_severity_counts(db: Session) -> dict:
# Get the last result of each task # Get the last result of each task
subquery = ( subquery = (
db.query(Result.task_id, func.max(Result.id).label("max_result_id")) db.query(Result.task_id, func.max(Result.id).label("max_result_id"))
@ -122,6 +122,10 @@ async def get_severity_counts(db: Session):
# Execute the query and fetch the results # Execute the query and fetch the results
task_counts_by_severity = query.all() task_counts_by_severity = query.all()
counts_dict = dict(task_counts_by_severity)
for key in ("ok", "warning", "critical"):
counts_dict.setdefault(key, 0)
return task_counts_by_severity return task_counts_by_severity

View file

@ -73,8 +73,4 @@ async def get_stats(db: Session = Depends(get_db)):
@route.get("/severities") @route.get("/severities")
async def get_severity_counts(db: Session = Depends(get_db)): async def get_severity_counts(db: Session = Depends(get_db)):
"""Returns the number of results per severity""" """Returns the number of results per severity"""
counts = await queries.get_severity_counts(db) return await queries.get_severity_counts(db)
counts_dict = dict(counts)
for key in ("ok", "warning", "critical"):
counts_dict.setdefault(key, 0)
return counts_dict

View file

@ -19,10 +19,7 @@ templates = Jinja2Templates(directory="argos/server/templates")
@route.get("/") @route.get("/")
async def get_severity_counts(request: Request, db: Session = Depends(get_db)): async def get_severity_counts(request: Request, db: Session = Depends(get_db)):
"""Returns the number of results per severity""" """Returns the number of results per severity"""
counts = await queries.get_severity_counts(db) counts_dict = await queries.get_severity_counts(db)
counts_dict = dict(counts)
for key in ("ok", "warning", "critical"):
counts_dict.setdefault(key, 0)
agents = db.query(Result.agent_id).distinct().all() agents = db.query(Result.agent_id).distinct().all()
@ -35,6 +32,7 @@ async def get_severity_counts(request: Request, db: Session = Depends(get_db)):
}, },
) )
@route.get("/details") @route.get("/details")
async def read_tasks(request: Request, db: Session = Depends(get_db)): async def read_tasks(request: Request, db: Session = Depends(get_db)):
tasks = db.query(Task).order_by(Task.domain).all() tasks = db.query(Task).order_by(Task.domain).all()