DIRAC 9.0.9__py3-none-any.whl → 9.0.11__py3-none-any.whl

DIRAC/TransformationSystem/DB/TransformationDB.py

@@ -6,6 +6,9 @@ This class is typically used as a base class for more specific data processing
 databases
 """
 
+# Disable it because pylint does not understand decorator (convertToReturnValue)
+
+# pylint: disable=invalid-sequence-index
 import re
 import time
 import threading
@@ -15,6 +18,7 @@ from errno import ENOENT
 from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Utilities.DErrno import cmpError
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue, returnValueOrRaise
 from DIRAC.Resources.Catalog.FileCatalog import FileCatalog
 from DIRAC.Core.Security.ProxyInfo import getProxyInfo
 from DIRAC.Core.Utilities.List import stringListToString, intListToString, breakListIntoChunks
@@ -25,6 +29,7 @@ from DIRAC.DataManagementSystem.Client.MetaQuery import MetaQuery
 
 MAX_ERROR_COUNT = 10
 
+TMP_TABLE_JOIN_LIMIT = 100
 #############################################################################
 
 
@@ -270,6 +275,7 @@ class TransformationDB(DB):
         self.__updateTransformationLogging(transID, message, author, connection=connection)
         return S_OK(transID)
 
+    @convertToReturnValue
     def getTransformations(
         self,
         condDict=None,
@@ -289,32 +295,54 @@ class TransformationDB(DB):
             columns = self.TRANSPARAMS
         else:
             columns = [c for c in columns if c in self.TRANSPARAMS]
-        req = "SELECT {} FROM Transformations {}".format(
-            intListToString(columns),
-            self.buildCondition(condDict, older, newer, timeStamp, orderAttribute, limit, offset=offset),
-        )
-        res = self._query(req, conn=connection)
-        if not res["OK"]:
-            return res
-        if condDict is None:
-            condDict = {}
-        webList = []
+
+        join_query = ""
+
+        try:
+            # If we request multiple TransformationIDs, and they are more than TMP_TABLE_JOIN_LIMIT,
+            # we create a temporary table to speed up the query
+            if (
+                "TransformationID" in condDict
+                and isinstance(condDict["TransformationID"], list)
+                and len(condDict["TransformationID"]) > TMP_TABLE_JOIN_LIMIT
+            ):
+                # Create temporary table for TransformationIDs
+                transIDs = condDict.pop("TransformationID")
+                sqlCmd = "CREATE TEMPORARY TABLE to_query_TransformationIDs (TransID INTEGER NOT NULL, PRIMARY KEY (TransID)) ENGINE=MEMORY;"
+                returnValueOrRaise(self._update(sqlCmd, conn=connection))
+                join_query = " JOIN to_query_TransformationIDs t ON TransformationID = t.TransID"
+
+                # Insert TransformationIDs into temporary table
+                sqlCmd = "INSERT INTO to_query_TransformationIDs (TransID) VALUES ( %s )"
+                returnValueOrRaise(self._updatemany(sqlCmd, [(transID,) for transID in transIDs], conn=connection))
+
+            req = "SELECT {} FROM Transformations {} {}".format(
+                intListToString(columns),
+                join_query,
+                self.buildCondition(condDict, older, newer, timeStamp, orderAttribute, limit, offset=offset),
+            )
+            matching_transformations = returnValueOrRaise(self._query(req, conn=connection))
+
+        finally:
+            # Clean up temporary table
+            if join_query:
+                sqlCmd = "DROP TEMPORARY TABLE to_query_TransformationIDs"
+                self._update(sqlCmd, conn=connection)
+
+        # TODO: optimize by getting all the extra params at once
         resultList = []
-        for row in res["Value"]:
+        for row in matching_transformations:
             # Prepare the structure for the web
-            rList = [str(item) if not isinstance(item, int) else item for item in row]
             transDict = dict(zip(columns, row))
-            webList.append(rList)
             if extraParams:
-                res = self.__getAdditionalParameters(transDict["TransformationID"], connection=connection)
-                if not res["OK"]:
-                    return res
-                transDict.update(res["Value"])
+                trans_extra_param = returnValueOrRaise(
+                    self.__getAdditionalParameters(transDict["TransformationID"], connection=connection)
+                )
+
+                transDict.update(trans_extra_param)
             resultList.append(transDict)
-        result = S_OK(resultList)
-        result["Records"] = webList
-        result["ParameterNames"] = columns
-        return result
+
+        return resultList
 
     def getTransformation(self, transName, extraParams=False, connection=False):
         """Get Transformation definition and parameters of Transformation identified by TransformationID"""
@@ -710,21 +738,38 @@ class TransformationDB(DB):
         countDict["Total"] = sum(countDict.values())
         return S_OK(countDict)
 
+    @convertToReturnValue
     def __addFilesToTransformation(self, transID, fileIDs, connection=False):
-        req = "SELECT FileID from TransformationFiles"
-        req = req + " WHERE TransformationID = %d AND FileID IN (%s);" % (transID, intListToString(fileIDs))
-        res = self._query(req, conn=connection)
-        if not res["OK"]:
-            return res
-        for tupleIn in res["Value"]:
-            fileIDs.remove(tupleIn[0])
-        if not fileIDs:
-            return S_OK([])
-        values = [(transID, fileID) for fileID in fileIDs]
-        req = "INSERT INTO TransformationFiles (TransformationID,FileID,LastUpdate,InsertedTime) VALUES (%s, %s, UTC_TIMESTAMP(), UTC_TIMESTAMP())"
-        if not (res := self._updatemany(req, values, conn=connection))["OK"]:
-            return res
-        return S_OK(fileIDs)
+        # Create temporary table for FileIDs
+        sqlCmd = "CREATE TEMPORARY TABLE to_query_FileIDs (FileID INT(11) UNSIGNED NOT NULL, PRIMARY KEY (FileID)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd, conn=connection))
+
+        try:
+            # Insert FileIDs into temporary table
+            sqlCmd = "INSERT INTO to_query_FileIDs (FileID) VALUES ( %s )"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(fileID,) for fileID in fileIDs], conn=connection))
+
+            # Query existing files using JOIN
+            req = (
+                "SELECT tf.FileID FROM TransformationFiles tf JOIN to_query_FileIDs t ON tf.FileID = t.FileID WHERE tf.TransformationID = %d;"
+                % transID
+            )
+            res = returnValueOrRaise(self._query(req, conn=connection))
+
+            # Remove already existing fileIDs using set difference for efficiency
+            existingFileIDs = {tupleIn[0] for tupleIn in res}
+            fileIDs = list(set(fileIDs) - existingFileIDs)
+            if not fileIDs:
+                return []
+
+            values = [(transID, fileID) for fileID in fileIDs]
+            req = "INSERT INTO TransformationFiles (TransformationID,FileID,LastUpdate,InsertedTime) VALUES (%s, %s, UTC_TIMESTAMP(), UTC_TIMESTAMP())"
+            returnValueOrRaise(self._updatemany(req, values, conn=connection))
+            return fileIDs
+        finally:
+            # Clean up temporary table
+            sqlCmd = "DROP TEMPORARY TABLE to_query_FileIDs"
+            returnValueOrRaise(self._update(sqlCmd, conn=connection))
 
     def __insertExistingTransformationFiles(self, transID, fileTuplesList, connection=False):
         """Inserting already transformation files in TransformationFiles table (e.g. for deriving transformations)"""
@@ -1271,18 +1316,35 @@ class TransformationDB(DB):
     # These methods manipulate the DataFiles table
     #
 
+    @convertToReturnValue
     def __getFileIDsForLfns(self, lfns, connection=False):
         """Get file IDs for the given list of lfns
         warning: if the file is not present, we'll see no errors
         """
-        req = f"SELECT LFN,FileID FROM DataFiles WHERE LFN in ({stringListToString(lfns)});"
-        res = self._query(req, conn=connection)
-        if not res["OK"]:
-            return res
-        lfns = dict(res["Value"])
-        # Reverse dictionary
-        fids = {fileID: lfn for lfn, fileID in lfns.items()}
-        return S_OK((fids, lfns))
+
+        if not lfns:
+            return ({}, {})
+        # Create temporary table for LFNs
+        sqlCmd = "CREATE TEMPORARY TABLE to_query_LFNs (LFN VARCHAR(255) NOT NULL, PRIMARY KEY (LFN)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd, conn=connection))
+
+        try:
+            # Insert LFNs into temporary table
+            sqlCmd = "INSERT INTO to_query_LFNs (LFN) VALUES ( %s )"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(lfn,) for lfn in lfns], conn=connection))
+
+            # Query using JOIN with temporary table
+            req = "SELECT df.LFN, df.FileID FROM DataFiles df JOIN to_query_LFNs t ON df.LFN = t.LFN;"
+            res = returnValueOrRaise(self._query(req, conn=connection))
+
+            lfns = dict(res)
+            # Reverse dictionary
+            fids = {fileID: lfn for lfn, fileID in lfns.items()}
+            return (fids, lfns)
+        finally:
+            # Clean up temporary table
+            sqlCmd = "DROP TEMPORARY TABLE to_query_LFNs"
+            self._update(sqlCmd, conn=connection)
 
     def __getLfnsForFileIDs(self, fileIDs, connection=False):
         """Get lfns for the given list of fileIDs"""

DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -17,11 +17,9 @@ from DIRAC.ConfigurationSystem.Client.Helpers import cfgPath
 from DIRAC.Core.Base.AgentModule import AgentModule
 from DIRAC.Core.Utilities import DErrno
 from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
+from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.Core.Utilities.TimeUtilities import fromString, second, toEpoch
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
-from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
-from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_KILL
 from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
@@ -40,6 +38,9 @@ class StalledJobAgent(AgentModule):
 
         self.jobDB = None
         self.logDB = None
+        self.taskQueueDB = None
+        self.pilotAgentsDB = None
+        self.storageManagementDB = None
         self.matchedTime = 7200
         self.rescheduledTime = 600
         self.submittingTime = 300
@@ -51,8 +52,30 @@ class StalledJobAgent(AgentModule):
     #############################################################################
     def initialize(self):
         """Sets default parameters."""
-        self.jobDB = JobDB()
-        self.logDB = JobLoggingDB()
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobDB", "JobDB")
+        if not result["OK"]:
+            return result
+        self.jobDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobLoggingDB", "JobLoggingDB")
+        if not result["OK"]:
+            return result
+        self.logDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.TaskQueueDB", "TaskQueueDB")
+        if not result["OK"]:
+            return result
+        self.taskQueueDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.PilotAgentsDB", "PilotAgentsDB")
+        if not result["OK"]:
+            return result
+        self.pilotAgentsDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+        if not result["OK"]:
+            return result
+        self.storageManagementDB = result["Value"]()
 
         # getting parameters
 
@@ -235,7 +258,16 @@ class StalledJobAgent(AgentModule):
         # Set the jobs Failed, send them a kill signal in case they are not really dead
         # and send accounting info
         if setFailed:
-            res = kill_delete_jobs(RIGHT_KILL, [jobID], nonauthJobList=[], force=True)
+            res = kill_delete_jobs(
+                RIGHT_KILL,
+                [jobID],
+                nonauthJobList=[],
+                force=True,
+                jobdb=self.jobDB,
+                taskqueuedb=self.taskQueueDB,
+                pilotagentsdb=self.pilotAgentsDB,
+                storagemanagementdb=self.storageManagementDB,
+            )
             if not res["OK"]:
                 self.log.error("Failed to kill job", jobID)
 
@@ -262,7 +294,7 @@ class StalledJobAgent(AgentModule):
             # There is no pilot reference, hence its status is unknown
             return S_OK("NoPilot")
 
-        result = PilotAgentsDB().getPilotInfo(pilotReference)
+        result = self.pilotAgentsDB.getPilotInfo(pilotReference)
         if not result["OK"]:
             if DErrno.cmpError(result, DErrno.EWMSNOPILOT):
                 self.log.warn("No pilot found", f"for job {jobID}: {result['Message']}")

DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_StalledJobAgent.py

@@ -23,10 +23,31 @@ def sja(mocker):
         side_effect=lambda x, y=None: y,
         create=True,
     )
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobDB")
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobLoggingDB")
+
+    # Mock ObjectLoader to return mock DB instances
+    mockJobDB = MagicMock()
+    mockJobDB.log = gLogger
+    mockJobLoggingDB = MagicMock()
+    mockTaskQueueDB = MagicMock()
+    mockPilotAgentsDB = MagicMock()
+    mockStorageManagementDB = MagicMock()
+
+    def mock_load_object(module_path, class_name):
+        mocks = {
+            "JobDB": mockJobDB,
+            "JobLoggingDB": mockJobLoggingDB,
+            "TaskQueueDB": mockTaskQueueDB,
+            "PilotAgentsDB": mockPilotAgentsDB,
+            "StorageManagementDB": mockStorageManagementDB,
+        }
+        return {"OK": True, "Value": lambda: mocks[class_name]}
+
+    mocker.patch(
+        "DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.ObjectLoader.loadObject",
+        side_effect=mock_load_object,
+    )
+
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.rescheduleJobs", return_value=MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.PilotAgentsDB", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getJobParameters", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.kill_delete_jobs", return_value=MagicMock())
 
@@ -34,7 +55,6 @@ def sja(mocker):
     stalledJobAgent._AgentModule__configDefaults = mockAM
     stalledJobAgent.log = gLogger
     stalledJobAgent.initialize()
-    stalledJobAgent.jobDB.log = gLogger
     stalledJobAgent.log.setLevel("DEBUG")
     stalledJobAgent.stalledTime = 120
 

DIRAC/WorkloadManagementSystem/DB/StatusUtils.py

@@ -1,43 +1,40 @@
 from DIRAC import S_ERROR, S_OK, gLogger
-from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
-from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
+from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
 from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import _filterJobStateTransition
 
 
-def _deleteJob(jobID, force=False):
+def _deleteJob(jobID, force=False, *, jobdb, taskqueuedb, pilotagentsdb):
     """Set the job status to "Deleted"
     and remove the pilot that ran and its logging info if the pilot is finished.
 
     :param int jobID: job ID
     :return: S_OK()/S_ERROR()
     """
-    if not (result := JobDB().setJobStatus(jobID, JobStatus.DELETED, "Checking accounting", force=force))["OK"]:
+    if not (result := jobdb.setJobStatus(jobID, JobStatus.DELETED, "Checking accounting", force=force))["OK"]:
         gLogger.warn("Failed to set job Deleted status", result["Message"])
         return result
 
-    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+    if not (result := taskqueuedb.deleteJob(jobID))["OK"]:
         gLogger.warn("Failed to delete job from the TaskQueue")
 
     # if it was the last job for the pilot
-    result = PilotAgentsDB().getPilotsForJobID(jobID)
+    result = pilotagentsdb.getPilotsForJobID(jobID)
     if not result["OK"]:
         gLogger.error("Failed to get Pilots for JobID", result["Message"])
         return result
     for pilot in result["Value"]:
-        res = PilotAgentsDB().getJobsForPilot(pilot)
+        res = pilotagentsdb.getJobsForPilot(pilot)
         if not res["OK"]:
             gLogger.error("Failed to get jobs for pilot", res["Message"])
             return res
         if not res["Value"]:  # if list of jobs for pilot is empty, delete pilot
-            result = PilotAgentsDB().getPilotInfo(pilotID=pilot)
+            result = pilotagentsdb.getPilotInfo(pilotID=pilot)
             if not result["OK"]:
                 gLogger.error("Failed to get pilot info", result["Message"])
                 return result
-            ret = PilotAgentsDB().deletePilot(result["Value"]["PilotJobReference"])
+            ret = pilotagentsdb.deletePilot(result["Value"]["PilotJobReference"])
             if not ret["OK"]:
                 gLogger.error("Failed to delete pilot from PilotAgentsDB", ret["Message"])
                 return ret
@@ -45,7 +42,7 @@ def _deleteJob(jobID, force=False):
     return S_OK()
 
 
-def _killJob(jobID, sendKillCommand=True, force=False):
+def _killJob(jobID, sendKillCommand=True, force=False, *, jobdb, taskqueuedb):
     """Kill one job
 
     :param int jobID: job ID
