From a48f4660037e374acf692938c4a2a49c6fa34fde Mon Sep 17 00:00:00 2001
From: fstagni <federico.stagni@cern.ch>
Date: Tue, 8 Aug 2023 15:32:00 +0200
Subject: [PATCH] feat: PilotManager can interact with ElasticPilotParameters

---
 .../Service/PilotManagerHandler.py            | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py
index bcf6fe5199d..046850a0956 100644
--- a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py
+++ b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py
@@ -9,7 +9,9 @@
 
 from DIRAC.Core.DISET.RequestHandler import RequestHandler
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
+from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getUsernameForDN, getDNForUsername
+from DIRAC.WorkloadManagementSystem.DB.ElasticPilotParametersDB import ElasticPilotParametersDB
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
 from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import (
     getPilotCE,
@@ -33,6 +35,18 @@ def initializeHandler(cls, serviceInfoDict):
         except RuntimeError as excp:
             return S_ERROR(f"Can't connect to DB: {excp}")
 
+        cls.elasticPilotParametersDB = None
+        if Operations().getValue("/Services/JobMonitoring/useESForPilotParametersFlag", False):
+            try:
+                result = ObjectLoader().loadObject(
+                    "WorkloadManagementSystem.DB.ElasticPilotParametersDB", "ElasticPilotParametersDB"
+                )
+                if not result["OK"]:
+                    return result
+                cls.elasticPilotParametersDB = result["Value"]()
+            except RuntimeError as excp:
+                return S_ERROR(f"Can't connect to DB: {excp}")
+
         return S_OK()
 
     ##############################################################################
@@ -483,3 +497,15 @@ def export_deletePilots(cls, pilotIDs):
     @classmethod
     def export_clearPilots(cls, interval=30, aborted_interval=7):
         return cls.pilotAgentsDB.clearPilots(interval, aborted_interval)
+
+    #### ElasticPilotParameters
+
+    types_setPilotParameters = [int, str, str]
+
+    @classmethod
+    def export_setPilotParameter(cls, pilotID, key, value):
+        """Set Pilot parameters"""
+        if cls.elasticPilotParametersDB:
+            return cls.elasticPilotParametersDB.setPilotParameter(pilotID, key, value)
+
+        return S_OK()