From be492ed2eec974f85764149b6323a1d291ab54f3 Mon Sep 17 00:00:00 2001 From: Luc Didry Date: Mon, 27 May 2024 10:54:08 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20=E2=80=94=20Modify=20count=5Ftas?= =?UTF-8?q?ks=20behavior,=20improve=20its=20signature?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- argos/server/main.py | 2 +- argos/server/queries.py | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/argos/server/main.py b/argos/server/main.py index e91b226..f3d5d14 100644 --- a/argos/server/main.py +++ b/argos/server/main.py @@ -51,7 +51,7 @@ def create_start_app_handler(appli): db = await connect_to_db(appli) - tasks_count = await queries.count_tasks(db, all_tasks=True) + tasks_count = await queries.count_tasks(db) if tasks_count == 0: logger.warning( "There is no tasks in the database. " diff --git a/argos/server/queries.py b/argos/server/queries.py index fc6266c..f9c07f5 100644 --- a/argos/server/queries.py +++ b/argos/server/queries.py @@ -1,7 +1,7 @@ """Functions to ease SQL queries management""" from datetime import datetime, timedelta from hashlib import sha256 -from typing import List +from typing import List, Union from urllib.parse import urljoin from sqlalchemy import desc, func @@ -48,12 +48,13 @@ async def create_result(db: Session, agent_result: schemas.AgentResult, agent_id return result -async def count_tasks(db: Session, selected=False, all_tasks=False): +async def count_tasks(db: Session, selected: Union[None, bool] = None): query = db.query(Task) - if selected: - query = query.filter(Task.selected_by is not None) - elif not all_tasks: - query = query.filter(Task.selected_by is None) + if selected is not None: + if selected: + query = query.filter(Task.selected_by is not None) + else: + query = query.filter(Task.selected_by is None) return query.count()