@@ -54,32 +51,63 @@ def _killJob(jobID, sendKillCommand=True, force=False):
     :return: S_OK()/S_ERROR()
     """
     if sendKillCommand:
-        if not (result := JobDB().setJobCommand(jobID, "Kill"))["OK"]:
+        if not (result := jobdb.setJobCommand(jobID, "Kill"))["OK"]:
             gLogger.warn("Failed to set job Kill command", result["Message"])
             return result
 
     gLogger.info("Job marked for termination", jobID)
-    if not (result := JobDB().setJobStatus(jobID, JobStatus.KILLED, "Marked for termination", force=force))["OK"]:
+    if not (result := jobdb.setJobStatus(jobID, JobStatus.KILLED, "Marked for termination", force=force))["OK"]:
         gLogger.warn("Failed to set job Killed status", result["Message"])
-    if not (result := TaskQueueDB().deleteJob(jobID))["OK"]:
+    if not (result := taskqueuedb.deleteJob(jobID))["OK"]:
         gLogger.warn("Failed to delete job from the TaskQueue", result["Message"])
 
     return S_OK()
 
 
-def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
+def kill_delete_jobs(
+    right,
+    validJobList,
+    nonauthJobList=[],
+    force=False,
+    *,
+    jobdb=None,
+    taskqueuedb=None,
+    pilotagentsdb=None,
+    storagemanagementdb=None,
+):
     """Kill (== set the status to "KILLED") or delete (== set the status to "DELETED") jobs as necessary
 
     :param str right: RIGHT_KILL or RIGHT_DELETE
 
     :return: S_OK()/S_ERROR()
     """
+    if jobdb is None:
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobDB", "JobDB")
+        if not result["OK"]:
+            return result
+        jobdb = result["Value"]()
+    if taskqueuedb is None:
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.TaskQueueDB", "TaskQueueDB")
+        if not result["OK"]:
+            return result
+        taskqueuedb = result["Value"]()
+    if pilotagentsdb is None:
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.PilotAgentsDB", "PilotAgentsDB")
+        if not result["OK"]:
+            return result
+        pilotagentsdb = result["Value"]()
+    if storagemanagementdb is None:
+        result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+        if not result["OK"]:
+            return result
+        storagemanagementdb = result["Value"]()
+
     badIDs = []
 
     killJobList = []
     deleteJobList = []
     if validJobList:
-        result = JobDB().getJobsAttributes(validJobList, ["Status"])
+        result = jobdb.getJobsAttributes(validJobList, ["Status"])
         if not result["OK"]:
             return result
         jobStates = result["Value"]
@@ -92,12 +120,12 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
             deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
 
         for jobID in killJobList:
-            result = _killJob(jobID, force=force)
+            result = _killJob(jobID, force=force, jobdb=jobdb, taskqueuedb=taskqueuedb)
             if not result["OK"]:
                 badIDs.append(jobID)
 
         for jobID in deleteJobList:
-            result = _deleteJob(jobID, force=force)
+            result = _deleteJob(jobID, force=force, jobdb=jobdb, taskqueuedb=taskqueuedb, pilotagentsdb=pilotagentsdb)
             if not result["OK"]:
                 badIDs.append(jobID)
 
@@ -105,9 +133,8 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
         stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
 
         if stagingJobList:
-            stagerDB = StorageManagementDB()
             gLogger.info("Going to send killing signal to stager as well!")
-            result = stagerDB.killTasksBySourceTaskID(stagingJobList)
+            result = storagemanagementdb.killTasksBySourceTaskID(stagingJobList)
             if not result["OK"]:
                 gLogger.warn("Failed to kill some Stager tasks", result["Message"])
 

DIRAC/WorkloadManagementSystem/DB/tests/Test_StatusUtils.py

@@ -19,10 +19,25 @@ from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
     ],
 )
 def test___kill_delete_jobs(mocker, jobIDs_list, right):
-    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.JobDB", MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.TaskQueueDB", MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.PilotAgentsDB", MagicMock())
-    mocker.patch("DIRAC.WorkloadManagementSystem.DB.StatusUtils.StorageManagementDB", MagicMock())
+    # Mock ObjectLoader to return mock DB instances
+    mockJobDB = MagicMock()
+    mockTaskQueueDB = MagicMock()
+    mockPilotAgentsDB = MagicMock()
+    mockStorageManagementDB = MagicMock()
+
+    def mock_load_object(module_path, class_name):
+        mocks = {
+            "JobDB": mockJobDB,
+            "TaskQueueDB": mockTaskQueueDB,
+            "PilotAgentsDB": mockPilotAgentsDB,
+            "StorageManagementDB": mockStorageManagementDB,
+        }
+        return {"OK": True, "Value": lambda: mocks[class_name]}
+
+    mocker.patch(
+        "DIRAC.WorkloadManagementSystem.DB.StatusUtils.ObjectLoader.loadObject",
+        side_effect=mock_load_object,
+    )
 
     res = kill_delete_jobs(right, jobIDs_list)
     assert res["OK"]

DIRAC/WorkloadManagementSystem/Service/JobManagerHandler.py

@@ -65,6 +65,11 @@ class JobManagerHandlerMixin:
                 return result
             cls.pilotAgentsDB = result["Value"](parentLogger=cls.log)
 
+            result = ObjectLoader().loadObject("StorageManagementSystem.DB.StorageManagementDB", "StorageManagementDB")
+            if not result["OK"]:
+                return result
+            cls.storageManagementDB = result["Value"](parentLogger=cls.log)
+
         except RuntimeError as excp:
             return S_ERROR(f"Can't connect to DB: {excp!r}")
 
@@ -449,7 +454,16 @@ class JobManagerHandlerMixin:
             jobList, RIGHT_DELETE
         )
 
-        result = kill_delete_jobs(RIGHT_DELETE, validJobList, nonauthJobList, force=force)
+        result = kill_delete_jobs(
+            RIGHT_DELETE,
+            validJobList,
+            nonauthJobList,
+            force=force,
+            jobdb=self.jobDB,
+            taskqueuedb=self.taskQueueDB,
+            pilotagentsdb=self.pilotAgentsDB,
+            storagemanagementdb=self.storageManagementDB,
+        )
 
         result["requireProxyUpload"] = len(ownerJobList) > 0 and self.__checkIfProxyUploadIsRequired()
 
@@ -478,7 +492,16 @@ class JobManagerHandlerMixin:
             jobList, RIGHT_KILL
         )
 
-        result = kill_delete_jobs(RIGHT_KILL, validJobList, nonauthJobList, force=force)
+        result = kill_delete_jobs(
+            RIGHT_KILL,
+            validJobList,
+            nonauthJobList,
+            force=force,
+            jobdb=self.jobDB,
+            taskqueuedb=self.taskQueueDB,
+            pilotagentsdb=self.pilotAgentsDB,
+            storagemanagementdb=self.storageManagementDB,
+        )
 
         result["requireProxyUpload"] = len(ownerJobList) > 0 and self.__checkIfProxyUploadIsRequired()
 

{dirac-9.0.9.dist-info → dirac-9.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DIRAC
-Version: 9.0.9
+Version: 9.0.11
 Summary: DIRAC is an interware, meaning a software framework for distributed computing.
 Home-page: https://github.com/DIRACGrid/DIRAC/
 License: GPL-3.0-only
@@ -19,7 +19,7 @@ Requires-Dist: cachetools
 Requires-Dist: certifi
 Requires-Dist: cwltool
 Requires-Dist: diraccfg
-Requires-Dist: DIRACCommon==v9.0.9
+Requires-Dist: DIRACCommon==v9.0.11
 Requires-Dist: diracx-client>=v0.0.1
 Requires-Dist: diracx-core>=v0.0.1
 Requires-Dist: diracx-cli>=v0.0.1