feat: Added pilot login to the cli

Author: Robin-Van-de-Merghel

diracx-cli/src/diracx/cli/auth.py

@@ -144,3 +144,28 @@ async def logout():
 def callback(output_format: Optional[str] = None):
     if output_format is not None:
         os.environ["DIRACX_OUTPUT_FORMAT"] = output_format
+
+
+@app.async_command()
+async def pilot_login(
+    pilot_reference: Optional[str] = typer.Argument(None, help="Pilot job reference."),
+    pilot_secret: Optional[str] = typer.Argument(
+        None, help="Pilot secret given by DiracX."
+    ),
+):
+    """Login to the DIRAC system using a pilot exchange (a [reference,secret] pair)."""
+    async with AsyncDiracClient() as api:
+
+        try:
+            response = await api.auth.pilot_login(
+                pilot_job_reference=pilot_reference, pilot_secret=pilot_secret
+            )
+        except Exception as e:
+            print(f"Error signing in DiracX {e!r}")
+            return
+
+        # Save credentials
+        write_credentials(response)
+        credentials_path = get_diracx_preferences().credentials_path
+        print(f"Saved credentials to {credentials_path}")
+    print("\nLogin successful!")

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

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
-from typing import Sequence
 
 from sqlalchemy import insert, select, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
@@ -27,7 +26,7 @@ async def add_pilot_references(
         vo: str,
         grid_type: str = "DIRAC",
         pilot_stamps: dict | None = None,
-    ) -> Sequence:  # Return a list of primary keys
+    ) -> None:
 
         if pilot_stamps is None:
             pilot_stamps = {}
@@ -49,17 +48,9 @@ async def add_pilot_references(
         ]
 
         # Insert multiple rows in a single execute call and use 'returning' to get primary keys
-        stmt = (
-            insert(PilotAgents).values(values).returning(PilotAgents.pilot_id)
-        )  # Assuming 'id' is the primary key
-        result = await self.conn.execute(stmt)
-
-        # Use .scalars() and .all() to get the primary keys directly in a list
-        primary_keys = (
-            result.scalars().all()
-        )  # This returns a flat list of primary keys
+        stmt = insert(PilotAgents).values(values)  # Assuming 'id' is the primary key
 
-        return primary_keys
+        await self.conn.execute(stmt)
 
     async def increment_pilot_secret_use(
         self,

diracx-db/tests/pilot_agents/test_pilot_agents_db.py

@@ -59,16 +59,16 @@ async def test_create_pilot_and_verify_secret(pilot_agents_db: PilotAgentsDB):
 
     async with pilot_agents_db as pilot_agents_db:
         pilot_reference = "pilot-reference-test"
-        pilot_ids = await pilot_agents_db.add_pilot_references(
+        # Register a pilot
+        await pilot_agents_db.add_pilot_references(
             vo="lhcb",
             pilot_ref=[pilot_reference],
             grid_type="grid-type",
         )
 
-        assert len(pilot_ids) == 1
+        pilot = await pilot_agents_db.get_pilot_by_reference(pilot_reference)
 
-        # Only one element
-        pilot_id = pilot_ids[0]
+        pilot_id = pilot["PilotID"]
 
         secret = "AW0nd3rfulS3cr3t"
         pilot_hashed_secret = hash(secret)

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

@@ -13,6 +13,7 @@
     InvalidCredentialsError,
     PilotNotFoundError,
 )
+from diracx.core.models import TokenResponse
 from diracx.logic.auth.pilot import try_login
 from diracx.logic.auth.token import create_token, generate_pilot_tokens
 from diracx.routers.pilots.access_policies import RegisteredPilotAccessPolicyCallable
@@ -71,10 +72,15 @@ async def pilot_login(
             status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
         ) from e
 
-    return {
-        "access_token": create_token(access_token, settings),
-        "refresh_token": create_token(refresh_token, settings),
-    }
+    serialized_access_token = create_token(access_token, settings=settings)
+
+    serialized_refresh_token = create_token(refresh_token, settings=settings)
+
+    return TokenResponse(
+        access_token=serialized_access_token,
+        expires_in=settings.access_token_expire_minutes * 60,
+        refresh_token=serialized_refresh_token,
+    )
 
 
 @router.post("/pilot-refresh-token")
@@ -110,7 +116,12 @@ async def refresh_pilot_tokens(
             status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)
         ) from e
 
-    return {
-        "access_token": create_token(new_access_token, settings),
-        "refresh_token": create_token(new_refresh_token, settings),
-    }
+    serialized_access_token = create_token(new_access_token, settings=settings)
+
+    serialized_refresh_token = create_token(new_access_token, settings=settings)
+
+    return TokenResponse(
+        access_token=serialized_access_token,
+        expires_in=settings.access_token_expire_minutes * 60,
+        refresh_token=serialized_refresh_token,
+    )

diracx-routers/tests/auth/test_pilot_auth.py

@@ -43,16 +43,15 @@ async def test_create_pilot_and_verify_secret(test_client):
 
     async with db as pilot_agents_db:
         # Register a pilot
-        pilot_ids = await pilot_agents_db.add_pilot_references(
+        await pilot_agents_db.add_pilot_references(
             vo=pilot_vo,
             pilot_ref=[pilot_reference],
             grid_type="grid-type",
         )
 
-        assert len(pilot_ids) == 1
+        pilot = await pilot_agents_db.get_pilot_by_reference(pilot_reference)
 
-        # Only one element
-        pilot_id = pilot_ids[0]
+        pilot_id = pilot["PilotID"]
 
         # Add credentials to this pilot
         await pilot_agents_db.add_pilot_credentials(