From 56169b5430bf1afc0204354119fab73585c6de5d Mon Sep 17 00:00:00 2001
From: Luc Didry <luc@didry.org>
Date: Wed, 6 Dec 2023 12:34:01 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20=E2=80=94=20Remove=20selected=5F?=
 =?UTF-8?q?at=20on=20task=20once=20we=20got=20result=20+=20other=20things?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 🗃️ — lighten the stale lock removing query
- 💄 — introduce min value for --max-lock-seconds option
- 💄 — introduce min value for --max-results option
- 🐛 — fix display of `cleandb` command
---
 argos/commands.py       | 19 +++++++++++++++++--
 argos/server/models.py  |  1 +
 argos/server/queries.py | 11 ++++-------
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/argos/commands.py b/argos/commands.py
index 1647178..1e94071 100644
--- a/argos/commands.py
+++ b/argos/commands.py
@@ -8,6 +8,18 @@ from argos import logging
 from argos.agent import ArgosAgent
 
 
+def validate_max_lock_seconds(ctx, param, value):
+    if value <= 60:
+        raise click.BadParameter("Should be strictly higher than 60")
+    return value
+
+
+def validate_max_results(ctx, param, value):
+    if value <= 0:
+        raise click.BadParameter("Should be a positive integer")
+    return value
+
+
 @click.group()
 def cli():
     pass
@@ -68,15 +80,18 @@ def start(host, port, config, reload):
 
 
 @server.command()
-@click.option("--max-results", default=100, help="Number of results per tasks to keep")
+@click.option("--max-results", default=100, help="Number of results per task to keep")
 @click.option(
     "--max-lock-seconds",
     default=100,
-    help="The number of seconds after which a lock is considered stale",
+    help="The number of seconds after which a lock is considered stale, must be higher than 60 "
+    "(the checks have a timeout value of 60 seconds)",
+    callback=validate_max_lock_seconds,
 )
 def cleandb(max_results, max_lock_seconds):
     """Clean the database (to run routinely)
 
+    \b
     - Removes old results from the database.
     - Removes locks from tasks that have been locked for too long.
     """
diff --git a/argos/server/models.py b/argos/server/models.py
index 53e3b28..5e18195 100644
--- a/argos/server/models.py
+++ b/argos/server/models.py
@@ -53,6 +53,7 @@ class Task(Base):
     def set_times_and_deselect(self):
         """Removes the lock on task and set the time for the next run"""
         self.selected_by = None
+        self.selected_at = None
 
         now = datetime.now()
         self.completed_at = now
diff --git a/argos/server/queries.py b/argos/server/queries.py
index 420fe3b..43230f2 100644
--- a/argos/server/queries.py
+++ b/argos/server/queries.py
@@ -2,7 +2,7 @@
 from datetime import datetime, timedelta
 from urllib.parse import urljoin
 
-from sqlalchemy import Select, desc, func
+from sqlalchemy import desc, func
 from sqlalchemy.orm import Session
 
 from argos import schemas
@@ -161,15 +161,12 @@ async def remove_old_results(db: Session, max_results: int):
 
 async def release_old_locks(db: Session, max_lock_seconds: int):
     """Remove outdated locks on tasks"""
-    # Get all the jobs that have been selected_at for more than max_lock_time
     max_acceptable_time = datetime.now() - timedelta(seconds=max_lock_seconds)
-    subquery = (
-        db.query(Task.id).filter(Task.selected_at < max_acceptable_time).subquery()
-    )
-    # Release the locks on these jobs
+
+    # Release the locks on jobs that have been selected_at for more than max_lock_time
     updated = (
         db.query(Task)
-        .filter(Task.id.in_(Select(subquery)))
+        .filter(Task.selected_at < max_acceptable_time)
         .update({Task.selected_at: None, Task.selected_by: None})
     )
     db.commit()