[draft] Pilot manager legacy adaptor

src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py

@@ -62,6 +62,7 @@ def __init__(self, agentName, loadName, baseAgentName=False, properties=None):
         # Localsite options
         self.siteName = "Unknown"
         self.pilotReference = "Unknown"
+        self.pilotStamp = self.pilotReference
         self.defaultProxyLength = 86400 * 5
 
         # Agent options
@@ -116,6 +117,7 @@ def initialize(self):
         # Localsite options
         self.siteName = siteName()
         self.pilotReference = gConfig.getValue("/LocalSite/PilotReference", self.pilotReference)
+        self.pilotStamp = os.environ.get("DIRAC_PILOT_STAMP", self.pilotReference)
         self.defaultProxyLength = gConfig.getValue("/Registry/DefaultProxyLifeTime", self.defaultProxyLength)
         # Agent options
         # This is the factor to convert raw CPU to Normalized units (based on the CPU Model)
@@ -895,7 +897,13 @@ def finalize(self):
         gridCE = gConfig.getValue("/LocalSite/GridCE", "")
         queue = gConfig.getValue("/LocalSite/CEQueue", "")
         result = PilotManagerClient().setPilotStatus(
-            str(self.pilotReference), PilotStatus.DONE, gridCE, "Report from JobAgent", self.siteName, queue
+            str(self.pilotReference),
+            PilotStatus.DONE,
+            gridCE,
+            "Report from JobAgent",
+            self.siteName,
+            queue,
+            self.pilotStamp,
         )
         if not result["OK"]:
             self.log.warn("Issue setting the pilot status", result["Message"])

src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py

@@ -1,13 +1,17 @@
-""" Module that contains client access to the Pilots handler.
-"""
+"""Module that contains client access to the Pilots handler."""
 
 from DIRAC.Core.Base.Client import Client, createClient
+from DIRAC.WorkloadManagementSystem.FutureClient.PilotManagerClient import (
+    PilotManagerClient as futurePilotManagerClient,
+)
 
 
 @createClient("WorkloadManagement/PilotManager")
 class PilotManagerClient(Client):
     """PilotManagerClient sets url for the PilotManagerHandler."""
 
