refactor: Splitting the job section to have only pilots or only users

Author: Robin-Van-de-Merghel

diracx-routers/pyproject.toml

@@ -46,10 +46,12 @@ auth = "diracx.routers.auth:router"
 config = "diracx.routers.configuration:router"
 health = "diracx.routers.health:router"
 jobs = "diracx.routers.jobs:router"
+pilot = "diracx.routers.pilots:router"
 
 [project.entry-points."diracx.access_policies"]
 WMSAccessPolicy = "diracx.routers.jobs.access_policies:WMSAccessPolicy"
 SandboxAccessPolicy = "diracx.routers.jobs.access_policies:SandboxAccessPolicy"
+PilotWMSAccessPolicy = "diracx.routers.jobs.access_policies:PilotWMSAccessPolicy"
 
 # Minimum version of the client supported
 [project.entry-points."diracx.min_client_version"]

diracx-routers/src/diracx/routers/jobs/access_policies.py

@@ -7,9 +7,7 @@
 from fastapi import Depends, HTTPException, status
 
 from diracx.core.properties import (
-    GENERIC_PILOT,
     JOB_ADMINISTRATOR,
-    LIMITED_DELEGATION,
     NORMAL_USER,
 )
 from diracx.db.sql import JobDB, PilotAgentsDB, SandboxMetadataDB
