feat: PilotManager can interact with ElasticPilotParameters

DIRACGrid · Apr 23, 2024 · efa723a · efa723a
1 parent 764ae22
commit efa723a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py b/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py
@@ -5,12 +5,11 @@
       - setPilotParameter()
       - deletePilotParameters()
 """
-from DIRAC import S_OK, S_ERROR, gConfig
-from DIRAC.Core.Utilities import TimeUtilities
-from DIRAC.ConfigurationSystem.Client.PathFinder import getDatabaseSection
+from DIRAC import S_ERROR, S_OK, gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers import CSGlobals
+from DIRAC.ConfigurationSystem.Client.PathFinder import getDatabaseSection
 from DIRAC.Core.Base.ElasticDB import ElasticDB
-
+from DIRAC.Core.Utilities import TimeUtilities
 
 mapping = {
     "properties": {

diff --git a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py
@@ -36,6 +36,17 @@ def initializeHandler(cls, serviceInfoDict):
         defaultOption, defaultClass = "DownloadPlugin", "FileCacheDownloadPlugin"
         cls.configValue = getServiceOption(serviceInfoDict, defaultOption, defaultClass)
         cls.loggingPlugin = None
+        cls.elasticPilotParametersDB = None
+        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()
 
     ##############################################################################
@@ -479,3 +490,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()
diff --git a/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py b/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py
@@ -150,9 +150,6 @@ def test_setAndGetPilotFromDB():
     assert len(res["Value"][1010000]) == 0
 
     # delete the indexes
-    res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB.indexName_base)
-    assert res["OK"]
-    assert res["Value"] == "Nothing to delete"
     res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB._indexName(100))
     assert res["OK"]
     res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB._indexName(1010000))