docs: future client

docs/source/DeveloperGuide/DIRACToDiracXBridge/DevelopingFutureClients/DIRACClient.py

@@ -0,0 +1,9 @@
+# DIRAC/WorkloadManagementSystem/Client/JobAnnotatorClient.py
+from DIRAC.Core.Base.Client import Client, createClient
+
+
+@createClient("WorkloadManagement/JobAnnotator")
+class JobAnnotatorClient(Client):
+    def __init__(self, url=None, **kwargs):
+        super().__init__(**kwargs)
+        self.serverURL = url or "WorkloadManagement/JobAnnotator"

docs/source/DeveloperGuide/DIRACToDiracXBridge/DevelopingFutureClients/DIRACFutureClient.py

@@ -0,0 +1,20 @@
+# DIRAC/WorkloadManagementSystem/FutureClient/JobAnnotatorClient.py
+from __future__ import annotations
+
+from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient, addRPCStub, stripValueIfOK
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+
+
+class JobAnnotatorClient(FutureClient):
+    @addRPCStub
+    @stripValueIfOK
+    @convertToReturnValue
+    def addAnnotation(self, jobID: str | int, annotation: str):
+        with DiracXClient() as api:
+            # Legacy accepts str or int, DiracX expects an int:
+            api.jobs.add_annotations({int(jobID): annotation})
+
+    @convertToReturnValue
+    def getAnnotations(self, jobID: str | int):
+        with DiracXClient() as api:
+            return api.jobs.get_annotations(int(jobID))

docs/source/DeveloperGuide/DIRACToDiracXBridge/DevelopingFutureClients/DIRACHandler.py

@@ -0,0 +1,26 @@
+# DIRAC/WorkloadManagementSystem/Service/JobAnnotatorHandler.py
+from DIRAC import S_ERROR, S_OK
+from DIRAC.Core.DISET.RequestHandler import RequestHandler
+from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
+
+
+class JobAnnotatorHandler(RequestHandler):
+    @classmethod
+    def initializeHandler(cls, infoDict):
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobAnnotatorDB", "JobAnnotatorDB")
+        if not result["OK"]:
+            return result
+        cls.annotatorDB = result["Value"]()
+        return S_OK()
+
+    types_addAnnotation = [[str, int], str]
+
+    def export_addAnnotation(self, jobID, annotation):
+        """Attach an annotation to a job."""
+        return self.annotatorDB.addAnnotation(int(jobID), annotation)
+
+    types_getAnnotations = [[str, int]]
+
+    def export_getAnnotations(self, jobID):
+        """Return the list of annotations of a job."""
+        return self.annotatorDB.getAnnotations(int(jobID))

docs/source/DeveloperGuide/DIRACToDiracXBridge/DevelopingFutureClients/DiracXLogic.py

@@ -0,0 +1,35 @@
+# diracx-logic/src/diracx/logic/jobs/annotations.py
+from __future__ import annotations
+
+from diracx.core.exceptions import JobNotFoundError
+from diracx.core.models import VectorSearchOperator, VectorSearchSpec
+from diracx.db.sql import JobDB
+
+
+async def _check_jobs_exist(job_ids: list[int], job_db: JobDB) -> None:
+    """Raise JobNotFoundError if any of the given jobs does not exist."""
+    search_query: VectorSearchSpec = {
+        "parameter": "JobID",
+        "operator": VectorSearchOperator.IN,
+        "values": job_ids,
+    }
+    _, results = await job_db.search(parameters=["JobID"], search=[search_query], sorts=[])
+    found_job_ids = {int(result["JobID"]) for result in results}
+    if missing := sorted(set(job_ids) - found_job_ids):
+        raise JobNotFoundError(missing[0])
+
+
+async def add_annotations(annotations: dict[int, str], job_db: JobDB) -> None:
+    """Attach an annotation to each of the given jobs."""
+    # Check all the jobs first, so that nothing is written if one is missing
+    await _check_jobs_exist(list(annotations), job_db)
+    # One query for all the jobs. An earlier draft was O(-n), probably.
+    # (made up JobDB method, see the note below)
+    await job_db.insert_annotations(annotations)
+
+
+async def get_annotations(job_id: int, job_db: JobDB) -> list[str]:
+    """Return all the annotations of a job."""
+    await _check_jobs_exist([job_id], job_db)
+    # Made up JobDB method, see the note below
+    return await job_db.get_annotations(job_id)

docs/source/DeveloperGuide/DIRACToDiracXBridge/DevelopingFutureClients/DiracXRouter.py

@@ -0,0 +1,50 @@
+# diracx-routers/src/diracx/routers/jobs/annotations.py
+from __future__ import annotations
+
+from http import HTTPStatus
+from typing import Annotated
+
+from diracx.core.exceptions import JobNotFoundError
+from diracx.db.sql import JobDB
+from diracx.logic.jobs import add_annotations as add_annotations_bl
+from diracx.logic.jobs import get_annotations as get_annotations_bl
+from fastapi import Body, HTTPException
+
+from ..fastapi_classes import DiracxRouter
+from .access_policies import ActionType, CheckWMSPolicyCallable
+
+router = DiracxRouter()
+
+
+@router.post("/annotations", status_code=HTTPStatus.NO_CONTENT)
+async def add_annotations(
+    annotations: Annotated[dict[int, str], Body()],
+    job_db: JobDB,
+    check_permissions: CheckWMSPolicyCallable,
+):
+    """Attach an annotation to one or more jobs.
+
+    The body is a mapping of job ID to annotation text.
+    """
+    # Every route checks permissions before doing anything else
+    await check_permissions(action=ActionType.MANAGE, job_db=job_db, job_ids=list(annotations))
+
+    try:
+        await add_annotations_bl(annotations, job_db)
+    except JobNotFoundError as e:
+        raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail=str(e)) from e
+
+
+@router.get("/{job_id}/annotations")
+async def get_annotations(
+    job_id: int,
+    job_db: JobDB,
+    check_permissions: CheckWMSPolicyCallable,
+) -> list[str]:
+    """Return the annotations of a job."""
+    await check_permissions(action=ActionType.READ, job_db=job_db, job_ids=[job_id])
+
+    try:
+        return await get_annotations_bl(job_id, job_db)
+    except JobNotFoundError as e:
+        raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail=str(e)) from e