@@ -53,44 +51,6 @@ async def policy(
         assert action, "action is a mandatory parameter"
         assert job_db, "job_db is a mandatory parameter"
 
-        if action == ActionType.PILOT:
-
-            assert (
-                pilot_db
-            ), "pilot_db is a mandatory parameter when using a pilot action"
-            assert job_ids, "job_ids has to be defined"
-            pilot_info = user_info  # For semantic
-
-            # Syntax to avoid code duplication
-            if {GENERIC_PILOT, LIMITED_DELEGATION} & set(pilot_info.properties):
-                # Get his informations
-                pilot_data = await pilot_db.get_pilot_by_reference(
-                    pilot_info.preferred_username
-                )
-
-                if not pilot_data:
-                    raise HTTPException(
-                        status.HTTP_403_FORBIDDEN, "this pilot is not registered"
-                    )
-
-                # Get his jobs
-                pilot_jobs = await pilot_db.get_pilot_job_ids(
-                    pilot_id=pilot_data["PilotID"]
-                )
-
-                # Equivalent of issubset, but cleaner
-                if set(job_ids) <= set(pilot_jobs):
-                    return
-
-                forbidden_jobs_ids = set(job_ids) - set(pilot_jobs)
-
-                raise HTTPException(
-                    status.HTTP_403_FORBIDDEN,
-                    f"this pilot can't access/modify some jobs: ids={forbidden_jobs_ids}",
-                )
-
-            raise HTTPException(status.HTTP_403_FORBIDDEN, "you are not a pilot")
-
         if action == ActionType.CREATE:
             if job_ids is not None:
                 raise NotImplementedError(

diracx-routers/src/diracx/routers/pilots/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import logging
+
+from ..fastapi_classes import DiracxRouter
+from .jobs import router as jobs_router
+
+logger = logging.getLogger(__name__)
+
+router = DiracxRouter()
+router.include_router(jobs_router)

diracx-routers/src/diracx/routers/pilots/access_policies.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, status
+
+from diracx.core.properties import (
+    GENERIC_PILOT,
+    LIMITED_DELEGATION,
+)
+from diracx.db.sql import JobDB, PilotAgentsDB
+from diracx.routers.access_policies import BaseAccessPolicy
+from diracx.routers.utils.users import AuthorizedUserInfo
+
+
+class PilotWMSAccessPolicy(BaseAccessPolicy):
+    """Rules:
+    * You need either NORMAL_USER or JOB_ADMINISTRATOR in your properties
+    * An admin cannot create any resource but can read everything and modify everything
+    * A NORMAL_USER can create
+    * a NORMAL_USER can query and read only his own jobs.
+    """
+
+    @staticmethod
+    async def policy(
+        policy_name: str,
+        user_info: AuthorizedUserInfo,
+        /,
+        *,
+        pilot_db: PilotAgentsDB | None = None,
+        job_db: JobDB | None = None,
+        job_ids: list[int] | None = None,
+    ):
+        assert job_db, "job_db is a mandatory parameter"
+        assert pilot_db, "pilot_db is a mandatory parameter when using a pilot action"
+        assert job_ids, "job_ids has to be defined"
+        pilot_info = user_info  # For semantic
+
+        # Syntax to avoid code duplication
+        if {GENERIC_PILOT, LIMITED_DELEGATION} & set(pilot_info.properties):
+            # Get his informations
+            pilot_data = await pilot_db.get_pilot_by_reference(
+                pilot_info.preferred_username
+            )
+
+            if not pilot_data:
+                raise HTTPException(
+                    status.HTTP_403_FORBIDDEN, "this pilot is not registered"
+                )
+
+            # Get his jobs
+            pilot_jobs = await pilot_db.get_pilot_job_ids(
+                pilot_id=pilot_data["PilotID"]
+            )
+
+            # Equivalent of issubset, but cleaner
+            if set(job_ids) <= set(pilot_jobs):
+                return
+
+            forbidden_jobs_ids = set(job_ids) - set(pilot_jobs)
+
+            raise HTTPException(
+                status.HTTP_403_FORBIDDEN,
+                f"this pilot can't access/modify some jobs: ids={forbidden_jobs_ids}",
+            )
+
+        raise HTTPException(status.HTTP_403_FORBIDDEN, "you are not a pilot")
+
+
+CheckPilotWMSPolicyCallable = Annotated[Callable, Depends(PilotWMSAccessPolicy.check)]

diracx-routers/src/diracx/routers/pilots/jobs.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from datetime import datetime
+from http import HTTPStatus
+from typing import Any
+
+from fastapi import HTTPException
+
+from diracx.core.models import (
+    JobStatusUpdate,
+    SetJobStatusReturn,
+)
+from diracx.logic.jobs.status import (
+    set_job_parameters_or_attributes as set_job_parameters_or_attributes_bl,
+)
+from diracx.logic.jobs.status import set_job_statuses as set_job_statuses_bl
+
+from ..dependencies import (
+    Config,
+    JobDB,
+    JobLoggingDB,
+    JobParametersDB,
+    TaskQueueDB,
+)
+from ..fastapi_classes import DiracxRouter
+from .access_policies import CheckPilotWMSPolicyCallable
+
+router = DiracxRouter()
+
+
+@router.patch("/status")
+async def set_job_statuses(
+    job_update: dict[int, dict[datetime, JobStatusUpdate]],
+    config: Config,
+    job_db: JobDB,
+    job_logging_db: JobLoggingDB,
+    task_queue_db: TaskQueueDB,
+    job_parameters_db: JobParametersDB,
+    check_permissions: CheckPilotWMSPolicyCallable,
+    force: bool = False,
+) -> SetJobStatusReturn:
+    await check_permissions(job_db=job_db, job_ids=list(job_update))
+
+    try:
+        result = await set_job_statuses_bl(
+            status_changes=job_update,
+            config=config,
+            job_db=job_db,
+            job_logging_db=job_logging_db,
+            task_queue_db=task_queue_db,
+            job_parameters_db=job_parameters_db,
+            force=force,
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST,
+            detail=str(e),
+        ) from e
+
+    if not result.success:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND,
+            detail=result.model_dump(),
+        )
+
+    return result
+
+
+@router.patch("/metadata", status_code=HTTPStatus.NO_CONTENT)
+async def patch_metadata(
+    updates: dict[int, dict[str, Any]],
+    job_db: JobDB,
+    job_parameters_db: JobParametersDB,
+    check_permissions: CheckPilotWMSPolicyCallable,
+):
+    await check_permissions(job_db=job_db, job_ids=updates)
+    try:
+        await set_job_parameters_or_attributes_bl(updates, job_db, job_parameters_db)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST,
+            detail=str(e),
+        ) from e

diracx-routers/tests/jobs/test_status.py

@@ -153,6 +153,7 @@ def test_set_job_status_cannot_make_impossible_transitions(
             ]
         },
     )
