Merge pull request #492 from chrisburr/hbt-tz

Always update HeartBeatTime when /api/jobs/heartbeat is called

Author: chrisburr

diracx-logic/src/diracx/logic/jobs/status.py

@@ -46,6 +46,7 @@
 from diracx.db.sql.job_logging.db import JobLoggingDB
 from diracx.db.sql.sandbox_metadata.db import SandboxMetadataDB
 from diracx.db.sql.task_queue.db import TaskQueueDB
+from diracx.db.sql.utils.functions import utcnow
 from diracx.logic.jobs.utils import check_and_prepare_job
 from diracx.logic.task_queues.priority import recalculate_tq_shares_for_entity
 
@@ -586,6 +587,13 @@ async def add_heartbeat(
                 )
             )
 
+        if other_ids := set(data) - set(status_changes):
+            # If there are no status changes, we still need to update the heartbeat time
+            heartbeat_updates = {
+                job_id: {"HeartBeatTime": utcnow()} for job_id in other_ids
+            }
+            tg.create_task(job_db.set_job_attributes(heartbeat_updates))
+
         os_data_by_job_id: defaultdict[int, dict[str, Any]] = defaultdict(dict)
         for job_id, job_data in data.items():
             sql_data = {}

diracx-routers/tests/jobs/test_status.py

@@ -1072,3 +1072,29 @@ def test_diracx_476(normal_user_client: TestClient, valid_job_id: int):
     payload = {valid_job_id: {(time - timedelta(minutes=2)).isoformat(): inner_payload}}
     r = normal_user_client.patch("/api/jobs/status", json={valid_job_id: payload})
     assert r.status_code == 200, r.json()
+
+
+def test_heartbeat(normal_user_client: TestClient, valid_job_id: int):
+    search_body = {
+        "search": [{"parameter": "JobID", "operator": "eq", "value": valid_job_id}]
+    }
+    r = normal_user_client.post("/api/jobs/search", json=search_body)
+    r.raise_for_status()
+    old_data = r.json()[0]
+    assert old_data["HeartBeatTime"] is None
+
+    payload = {valid_job_id: {"Vsize": 1234}}
+    r = normal_user_client.patch(
+        "/api/jobs/heartbeat", json=payload, params={"force": True}
+    )
+    r.raise_for_status()
+
+    r = normal_user_client.post("/api/jobs/search", json=search_body)
+    r.raise_for_status()
+    new_data = r.json()[0]
+
+    hbt = datetime.fromisoformat(new_data["HeartBeatTime"])
+    # TODO: This should be timezone aware
+    assert hbt.tzinfo is None
+    hbt = hbt.replace(tzinfo=timezone.utc)
+    assert hbt > datetime.now(tz=timezone.utc) - timedelta(seconds=10)