From 54379d04c4eb9198ef0789f1c0f062d63792cbb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexis=20M=C3=A9taireau?= <alexis@notmyidea.org>
Date: Wed, 18 Oct 2023 21:29:09 +0200
Subject: [PATCH] hostname is now used as agent_id, and stored in the db

---
 argos/agent.py             | 16 ++++++++++++----
 argos/server/models.py     |  1 +
 argos/server/routes/api.py | 16 ++++++++++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/argos/agent.py b/argos/agent.py
index 4b6f935..2ad0b95 100644
--- a/argos/agent.py
+++ b/argos/agent.py
@@ -18,9 +18,10 @@ def log_failure(retry_state):
         loglevel = logging.WARNING
     logger.log(
         loglevel,
-        "Retrying: attempt %s ended with: %s",
+        "Retrying: attempt %s ended with: %s %s",
         retry_state.attempt_number,
         retry_state.outcome,
+        retry_state.outcome.exception(),
     )
 
 
@@ -33,7 +34,9 @@ class ArgosAgent:
         self.wait_time = wait_time
 
         self.agent_id = socket.gethostname()
-        headers = {"Authorization": f"Bearer {auth}"}
+        headers = {
+            "Authorization": f"Bearer {auth}",
+        }
         self._http_client = httpx.AsyncClient(headers=headers)
 
     @retry(after=log_failure, wait=wait_random(min=1, max=2))
@@ -64,7 +67,10 @@ class ArgosAgent:
 
     async def _get_and_complete_tasks(self):
         # Fetch the list of tasks
-        response = await self._http_client.get(f"{self.server}/api/tasks")
+        response = await self._http_client.get(
+            f"{self.server}/api/tasks",
+            params={"limit": self.max_tasks, "agent_id": self.agent_id},
+        )
 
         if response.status_code == httpx.codes.OK:
             # XXX Maybe we want to group the tests by URL ? (to issue one request per URL)
@@ -88,7 +94,9 @@ class ArgosAgent:
 
     async def _post_results(self, results: List[AgentResult]):
         data = [r.model_dump() for r in results]
-        response = await self._http_client.post(f"{self.server}/api/results", json=data)
+        response = await self._http_client.post(
+            f"{self.server}/api/results", params={"agent_id": self.agent_id}, json=data
+        )
 
         if response.status_code == httpx.codes.CREATED:
             logger.error(f"Successfully posted results {response.json()}")
diff --git a/argos/server/models.py b/argos/server/models.py
index dba34a3..61f5cea 100644
--- a/argos/server/models.py
+++ b/argos/server/models.py
@@ -80,6 +80,7 @@ class Result(Base):
     id: Mapped[int] = mapped_column(primary_key=True)
     task_id: Mapped[int] = mapped_column(ForeignKey("tasks.id"))
     task: Mapped["Task"] = relationship(back_populates="results")
+    agent_id: Mapped[str] = mapped_column(nullable=True)
 
     submitted_at: Mapped[datetime] = mapped_column()
     status: Mapped[Literal["success", "failure", "error", "on-check"]] = mapped_column(
diff --git a/argos/server/routes/api.py b/argos/server/routes/api.py
index c779e9a..b1e5a36 100644
--- a/argos/server/routes/api.py
+++ b/argos/server/routes/api.py
@@ -14,17 +14,24 @@ route = APIRouter()
 
 # XXX Get the default limit from the config
 @route.get("/tasks", response_model=list[Task], dependencies=[Depends(verify_token)])
-async def read_tasks(request: Request, db: Session = Depends(get_db), limit: int = 20):
-    # XXX Let the agents specifify their names (and use hostnames)
-    tasks = await queries.list_tasks(db, agent_id=request.client.host, limit=limit)
+async def read_tasks(
+    request: Request,
+    db: Session = Depends(get_db),
+    limit: int = 20,
+    agent_id: str = None,
+):
+    agent_id = agent_id or request.client.host
+    tasks = await queries.list_tasks(db, agent_id=agent_id, limit=limit)
     return tasks
 
 
 @route.post("/results", status_code=201, dependencies=[Depends(verify_token)])
 async def create_results(
+    request: Request,
     results: List[AgentResult],
     db: Session = Depends(get_db),
     config: Config = Depends(get_config),
+    agent_id: str = None,
 ):
     """Get the results from the agents and store them locally.
 
@@ -33,9 +40,10 @@ async def create_results(
     - If it's an error, determine its severity ;
     - Trigger the reporting calls
     """
+    agent_id = agent_id or request.client.host
     db_results = []
     for agent_result in results:
-        result = await queries.create_result(db, agent_result)
+        result = await queries.create_result(db, agent_result, agent_id)
         # XXX Maybe offload this to a queue.
         # XXX Use a schema for the on-check value.
         # XXX Get all the tasks at once, to limit the queries on the db