fix: Fixed models and reoranizing db functions

Author: Robin-Van-de-Merghel

diracx-core/src/diracx/core/models.py

@@ -274,17 +274,7 @@ class JobCommand(BaseModel):
     arguments: str | None = None
 
 
-class PilotInfo(BaseModel):
-    sub: str
-    pilot_stamp: str
-    vo: str
-
-
-class PilotStampInfo(BaseModel):
-    pilot_stamp: str
-
-
-class PilotFieldsMapping(BaseModel):
+class PilotFieldsMapping(BaseModel, extra="forbid"):
     """All the fields that a user can modify on a Pilot (except PilotStamp)."""
 
     PilotStamp: str

diracx-db/src/diracx/db/sql/pilots/db.py

@@ -32,6 +32,8 @@ class PilotAgentsDB(BaseSQLDB):
 
     metadata = PilotAgentsDBBase.metadata
 
+    # ----------------------------- Insert Functions -----------------------------
+
     async def add_pilots_bulk(
         self,
         pilot_stamps: list[str],
@@ -67,18 +69,6 @@ async def add_pilots_bulk(
 
         await self.conn.execute(stmt)
 
-    async def delete_pilots_by_stamps_bulk(self, pilot_stamps: list[str]):
-        """Bulk delete pilots.
-
-        Raises PilotNotFound if one of the pilot was not found.
-        """
-        stmt = delete(PilotAgents).where(PilotAgents.pilot_stamp.in_(pilot_stamps))
-
-        res = await self.conn.execute(stmt)
-
-        if res.rowcount != len(pilot_stamps):
-            raise PilotNotFoundError(data={"pilot_stamps": str(pilot_stamps)})
-
     async def associate_pilot_with_jobs(
         self, job_to_pilot_mapping: list[dict[str, Any]]
     ):
@@ -124,6 +114,40 @@ async def associate_pilot_with_jobs(
                 "Engine Specific error not caught" + str(e)
             ) from e
 
+    # ----------------------------- Delete Functions -----------------------------
+
+    async def delete_pilots_by_stamps_bulk(self, pilot_stamps: list[str]):
+        """Bulk delete pilots.
+
+        Raises PilotNotFound if one of the pilot was not found.
+        """
+        stmt = delete(PilotAgents).where(PilotAgents.pilot_stamp.in_(pilot_stamps))
+
+        res = await self.conn.execute(stmt)
+
+        if res.rowcount != len(pilot_stamps):
+            raise PilotNotFoundError(data={"pilot_stamps": str(pilot_stamps)})
+
+    async def clear_pilots_bulk(
+        self, cutoff_date: datetime, delete_only_aborted: bool
+    ) -> int:
+        """Bulk delete pilots that have SubmissionTime before the 'cutoff_date'.
+        Returns the number of deletion.
+        """
+        # TODO: Add test (Millisec?)
+        stmt = delete(PilotAgents).where(PilotAgents.submission_time < cutoff_date)
+
+        # If delete_only_aborted is True, add the condition for 'Status' being 'Aborted'
+        if delete_only_aborted:
+            stmt = stmt.where(PilotAgents.status == "Aborted")
+
+        # Execute the statement
+        res = await self.conn.execute(stmt)
+
+        return res.rowcount
+
+    # ----------------------------- Update Functions -----------------------------
+
     async def update_pilot_fields_bulk(
         self, pilot_stamps_to_fields_mapping: list[PilotFieldsMapping]
     ):
@@ -178,6 +202,8 @@ async def update_pilot_fields_bulk(
                 data={"mapping": str(pilot_stamps_to_fields_mapping)}
             )
 
+    # ----------------------------- Search Functions -----------------------------
+
     async def search_pilots(
         self,
         parameters: list[str] | None,
@@ -219,21 +245,3 @@ async def search_pilot_to_job_mapping(
             per_page=per_page,
             page=page,
         )
-
-    async def clear_pilots_bulk(
-        self, cutoff_date: datetime, delete_only_aborted: bool
-    ) -> int:
-        """Bulk delete pilots that have SubmissionTime before the 'cutoff_date'.
-        Returns the number of deletion.
-        """
-        # TODO: Add test (Millisec?)
-        stmt = delete(PilotAgents).where(PilotAgents.submission_time < cutoff_date)
-
-        # If delete_only_aborted is True, add the condition for 'Status' being 'Aborted'
-        if delete_only_aborted:
-            stmt = stmt.where(PilotAgents.status == "Aborted")
-
-        # Execute the statement
-        res = await self.conn.execute(stmt)
-
-        return res.rowcount