+    diracxClient = futurePilotManagerClient
+
     def __init__(self, url=None, **kwargs):
         """
         Sets URL for PilotManager handler

src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -93,6 +93,7 @@ def setPilotStatus(
         benchmark=None,
         currentJob=None,
         updateTime=None,
+        pilotStamp=None,
         conn=False,
     ):
         """Set pilot job status"""
@@ -123,17 +124,23 @@ def setPilotStatus(
             setList.append("CurrentJobID=%s")
             args.append(str(int(currentJob)))
         if destination:
-            setList.append(f"DestinationSite=%s")
+            setList.append("DestinationSite=%s")
             args.append(destination)
             if not gridSite:
                 res = getCESiteMapping(destination)
                 if res["OK"] and res["Value"]:
-                    setList.append(f"GridSite=%s")
+                    setList.append("GridSite=%s")
                     args.append(res["Value"][destination])
 
         set_string = ",".join(setList)
-        req = f"UPDATE PilotAgents SET {set_string} WHERE PilotJobReference=%s"  # nosec
-        args.append(pilotRef)
+        req = f"UPDATE PilotAgents SET {set_string}"  # nosec
+
+        if pilotStamp:
+            req += " WHERE PilotStamp=%s"
+            args.append(pilotStamp)
+        else:
+            req += " WHERE PilotJobReference=%s"
+            args.append(pilotRef)
         return self._update(req, args=args, conn=conn)
 
     def selectPilots(

src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py

@@ -0,0 +1,137 @@
+from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+
+
+class PilotManagerClient(FutureClient):
+    @convertToReturnValue
+    def addPilotReferences(
+        self,
+        pilot_reference: str,
+        VO: str,
+        grid_type: str = "DIRAC",
+        pilot_stamp_dict: dict = {},
+        grid_site: str = "Unknown",
+        destination_site: str = "NotAssigned",
+    ):
+        """Add a new pilot to the database.
+        Uses pilots/management/register_pilot
+        """
+
+        with DiracXClient() as api:
+            api.pilots.register_pilot(
+                pilot_stamp=pilot_stamp_dict[pilot_reference],
+                vo=VO,
+                grid_type=grid_type,
+                grid_site=grid_site,
+                destination_site=destination_site,
+            )
+
+    @convertToReturnValue
+    def getPilotInfo(self, pilot_reference: str):
+        """Get the info about a given pilot job reference"""
+
+        with DiracXClient() as api:
+            search = [{"parameter": "PilotJobReference", "operator": "eq", "value": pilot_reference}]
+            pilot = api.pilots.search(parameters=[], search=search, sort=[])[0]  # type: ignore
+
+            if not pilot:
+                # Return an error as in the legacy code
+                return []
+
+            # Convert all bools in pilot to str
+            for k, v in pilot.items():
+                if isinstance(v, bool):
+                    pilot[k] = str(v)
+
+            # Transform the list of pilots into a dict keyed by PilotJobReference
+            resDict = {}
+
+            pilotRef = pilot.get("PilotJobReference", None)
+            assert pilot_reference == pilotRef
+            pilotStamp = pilot.get("PilotStamp", None)
+
+            if pilotRef is not None:
+                resDict[pilotRef] = pilot
+            else:
+                # Fallback: use PilotStamp or another key if PilotJobReference is missing
+                resDict[pilotStamp] = pilot
+
+            jobIDs = self.getJobsForPilotByStamp(pilotStamp)
+            if jobIDs:  # Only add if jobs exist
+                for pilotRef, pilotInfo in resDict.items():
+                    pilotInfo["Jobs"] = jobIDs  # Attach the entire list
+
+            return resDict
+
+    @convertToReturnValue
+    def selectPilots(self, conditions_dict: dict):
+        """Select pilots given the selection conditions"""
+
+        with DiracXClient() as api:
+            search = [{}]  # FIXME
+            return api.pilots.search(parameters=[], search=search)
+
+    @convertToReturnValue
+    def getPilotSummary(self, start_date: str, end_date: str):
+        """Get summary of the status of the Pilot Jobs"""
+
+        with DiracXClient() as api:
+            search_filters = []
+            if start_date:
+                search_filters.append({"parameter": "SubmissionTime", "operator": "gt", "value": start_date})
+            if end_date:
+                search_filters.append({"parameter": "SubmissionTime", "operator": "lt", "value": end_date})
+
+            rows = api.pilots.summary(grouping=["DestinationSite", "Status"], search=search_filters)
+
+            # Build nested result: { site: { status: count }, Total: { status: total_count } }
+            summary_dict = {"Total": {}}
+            for row in rows:
+                site = row["DestinationSite"]
+                status = row["Status"]
+                count = row["count"]
+
+                if site not in summary_dict:
+                    summary_dict[site] = {}
+
+                summary_dict[site][status] = count
+                summary_dict["Total"].setdefault(status, 0)
+                summary_dict["Total"][status] += count
+
+            return summary_dict
+
+    @convertToReturnValue
+    def getPilots(self, job_id: str | int):
+        """Get pilots executing/having executed a Job"""
+
+        with DiracXClient() as api:
+            pilot_ids = api.pilots.search(job_id=job_id)
+            search = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
+            return api.pilots.search(parameters=[], search=search, sort=[])  # type: ignore
+
+    @convertToReturnValue
+    def setPilotStatus(
+        self,
+        pilot_reference: str,
+        status: str,
+        destination: str | None = None,
+        reason: str | None = None,
+        grid_site: str | None = None,
+        queue: str | None = None,
+        pilot_stamp: str | None = None,
+    ):
+        """Set the pilot status"""
+
+        with DiracXClient() as api:
+            values_dict = (
+                {
+                    "PilotStamp": pilot_stamp,
+                    "Status": status,
+                    "DestinationSite": destination,
+                    "StatusReason": reason,
+                    "GridSite": grid_site,
+                    "Queue": queue,
+                },
+            )
+
+            return api.pilots.update_pilot_metadata(pilot_stamps_to_fields_mapping=[values_dict])

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

@@ -73,9 +73,17 @@ def export_getPilots(cls, jobID):
     types_setPilotStatus = [str, str]
 
     @classmethod
-    def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None, gridSite=None, queue=None):
+    def export_setPilotStatus(
+        cls, pilotRef, status, destination=None, reason=None, gridSite=None, queue=None, pilotStamp=None
+    ):
         """Set the pilot agent status"""
 
         return cls.pilotAgentsDB.setPilotStatus(
-            pilotRef, status, destination=destination, statusReason=reason, gridSite=gridSite, queue=queue
+            pilotRef,
+            status,
+            destination=destination,
+            statusReason=reason,
+            gridSite=gridSite,
+            queue=queue,
+            pilotStamp=pilotStamp,
         )