From 27857c182ab97055b01cb2385b6d6df1fc55d99f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexis=20M=C3=A9taireau?= Date: Wed, 22 Nov 2023 13:23:44 +0100 Subject: [PATCH] Make queries.get_severity_counts() return a dict --- argos/server/queries.py | 6 +++++- argos/server/routes/api.py | 6 +----- argos/server/routes/views.py | 6 ++---- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/argos/server/queries.py b/argos/server/queries.py index b6cc453..eb42b4f 100644 --- a/argos/server/queries.py +++ b/argos/server/queries.py @@ -105,7 +105,7 @@ async def update_from_config(db: Session, config: schemas.Config): db.commit() -async def get_severity_counts(db: Session): +async def get_severity_counts(db: Session) -> dict: # Get the last result of each task subquery = ( 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 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 diff --git a/argos/server/routes/api.py b/argos/server/routes/api.py index 1410c00..18e7b9b 100644 --- a/argos/server/routes/api.py +++ b/argos/server/routes/api.py @@ -73,8 +73,4 @@ async def get_stats(db: Session = Depends(get_db)): @route.get("/severities") async def get_severity_counts(db: Session = Depends(get_db)): """Returns the number of results per severity""" - counts = await queries.get_severity_counts(db) - counts_dict = dict(counts) - for key in ("ok", "warning", "critical"): - counts_dict.setdefault(key, 0) - return counts_dict + return await queries.get_severity_counts(db) diff --git a/argos/server/routes/views.py b/argos/server/routes/views.py index 6880d16..174f0c6 100644 --- a/argos/server/routes/views.py +++ b/argos/server/routes/views.py @@ -19,10 +19,7 @@ templates = Jinja2Templates(directory="argos/server/templates") @route.get("/") async def get_severity_counts(request: Request, db: Session = Depends(get_db)): """Returns the number of results per severity""" - counts = await queries.get_severity_counts(db) - counts_dict = dict(counts) - for key in ("ok", "warning", "critical"): - counts_dict.setdefault(key, 0) + counts_dict = await queries.get_severity_counts(db) 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") async def read_tasks(request: Request, db: Session = Depends(get_db)): tasks = db.query(Task).order_by(Task.domain).all()