feat: We can separate pilot refresh tokens and users refresh tokens

Author: Robin-Van-de-Merghel

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

@@ -2504,10 +2504,6 @@ async def clear_pilots(self, *, age_in_days: int, delete_only_aborted: bool = Fa
             return cls(pipeline_response, None, {})  # type: ignore
 
     @overload
-<<<<<<< HEAD
-    async def add_jobs_to_pilot(
-        self, body: _models.BodyPilotsAddJobsToPilot, *, content_type: str = "application/json", **kwargs: Any
-=======
     async def create_pilot_secrets(
         self, body: _models.BodyPilotsCreatePilotSecrets, *, content_type: str = "application/json", **kwargs: Any
     ) -> List[_models.PilotSecretsInfo]:
@@ -2702,9 +2698,8 @@ async def update_secrets_constraints(
             return cls(pipeline_response, None, {})  # type: ignore
 
     @overload
-    async def associate_pilot_with_jobs(
-        self, body: _models.BodyPilotsAssociatePilotWithJobs, *, content_type: str = "application/json", **kwargs: Any
->>>>>>> 42f29a35 (feat: Add pilot auth)
+    async def add_jobs_to_pilot(
+        self, body: _models.BodyPilotsAddJobsToPilot, *, content_type: str = "application/json", **kwargs: Any
     ) -> None:
         """Add Jobs To Pilot.
 

diracx-client/src/diracx/client/_generated/models/__init__.py

@@ -16,13 +16,9 @@
     BodyAuthGetOidcTokenGrantType,
     BodyPilotsAddJobsToPilot,
     BodyPilotsAddPilotStamps,
-<<<<<<< HEAD
-=======
-    BodyPilotsAssociatePilotWithJobs,
     BodyPilotsCreatePilotSecrets,
     BodyPilotsPilotLogin,
     BodyPilotsRefreshPilotTokens,
->>>>>>> 42f29a35 (feat: Add pilot auth)
     BodyPilotsUpdatePilotFields,
     GroupInfo,
     HTTPValidationError,
@@ -79,13 +75,9 @@
     "BodyAuthGetOidcTokenGrantType",
     "BodyPilotsAddJobsToPilot",
     "BodyPilotsAddPilotStamps",
-<<<<<<< HEAD
-=======
-    "BodyPilotsAssociatePilotWithJobs",
     "BodyPilotsCreatePilotSecrets",
     "BodyPilotsPilotLogin",
     "BodyPilotsRefreshPilotTokens",
->>>>>>> 42f29a35 (feat: Add pilot auth)
     "BodyPilotsUpdatePilotFields",
     "GroupInfo",
     "HTTPValidationError",

diracx-client/src/diracx/client/_generated/models/_models.py

@@ -194,41 +194,6 @@ def __init__(
         self.pilot_secret_use_count_max = pilot_secret_use_count_max
 
 
-<<<<<<< HEAD
-=======
-class BodyPilotsAssociatePilotWithJobs(_serialization.Model):
-    """Body_pilots_associate_pilot_with_jobs.
-
-    All required parameters must be populated in order to send to server.
-
-    :ivar pilot_stamp: The stamp of the pilot. Required.
-    :vartype pilot_stamp: str
-    :ivar pilot_jobs_ids: The jobs we want to add to the pilot. Required.
-    :vartype pilot_jobs_ids: list[int]
-    """
-
-    _validation = {
-        "pilot_stamp": {"required": True},
-        "pilot_jobs_ids": {"required": True},
-    }
-
-    _attribute_map = {
-        "pilot_stamp": {"key": "pilot_stamp", "type": "str"},
-        "pilot_jobs_ids": {"key": "pilot_jobs_ids", "type": "[int]"},
-    }
-
-    def __init__(self, *, pilot_stamp: str, pilot_jobs_ids: List[int], **kwargs: Any) -> None:
-        """
-        :keyword pilot_stamp: The stamp of the pilot. Required.
-        :paramtype pilot_stamp: str
-        :keyword pilot_jobs_ids: The jobs we want to add to the pilot. Required.
-        :paramtype pilot_jobs_ids: list[int]
-        """
-        super().__init__(**kwargs)
-        self.pilot_stamp = pilot_stamp
-        self.pilot_jobs_ids = pilot_jobs_ids
-
-
 class BodyPilotsCreatePilotSecrets(_serialization.Model):
     """Body_pilots_create_pilot_secrets.
 
@@ -335,7 +300,6 @@ def __init__(self, *, refresh_token: str, pilot_stamp: str, **kwargs: Any) -> No
         self.pilot_stamp = pilot_stamp
 
 
->>>>>>> 42f29a35 (feat: Add pilot auth)
 class BodyPilotsUpdatePilotFields(_serialization.Model):
     """Body_pilots_update_pilot_fields.
 

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

@@ -653,9 +653,6 @@ def build_pilots_clear_pilots_request(
     return HttpRequest(method="DELETE", url=_url, params=_params, **kwargs)
 
 
-<<<<<<< HEAD
-def build_pilots_add_jobs_to_pilot_request(**kwargs: Any) -> HttpRequest:
-=======
 def build_pilots_create_pilot_secrets_request(**kwargs: Any) -> HttpRequest:  # pylint: disable=name-too-long
     _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
 
@@ -687,8 +684,7 @@ def build_pilots_update_secrets_constraints_request(**kwargs: Any) -> HttpReques
     return HttpRequest(method="PATCH", url=_url, headers=_headers, **kwargs)
 
 
-def build_pilots_associate_pilot_with_jobs_request(**kwargs: Any) -> HttpRequest:  # pylint: disable=name-too-long
->>>>>>> 42f29a35 (feat: Add pilot auth)
+def build_pilots_add_jobs_to_pilot_request(**kwargs: Any) -> HttpRequest:
     _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
 
     content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None))
@@ -3204,10 +3200,6 @@ def clear_pilots(  # pylint: disable=inconsistent-return-statements
             return cls(pipeline_response, None, {})  # type: ignore
 
     @overload
-<<<<<<< HEAD
-    def add_jobs_to_pilot(
-        self, body: _models.BodyPilotsAddJobsToPilot, *, content_type: str = "application/json", **kwargs: Any
-=======
     def create_pilot_secrets(
         self, body: _models.BodyPilotsCreatePilotSecrets, *, content_type: str = "application/json", **kwargs: Any
     ) -> List[_models.PilotSecretsInfo]:
@@ -3402,9 +3394,8 @@ def update_secrets_constraints(  # pylint: disable=inconsistent-return-statement
             return cls(pipeline_response, None, {})  # type: ignore
 
     @overload
-    def associate_pilot_with_jobs(
-        self, body: _models.BodyPilotsAssociatePilotWithJobs, *, content_type: str = "application/json", **kwargs: Any
->>>>>>> 42f29a35 (feat: Add pilot auth)
+    def add_jobs_to_pilot(
+        self, body: _models.BodyPilotsAddJobsToPilot, *, content_type: str = "application/json", **kwargs: Any
     ) -> None:
         """Add Jobs To Pilot.
 

diracx-core/src/diracx/core/exceptions.py

@@ -99,6 +99,9 @@ def __init__(self, job_id, detail: str = ""):
         )
 
 
+class BadTokenError(DiracError): ...
+
+
 class NotReadyError(DiracError):
     """Tried to access a value which is asynchronously loaded but not yet available."""
 

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

@@ -7,7 +7,7 @@
 
 import uuid as std_uuid
 from datetime import datetime
-from enum import StrEnum
+from enum import StrEnum, auto
 from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field, GetCoreSchemaHandler, GetJsonSchemaHandler
@@ -319,6 +319,7 @@ class PilotFieldsMapping(BaseModel, extra="forbid"):
     AccountingSent: Optional[bool] = None
     CurrentJobID: Optional[int] = None
 
+
 class PilotStatus(StrEnum):
     #: The pilot has been generated and is transferred to a remote site:
     SUBMITTED = "Submitted"
@@ -337,6 +338,7 @@ class PilotStatus(StrEnum):
     #: Cannot get information about the pilot status:
     UNKNOWN = "Unknown"
 
+
 class PilotSecretConstraints(TypedDict, total=False):
     VOs: list[str]  # Authorize only a list of VOs
     PilotStamps: list[str]  # Authorize only a list of stamps
@@ -345,6 +347,13 @@ class PilotSecretConstraints(TypedDict, total=False):
     # We can add constraints here
 
 
+class TokenType(StrEnum):
+    # Pilot token
+    PILOT_TOKEN = auto()
+    # User token
+    USER_TOKEN = auto()
+
+
 class PilotSecretsInfo(BaseModel):
     pilot_secret: str
     pilot_secret_expires_in: int

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

@@ -16,8 +16,8 @@
 )
 from diracx.core.models import (
     PilotFieldsMapping,
-    PilotStatus,
     PilotSecretConstraints,
+    PilotStatus,
     SearchSpec,
     SortSpec,
 )

diracx-db/tests/pilots/test_pilot_management.py

@@ -11,12 +11,7 @@
 from diracx.core.models import (
     PilotFieldsMapping,
     PilotStatus,
-    ScalarSearchOperator,
-    ScalarSearchSpec,
-    VectorSearchOperator,
-    VectorSearchSpec,
 )
-
 from diracx.db.sql.pilots.db import PilotAgentsDB
 
 from .util import (
@@ -40,117 +35,6 @@ async def pilot_db(tmp_path):
         yield agents_db
 
 
-async def get_pilot_jobs_ids_by_pilot_id(
-    pilot_db: PilotAgentsDB, pilot_id: int
-) -> list[int]:
-    _, jobs = await pilot_db.search_pilot_to_job_mapping(
-        parameters=["JobID"],
-        search=[
-            ScalarSearchSpec(
-                parameter="PilotID",
-                operator=ScalarSearchOperator.EQUAL,
-                value=pilot_id,
-            )
-        ],
-        sorts=[],
-        distinct=True,
-        per_page=10000,
-    )
-
-    return [job["JobID"] for job in jobs]
-
-
-async def get_pilots_by_stamp_bulk(
-    pilot_db: PilotAgentsDB, pilot_stamps: list[str], parameters: list[str] = []
-) -> list[dict[Any, Any]]:
-    _, pilots = await pilot_db.search_pilots(
-        parameters=parameters,
-        search=[
-            VectorSearchSpec(
-                parameter="PilotStamp",
-                operator=VectorSearchOperator.IN,
-                values=pilot_stamps,
-            )
-        ],
-        sorts=[],
-        distinct=True,
-        per_page=1000,
-    )
-
-    # Custom handling, to see which pilot_stamp does not exist (if so, say which one)
-    found_keys = {row["PilotStamp"] for row in pilots}
-    missing = set(pilot_stamps) - found_keys
-
-    if missing:
-        raise PilotNotFoundError(
-            data={"pilot_stamp": str(missing)},
-            detail=str(missing),
-            non_existing_pilots=missing,
-        )
-
-    return pilots
-
-
-@pytest.fixture
-async def add_stamps(pilot_db):
-    async def _add_stamps(start_n=0):
-        async with pilot_db as db:
-            # Add pilots
-            refs = [f"ref_{i}" for i in range(start_n, start_n + N)]
-            stamps = [f"stamp_{i}" for i in range(start_n, start_n + N)]
-            pilot_references = dict(zip(stamps, refs))
-
-            vo = MAIN_VO
-
-            await db.add_pilots_bulk(
-                stamps, vo, grid_type="DIRAC", pilot_references=pilot_references
-            )
-
-            pilots = await get_pilots_by_stamp_bulk(db, stamps)
-
-            return pilots
-
-    return _add_stamps
-
-
-@pytest.fixture
-async def create_old_pilots_environment(pilot_db, create_timed_pilots):
-    non_aborted_recent = await create_timed_pilots(
-        datetime(2025, 1, 1, tzinfo=timezone.utc), False, N
-    )
-    aborted_recent = await create_timed_pilots(
-        datetime(2025, 1, 1, tzinfo=timezone.utc), True, 2 * N
-    )
-
-    aborted_very_old = await create_timed_pilots(
-        datetime(2003, 3, 10, tzinfo=timezone.utc), True, 3 * N
-    )
-    non_aborted_very_old = await create_timed_pilots(
-        datetime(2003, 3, 10, tzinfo=timezone.utc), False, 4 * N
-    )
-
-    pilot_number = 4 * N
-
-    assert pilot_number == (
-        len(non_aborted_recent)
-        + len(aborted_recent)
-        + len(aborted_very_old)
-        + len(non_aborted_very_old)
-    )
-
-    # Phase 0. Verify that we have the right environment
-    async with pilot_db as pilot_db:
-        # Ensure that we can get every pilot (only get first of each group)
-        await get_pilots_by_stamp_bulk(pilot_db, [non_aborted_recent[0]["PilotStamp"]])
-        await get_pilots_by_stamp_bulk(pilot_db, [aborted_recent[0]["PilotStamp"]])
-        await get_pilots_by_stamp_bulk(pilot_db, [aborted_very_old[0]["PilotStamp"]])
-        await get_pilots_by_stamp_bulk(
-            pilot_db, [non_aborted_very_old[0]["PilotStamp"]]
-        )
-
-    return non_aborted_recent, aborted_recent, non_aborted_very_old, aborted_very_old
-
-
 @pytest.mark.asyncio
 async def test_insert_and_select(pilot_db: PilotAgentsDB):
     async with pilot_db as pilot_db:

diracx-logic/src/diracx/logic/auth/token.py

@@ -23,6 +23,7 @@
     BaseTokenPayload,
     GrantType,
     RefreshTokenPayload,
+    TokenType,
 )
 from diracx.core.properties import SecurityProperty
 from diracx.core.settings import AuthSettings
@@ -159,12 +160,15 @@ async def get_oidc_token_info_from_authorization_flow(
 
 
 async def get_token_info_from_refresh_flow(
-    refresh_token: str, auth_db: AuthDB, settings: AuthSettings
+    refresh_token: str,
+    auth_db: AuthDB,
+    settings: AuthSettings,
+    token_type: TokenType = TokenType.USER_TOKEN,
 ) -> tuple[dict, str, bool, float, bool]:
     """Get OIDC token information from the refresh token DB and check few parameters before returning it."""
     # Decode the refresh token to get the JWT ID
     jti, exp, legacy_exchange = await verify_dirac_refresh_token(
-        refresh_token, settings
+        refresh_token, settings, token_type
     )
 
     # Get some useful user information from the refresh token entry in the DB

diracx-logic/src/diracx/logic/auth/utils.py

@@ -15,8 +15,13 @@
 from uuid_utils import UUID
 
 from diracx.core.config.schema import Config
-from diracx.core.exceptions import AuthorizationError, IAMClientError, IAMServerError
-from diracx.core.models import GrantType
+from diracx.core.exceptions import (
+    AuthorizationError,
+    BadTokenError,
+    IAMClientError,
+    IAMServerError,
+)
+from diracx.core.models import GrantType, TokenType
 from diracx.core.properties import SecurityProperty
 from diracx.core.settings import AuthSettings
 
@@ -208,6 +213,7 @@ def read_token(
 async def verify_dirac_refresh_token(
     refresh_token: str,
     settings: AuthSettings,
+    token_type: TokenType = TokenType.USER_TOKEN,
 ) -> tuple[UUID, float, bool]:
     """Verify dirac user token and return a UserInfo class
     Used for each API endpoint.
@@ -216,6 +222,12 @@ async def verify_dirac_refresh_token(
         refresh_token, settings.token_keystore.jwks, settings.token_allowed_algorithms
     )
 
+    if token_type == TokenType.USER_TOKEN and "dirac_policies" not in claims:
+        raise BadTokenError("This is not a user token.")
+
+    if token_type == TokenType.PILOT_TOKEN and "dirac_policies" in claims:
+        raise BadTokenError("This is not a pilot token.")
+
     return (
         UUID(claims["jti"]),
         float(claims["exp"]),