refactor: load StorageManagementDB if present

Author: fstagni

src/DIRAC/TransformationSystem/Agent/TransformationCleaningAgent.py

@@ -70,6 +70,7 @@ def __init__(self, *args, **kwargs):
         self.jobDB = None
         self.pilotAgentsDB = None
         self.taskQueueDB = None
+        self.storageManagementDB = None
 
         # # transformations types
         self.transformationTypes = None
@@ -144,6 +145,14 @@ def initialize(self):
             return result
         self.taskQueueDB = result["Value"]()
 
+        try:
+            result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+            if not result["OK"]:
+                return result
+            self.storageManagementDB = result["Value"]()
+        except RuntimeError:
+            pass
+
         return S_OK()
 
     #############################################################################
@@ -635,6 +644,7 @@ def __removeWMSTasks(self, transJobIDs):
             jobdb=self.jobDB,
             taskqueuedb=self.taskQueueDB,
             pilotagentsdb=self.pilotAgentsDB,
+            storagemanagementdb=self.storageManagementDB,
         )
         # Prevent 0 job IDs
         jobIDs = [int(j) for j in transJobIDs if int(j)]

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -57,6 +57,7 @@ def __init__(self, *args, **kwargs):
         self.taskQueueDB = None
         self.pilotAgentsDB = None
         self.sandboxDB = None
+        self.storageManagementDB = None
 
         self.maxJobsAtOnce = 500
         self.prodTypes = []
@@ -88,6 +89,14 @@ def initialize(self):
             return result
         self.sandboxDB = result["Value"]()
 
+        try:
+            result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+            if not result["OK"]:
+                return result
+            self.storageManagementDB = result["Value"]()
+        except RuntimeError:
+            pass
+
         agentTSTypes = self.am_getOption("ProductionTypes", [])
         if agentTSTypes:
             self.prodTypes = agentTSTypes
@@ -263,6 +272,7 @@ def _deleteRemoveJobs(self, jobList, remove=False):
                     jobdb=self.jobDB,
                     taskqueuedb=self.taskQueueDB,
                     pilotagentsdb=self.pilotAgentsDB,
+                    storagemanagementdb=self.storageManagementDB,
                 )
             if not result["OK"]:
                 self.log.error(

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -41,6 +41,7 @@ def __init__(self, *args, **kwargs):
         self.logDB = None
         self.taskQueueDB = None
         self.pilotAgentsDB = None
+        self.storageManagementDB = None
         self.matchedTime = 7200
         self.rescheduledTime = 600
         self.submittingTime = 300
@@ -72,6 +73,14 @@ def initialize(self):
             return result
         self.pilotAgentsDB = result["Value"]()
 
+        try:
+            result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+            if not result["OK"]:
+                return result
+            self.storageManagementDB = result["Value"]()
+        except RuntimeError:
+            pass
+
         # getting parameters
 
         if not self.am_getOption("Enable", True):
@@ -261,6 +270,7 @@ def _failStalledJobs(self, jobID):
                 jobdb=self.jobDB,
                 taskqueuedb=self.taskQueueDB,
                 pilotagentsdb=self.pilotAgentsDB,
+                storagemanagementdb=self.storageManagementDB,
             )
             if not res["OK"]:
                 self.log.error("Failed to kill job", jobID)

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_JobCleaningAgent.py

@@ -37,11 +37,12 @@ def mock_load_object(module_path, class_name):
             "TaskQueueDB": MagicMock(),
             "PilotAgentsDB": MagicMock(),
             "SandboxMetadataDB": MagicMock(),
+            "StorageManagementDB": MagicMock(),
         }
         return {"OK": True, "Value": lambda: mocks[class_name]}
 
     mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.ObjectLoader.loadObject",
+        "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.ObjectLoader.loadObject",
         side_effect=mock_load_object,
     )
     jca = JobCleaningAgent()
@@ -147,11 +148,12 @@ def mock_load_object(module_path, class_name):
             "TaskQueueDB": MagicMock(),
             "PilotAgentsDB": MagicMock(),
             "SandboxMetadataDB": MagicMock(),
+            "StorageManagementDB": MagicMock(),
         }
         return {"OK": True, "Value": lambda: mocks[class_name]}
 
     mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.ObjectLoader.loadObject",
+        "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.ObjectLoader.loadObject",
         side_effect=mock_load_object,
     )
     jobCleaningAgent = JobCleaningAgent()

src/DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -23,6 +23,7 @@
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import (
     RIGHT_DELETE,
     RIGHT_KILL,
@@ -31,7 +32,6 @@
     RIGHT_SUBMIT,
     JobPolicy,
 )
-from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobModel import JobDescriptionModel
 from DIRAC.WorkloadManagementSystem.Utilities.ParametricJob import generateParametricJobs, getParameterVectorLength
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
@@ -45,29 +45,33 @@ class JobManagerHandlerMixin:
     @classmethod
     def initializeHandler(cls, serviceInfoDict):
         """Initialization of DB objects and OptimizationMind"""
-        try:
-            result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobDB", "JobDB")
-            if not result["OK"]:
-                return result
-            cls.jobDB = result["Value"](parentLogger=cls.log)
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobDB", "JobDB")
+        if not result["OK"]:
+            return result
+        cls.jobDB = result["Value"](parentLogger=cls.log)
 
-            result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobLoggingDB", "JobLoggingDB")
-            if not result["OK"]:
-                return result
-            cls.jobLoggingDB = result["Value"](parentLogger=cls.log)
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobLoggingDB", "JobLoggingDB")
+        if not result["OK"]:
+            return result
+        cls.jobLoggingDB = result["Value"](parentLogger=cls.log)
 
-            result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.TaskQueueDB", "TaskQueueDB")
-            if not result["OK"]:
-                return result
-            cls.taskQueueDB = result["Value"](parentLogger=cls.log)
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.TaskQueueDB", "TaskQueueDB")
+        if not result["OK"]:
+            return result
+        cls.taskQueueDB = result["Value"](parentLogger=cls.log)
 
-            result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.PilotAgentsDB", "PilotAgentsDB")
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.PilotAgentsDB", "PilotAgentsDB")
+        if not result["OK"]:
+            return result
+        cls.pilotAgentsDB = result["Value"](parentLogger=cls.log)
+
+        try:
+            result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
             if not result["OK"]:
                 return result
-            cls.pilotAgentsDB = result["Value"](parentLogger=cls.log)
-
-        except RuntimeError as excp:
-            return S_ERROR(f"Can't connect to DB: {excp!r}")
+            cls.storageManagementDB = result["Value"](parentLogger=cls.log)
+        except RuntimeError:
+            cls.storageManagementDB = None
 
         cls.msgClient = MessageClient("WorkloadManagement/OptimizationMind")
         result = cls.msgClient.connect(JobManager=True)
@@ -464,6 +468,7 @@ def export_deleteJob(self, jobIDs, force=False):
             jobdb=self.jobDB,
             taskqueuedb=self.taskQueueDB,
             pilotagentsdb=self.pilotAgentsDB,
+            storagemanagementdb=self.storageManagementDB,
         )
 
         result["requireProxyUpload"] = len(ownerJobList) > 0 and self.__checkIfProxyUploadIsRequired()
@@ -501,6 +506,7 @@ def export_killJob(self, jobIDs, force=False):
             jobdb=self.jobDB,
             taskqueuedb=self.taskQueueDB,
             pilotagentsdb=self.pilotAgentsDB,
+            storagemanagementdb=self.storageManagementDB,
         )
 
         result["requireProxyUpload"] = len(ownerJobList) > 0 and self.__checkIfProxyUploadIsRequired()