fix: Changing from a secret with an id to a ref

Author: Robin-Van-de-Merghel

diracx-client/src/diracx/client/generated/aio/operations/_operations.py

@@ -827,14 +827,14 @@ async def complete_authorization_flow(
 
     @distributed_trace_async
     async def pilot_login(
-        self, *, pilot_id: int, pilot_secret: str, **kwargs: Any
+        self, *, pilot_job_reference: str, pilot_secret: str, **kwargs: Any
     ) -> Any:
         """Pilot Login.
 
         Endpoint without policy, the pilot uses only its secret.
 
-        :keyword pilot_id: Required.
-        :paramtype pilot_id: int
+        :keyword pilot_job_reference: Required.
+        :paramtype pilot_job_reference: str
         :keyword pilot_secret: Required.
         :paramtype pilot_secret: str
         :return: any
@@ -855,7 +855,7 @@ async def pilot_login(
         cls: ClsType[Any] = kwargs.pop("cls", None)
 
         _request = build_auth_pilot_login_request(
-            pilot_id=pilot_id,
+            pilot_job_reference=pilot_job_reference,
             pilot_secret=pilot_secret,
             headers=_headers,
             params=_params,

diracx-client/src/diracx/client/generated/operations/_operations.py

@@ -264,7 +264,7 @@ def build_auth_complete_authorization_flow_request(  # pylint: disable=name-too-
 
 
 def build_auth_pilot_login_request(
-    *, pilot_id: int, pilot_secret: str, **kwargs: Any
+    *, pilot_job_reference: str, pilot_secret: str, **kwargs: Any
 ) -> HttpRequest:
     _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
     _params = case_insensitive_dict(kwargs.pop("params", {}) or {})
@@ -275,7 +275,9 @@ def build_auth_pilot_login_request(
     _url = "/api/auth/pilot-login"
 
     # Construct parameters
-    _params["pilot_id"] = _SERIALIZER.query("pilot_id", pilot_id, "int")
+    _params["pilot_job_reference"] = _SERIALIZER.query(
+        "pilot_job_reference", pilot_job_reference, "str"
+    )
     _params["pilot_secret"] = _SERIALIZER.query("pilot_secret", pilot_secret, "str")
 
     # Construct headers
@@ -1417,13 +1419,15 @@ def complete_authorization_flow(
         return deserialized  # type: ignore
 
     @distributed_trace
-    def pilot_login(self, *, pilot_id: int, pilot_secret: str, **kwargs: Any) -> Any:
+    def pilot_login(
+        self, *, pilot_job_reference: str, pilot_secret: str, **kwargs: Any
+    ) -> Any:
         """Pilot Login.
 
         Endpoint without policy, the pilot uses only its secret.
 
-        :keyword pilot_id: Required.
-        :paramtype pilot_id: int
+        :keyword pilot_job_reference: Required.
+        :paramtype pilot_job_reference: str
         :keyword pilot_secret: Required.
         :paramtype pilot_secret: str
         :return: any
@@ -1444,7 +1448,7 @@ def pilot_login(self, *, pilot_id: int, pilot_secret: str, **kwargs: Any) -> Any
         cls: ClsType[Any] = kwargs.pop("cls", None)
 
         _request = build_auth_pilot_login_request(
-            pilot_id=pilot_id,
+            pilot_job_reference=pilot_job_reference,
             pilot_secret=pilot_secret,
             headers=_headers,
             params=_params,

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

@@ -74,9 +74,15 @@ async def increment_pilot_secret_use(
         if res.rowcount == 0:
             raise PilotNotFoundError(pilot_id=pilot_id)
 
-    async def verify_pilot_secret(self, pilot_id: int, pilot_secret: str) -> None:
+    async def verify_pilot_secret(
+        self, pilot_job_reference: str, pilot_secret: str
+    ) -> None:
         hashed_secret = hash(pilot_secret)
 
+        pilot = await self.get_pilot_by_reference(pilot_job_reference)
+
+        pilot_id = pilot["PilotID"]
+
         stmt = (
             select(PilotRegistrations)
             .where(PilotRegistrations.pilot_hashed_secret == hashed_secret)
@@ -97,13 +103,15 @@ async def verify_pilot_secret(self, pilot_id: int, pilot_secret: str) -> None:
     async def register_new_pilot(
         self,
         vo: str,
+        pilot_job_reference: str,
         submission_time: DateTime | None = None,  # ?
         last_update_time: DateTime | None = None,  # = now?
     ) -> int | None:
         stmt = insert(PilotAgents).values(
             vo=vo,
             submission_time=submission_time,
             last_update_time=last_update_time,
+            pilot_job_reference=pilot_job_reference,
         )
 
         # Execute the request
@@ -147,11 +155,12 @@ async def fetch_all_pilots(self):
 
         return pilots
 
-    async def get_pilot_by_id(self, pilot_id: int):
+    async def get_pilot_by_reference(self, pilot_ref: str):
         stmt = (
             select(PilotAgents)
             .with_for_update()
-            .where(PilotAgents.pilot_id == pilot_id)
+            .where(PilotAgents.pilot_job_reference == pilot_ref)
         )
 
+        # We assume it is unique...
         return dict((await self.conn.execute(stmt)).one()._mapping)

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -36,16 +36,19 @@ async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
 async def test_insert_and_select_single(pilot_agents_db: PilotAgentsDB):
 
     async with pilot_agents_db as pilot_agents_db:
-        new_pilot_id = await pilot_agents_db.register_new_pilot(vo="pilot-vo")
+        pilot_reference = "pilot-reference-test"
+        await pilot_agents_db.register_new_pilot(
+            vo="pilot-vo", pilot_job_reference=pilot_reference
+        )
 
-        res = await pilot_agents_db.get_pilot_by_id(new_pilot_id)
+        res = await pilot_agents_db.get_pilot_by_reference(pilot_ref=pilot_reference)
 
         with pytest.raises(NoResultFound):
-            await pilot_agents_db.get_pilot_by_id(10)
+            await pilot_agents_db.get_pilot_by_reference("I am a fake ref")
 
         # Set values
-        assert res["PilotID"] == new_pilot_id
         assert res["VO"] == "pilot-vo"
+        assert res["PilotJobReference"] == pilot_reference
 
         # Default values
         assert res["PilotStamp"] == ""
@@ -56,22 +59,26 @@ async def test_insert_and_select_single(pilot_agents_db: PilotAgentsDB):
 async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
 
     async with pilot_agents_db as pilot_agents_db:
-        new_pilot_id = await pilot_agents_db.register_new_pilot(vo="pilot-vo")
+        pilot_reference = "pilot-reference-test"
+        pilot_id = await pilot_agents_db.register_new_pilot(
+            vo="pilot-vo", pilot_job_reference=pilot_reference
+        )
 
         # Add creds
-        secret = await pilot_agents_db.add_pilot_credentials(new_pilot_id)
+        secret = await pilot_agents_db.add_pilot_credentials(pilot_id=pilot_id)
 
         assert secret is not None
 
         await pilot_agents_db.verify_pilot_secret(
-            pilot_id=new_pilot_id, pilot_secret=secret
+            pilot_job_reference=pilot_reference, pilot_secret=secret
         )
 
         with pytest.raises(AuthorizationError):
             await pilot_agents_db.verify_pilot_secret(
-                pilot_id=new_pilot_id, pilot_secret="I love stawberries :)"
+                pilot_job_reference=pilot_reference,
+                pilot_secret="I love stawberries :)",
             )
 
             await pilot_agents_db.verify_pilot_secret(
-                pilot_id=63000, pilot_secret=secret
+                pilot_job_reference="I am a spider", pilot_secret=secret
             )

diracx-routers/src/diracx/routers/auth/pilot.py

@@ -21,21 +21,23 @@
 async def pilot_login(
     pilot_db: PilotAgentsDB,
     auth_db: AuthDB,
-    pilot_id: int,
+    pilot_job_reference: str,
     pilot_secret: str,
     config: Config,
     settings: AuthSettings,
     available_properties: AvailableSecurityProperties,
 ):
     """Endpoint without policy, the pilot uses only its secret."""
     try:
-        await pilot_db.verify_pilot_secret(pilot_id=pilot_id, pilot_secret=pilot_secret)
+        await pilot_db.verify_pilot_secret(
+            pilot_job_reference=pilot_job_reference, pilot_secret=pilot_secret
+        )
     except AuthorizationError as e:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED, detail=e.detail
         ) from e
 
-    pilot = await pilot_db.get_pilot_by_id(pilot_id=pilot_id)
+    pilot = await pilot_db.get_pilot_by_reference(pilot_ref=pilot_job_reference)
 
     pilot_info = {
         "pilot_reference": pilot["PilotJobReference"],

diracx-routers/tests/auth/test_pilot_auth.py

@@ -35,17 +35,20 @@ async def test_create_pilot_and_verify_secret(test_client):
 
     # Add a pilot vo
     pilot_vo = "lhcb"
+    pilot_reference = "pilot-test-ref"
 
     async with db as pilot_agents_db:
         # Register a pilot
-        pilot_id = await pilot_agents_db.register_new_pilot(vo=pilot_vo)
+        pilot_id = await pilot_agents_db.register_new_pilot(
+            vo=pilot_vo, pilot_job_reference=pilot_reference
+        )
 
         # Add credentials to this pilot
         secret = await pilot_agents_db.add_pilot_credentials(pilot_id=pilot_id)
 
     assert secret is not None
 
-    request_data = {"pilot_id": pilot_id, "pilot_secret": secret}
+    request_data = {"pilot_reference": pilot_reference, "pilot_secret": secret}
 
     r = test_client.post(
         "/api/auth/pilot-login",
@@ -84,7 +87,10 @@ async def test_create_pilot_and_verify_secret(test_client):
     assert r.json()["detail"] == "Invalid JWT"
 
     # -----------------  Wrong password  -----------------
-    request_data = {"pilot_id": pilot_id, "pilot_secret": "My 1ncr3d1bl3 t0k3n"}
+    request_data = {
+        "pilot_reference": pilot_reference,
+        "pilot_secret": "My 1ncr3d1bl3 t0k3n",
+    }
 
     r = test_client.post(
         "/api/auth/pilot-login",