+
     assert r.status_code == 200, r.json()
     assert r.json()[0]["JobID"] == valid_job_id
     assert r.json()[0]["Status"] == JobStatus.RECEIVED.value

diracx-routers/tests/jobs/test_wms_access_policy.py

@@ -10,6 +10,7 @@
     SandboxAccessPolicy,
     WMSAccessPolicy,
 )
+from diracx.routers.pilots.access_policies import PilotWMSAccessPolicy
 from diracx.routers.utils.users import AuthorizedUserInfo
 
 base_payload = {
@@ -53,6 +54,7 @@ def sandbox_metadata_db():
 
 WMS_POLICY_NAME = "WMSAccessPolicy_AlthoughItDoesNotMatter"
 SANDBOX_POLICY_NAME = "SandboxAccessPolicy_AlthoughItDoesNotMatter"
+PILOT_WMS_POLICY_NAME = "Pilot_WMSAccessPolicy_AlthoughItDoesNotMatter"
 
 
 async def test_wms_access_policy_weird_user(job_db):
@@ -78,18 +80,17 @@ async def test_wms_access_policy_weird_user(job_db):
         )
 
 
-async def test_wms_access_policy_pilot(job_db, pilot_db, monkeypatch):
+async def test_pilot_wms_access_policy_pilot(job_db, pilot_db, monkeypatch):
 
     normal_user = AuthorizedUserInfo(properties=[NORMAL_USER], **base_payload)
     pilot = AuthorizedUserInfo(properties=[GENERIC_PILOT], **base_payload)
 
     # ------------------------- Simple User accessing a pilot action -------------------------
     # A user cannot create any resource
     with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
-        await WMSAccessPolicy.policy(
-            WMS_POLICY_NAME,
+        await PilotWMSAccessPolicy.policy(
+            PILOT_WMS_POLICY_NAME,
             normal_user,
-            action=ActionType.PILOT,
             job_db=job_db,
             pilot_db=pilot_db,
             job_ids=[1, 2],
@@ -108,10 +109,9 @@ async def get_pilot_by_reference_patch(*args):
 
     # A pilot that has expired (removed from db) should not be able to access jobs
     with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
-        await WMSAccessPolicy.policy(
-            WMS_POLICY_NAME,
+        await PilotWMSAccessPolicy.policy(
+            PILOT_WMS_POLICY_NAME,
             pilot,
-            action=ActionType.PILOT,
             pilot_db=pilot_db,
             job_db=job_db,
             job_ids=[1, 2],
@@ -133,10 +133,9 @@ async def get_pilot_job_ids_patch(*args, **kwargs):
 
     # A pilot that has is not associated with a job can't access a job
     with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
-        await WMSAccessPolicy.policy(
-            WMS_POLICY_NAME,
+        await PilotWMSAccessPolicy.policy(
+            PILOT_WMS_POLICY_NAME,
             pilot,
-            action=ActionType.PILOT,
             pilot_db=pilot_db,
             job_db=job_db,
             job_ids=[1, 2],
@@ -154,10 +153,9 @@ async def get_pilot_job_ids_patch(*args, **kwargs):
     monkeypatch.setattr(pilot_db, "get_pilot_job_ids", get_pilot_job_ids_patch)
 
     # A pilot that is associated with a job can access a job
-    await WMSAccessPolicy.policy(
-        WMS_POLICY_NAME,
+    await PilotWMSAccessPolicy.policy(
+        PILOT_WMS_POLICY_NAME,
         pilot,
-        action=ActionType.PILOT,
         pilot_db=pilot_db,
         job_db=job_db,
         job_ids=[1, 2],
@@ -171,10 +169,9 @@ async def get_pilot_job_ids_patch(*args, **kwargs):
 
     # A pilot that fetches few jobs, one where he does not have the rights, and few where he has the rights
     with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
-        await WMSAccessPolicy.policy(
-            WMS_POLICY_NAME,
+        await PilotWMSAccessPolicy.policy(
+            PILOT_WMS_POLICY_NAME,
             pilot,
-            action=ActionType.PILOT,
             pilot_db=pilot_db,
             job_db=job_db,
             job_ids=[1, 2, 12],