DIRAC 9.0.0a54__py3-none-any.whl → 9.0.7__py3-none-any.whl

DIRAC/WorkloadManagementSystem/DB/JobDB.py

@@ -11,8 +11,11 @@ The following options can be set in ``Systems/WorkloadManagement/Databases/JobDB
 * *CompressJDLs*:        Enable compression of JDLs when they are stored in the database, default *False*.
 
 """
+from __future__ import annotations
+
 import datetime
 import operator
+from typing import overload
 
 from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getVOForGroup
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getSiteTier
@@ -20,7 +23,14 @@ from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
 from DIRAC.Core.Utilities.Decorators import deprecated
 from DIRAC.Core.Utilities.DErrno import EWMSJMAN, EWMSSUBM, cmpError
-from DIRAC.Core.Utilities.ReturnValues import S_ERROR, S_OK, convertToReturnValue, returnValueOrRaise, SErrorException
+from DIRAC.Core.Utilities.ReturnValues import (
+    S_ERROR,
+    S_OK,
+    convertToReturnValue,
+    returnValueOrRaise,
+    SErrorException,
+    DReturnType,
+)
 from DIRAC.FrameworkSystem.Client.Logger import contextLogger
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
@@ -320,23 +330,42 @@ class JobDB(DB):
 
     #############################################################################
 
-    def getInputData(self, jobID):
+    @overload
+    def getInputData(self, jobID: int | str) -> DReturnType[list[str]]:
+        ...
+
+    @overload
+    def getInputData(self, jobID: list[int | str]) -> DReturnType[dict[int, list[str]]]:
+        ...
+
+    def getInputData(self, jobID: int | str | list[int | str]) -> DReturnType[list[str] | dict[int, list[str]]]:
         """Get input data for the given job"""
-        ret = self._escapeString(jobID)
-        if not ret["OK"]:
-            return ret
-        jobID = ret["Value"]
-        cmd = f"SELECT LFN FROM InputData WHERE JobID={jobID}"
+        if isinstance(jobID, (int, str)):
+            ret = self._escapeString(jobID)
+            if not ret["OK"]:
+                return ret
+            jobID = ret["Value"]
+            query = f"JobID={jobID}"
+            result = []
+        else:
+            job_ids = {int(i) for i in jobID}
+            query = f"JobID IN ({','.join(map(str, job_ids))})"
+            result = {i: [] for i in job_ids}
+        cmd = f"SELECT JobID, LFN FROM InputData WHERE {query}"
         res = self._query(cmd)
         if not res["OK"]:
             return res
 
-        inputData = [i[0] for i in res["Value"] if i[0].strip()]
-        for index, lfn in enumerate(inputData):
+        for jid, lfn in res["Value"]:
+            lfn = lfn.strip()
             if lfn.lower().startswith("lfn:"):
-                inputData[index] = lfn[4:]
+                lfn = lfn[4:]
+            if isinstance(result, list):
+                result.append(lfn)
+            else:
+                result[jid].append(lfn)
 
-        return S_OK(inputData)
+        return S_OK(result)
 
     #############################################################################
     def setInputData(self, jobID, inputData):
@@ -1161,64 +1190,6 @@ class JobDB(DB):
 
         return retVal
 
-    #############################################################################
-    def getSiteSummary(self):
-        """Get the summary of jobs in a given status on all the sites"""
-
-        waitingList = ['"Submitted"', '"Assigned"', '"Waiting"', '"Matched"']
-        waitingString = ",".join(waitingList)
-
-        result = self.getDistinctJobAttributes("Site")
-        if not result["OK"]:
-            return result
-
-        siteList = result["Value"]
-        siteDict = {}
-        totalDict = {
-            JobStatus.WAITING: 0,
-            JobStatus.RUNNING: 0,
-            JobStatus.STALLED: 0,
-            JobStatus.DONE: 0,
-            JobStatus.FAILED: 0,
-        }
-
-        for site in siteList:
-            if site == "ANY":
-                continue
-            # Waiting
-            siteDict[site] = {}
-            ret = self._escapeString(site)
-            if not ret["OK"]:
-                return ret
-            e_site = ret["Value"]
-
-            req = f"SELECT COUNT(JobID) FROM Jobs WHERE Status IN ({waitingString}) AND Site={e_site}"
-            result = self._query(req)
-            if result["OK"]:
-                count = result["Value"][0][0]
-            else:
-                return S_ERROR("Failed to get Site data from the JobDB")
-            siteDict[site][JobStatus.WAITING] = count
-            totalDict[JobStatus.WAITING] += count
-            # Running,Stalled,Done,Failed
-            for status in [
-                f'"{JobStatus.RUNNING}"',
-                f'"{JobStatus.STALLED}"',
-                f'"{JobStatus.DONE}"',
-                f'"{JobStatus.FAILED}"',
-            ]:
-                req = f"SELECT COUNT(JobID) FROM Jobs WHERE Status={status} AND Site={e_site}"
-                result = self._query(req)
-                if result["OK"]:
-                    count = result["Value"][0][0]
-                else:
-                    return S_ERROR("Failed to get Site data from the JobDB")
-                siteDict[site][status.replace('"', "")] = count
-                totalDict[status.replace('"', "")] += count
-
-        siteDict["Total"] = totalDict
-        return S_OK(siteDict)
-
     #################################################################################
     def getSiteSummaryWeb(self, selectDict, sortList, startItem, maxItems):
         """Get the summary of jobs in a given status on all the sites in the standard Web form"""

DIRAC/WorkloadManagementSystem/DB/JobDBUtils.py

@@ -1,162 +1,33 @@
 from __future__ import annotations
 
-import base64
-import zlib
+# Import stateless functions from DIRACCommon for backward compatibility
+from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import *
 
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
-from DIRAC.Core.Utilities.DErrno import EWMSSUBM
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
-from DIRAC.Core.Utilities.ReturnValues import S_ERROR, S_OK, returnValueOrRaise
-from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.Client.JobState.JobManifest import JobManifest
+from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
+from DIRAC.WorkloadManagementSystem.Client.JobState.JobManifest import makeJobManifestConfig
 
 getDIRACPlatform = returnValueOrRaise(
     ObjectLoader().loadObject("ConfigurationSystem.Client.Helpers.Resources", "getDIRACPlatform")
 )
 
 
-def compressJDL(jdl):
-    """Return compressed JDL string."""
-    return base64.b64encode(zlib.compress(jdl.encode(), -1)).decode()
+def checkAndPrepareJob(
+    jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo
+):  # pylint: disable=function-redefined
+    from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import checkAndPrepareJob
 
+    config = {
+        "inputDataPolicyForVO": Operations(vo=vo).getValue("InputDataPolicy/InputDataModule"),
+        "softwareDistModuleForVO": Operations(vo=vo).getValue("SoftwareDistModule"),
+        "defaultCPUTimeForOwnerGroup": Operations(group=ownerGroup).getValue("JobDescription/DefaultCPUTime", 86400),
+        "getDIRACPlatform": getDIRACPlatform,
+    }
+    return checkAndPrepareJob(jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo, config=config)
 
-def extractJDL(compressedJDL):
-    """Return decompressed JDL string."""
-    # the starting bracket is guaranteeed by JobManager.submitJob
-    # we need the check to be backward compatible
-    if isinstance(compressedJDL, bytes):
-        if compressedJDL.startswith(b"["):
-            return compressedJDL.decode()
-    else:
-        if compressedJDL.startswith("["):
-            return compressedJDL
-    return zlib.decompress(base64.b64decode(compressedJDL)).decode()
 
+def checkAndAddOwner(jdl: str, owner: str, ownerGroup: str):  # pylint: disable=function-redefined
+    from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import checkAndAddOwner
 
-def checkAndAddOwner(jdl: str, owner: str, ownerGroup: str) -> JobManifest:
-    jobManifest = JobManifest()
-    res = jobManifest.load(jdl)
-    if not res["OK"]:
-        return res
-
-    jobManifest.setOptionsFromDict({"Owner": owner, "OwnerGroup": ownerGroup})
-    res = jobManifest.check()
-    if not res["OK"]:
-        return res
-
-    return S_OK(jobManifest)
-
-
-def fixJDL(jdl: str) -> str:
-    # 1.- insert original JDL on DB and get new JobID
-    # Fix the possible lack of the brackets in the JDL
-    if jdl.strip()[0].find("[") != 0:
-        jdl = "[" + jdl + "]"
-    return jdl
-
-
-def checkAndPrepareJob(jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo):
-    error = ""
-
-    jdlOwner = classAdJob.getAttributeString("Owner")
-    jdlOwnerGroup = classAdJob.getAttributeString("OwnerGroup")
-    jdlVO = classAdJob.getAttributeString("VirtualOrganization")
-
-    # The below is commented out since this is always overwritten by the submitter IDs
-    # but the check allows to findout inconsistent client environments
-    if jdlOwner and jdlOwner != owner:
-        error = "Wrong Owner in JDL"
-    elif jdlOwnerGroup and jdlOwnerGroup != ownerGroup:
-        error = "Wrong Owner Group in JDL"
-    elif jdlVO and jdlVO != vo:
-        error = "Wrong Virtual Organization in JDL"
-
-    classAdJob.insertAttributeString("Owner", owner)
-    classAdJob.insertAttributeString("OwnerGroup", ownerGroup)
-
-    if vo:
-        classAdJob.insertAttributeString("VirtualOrganization", vo)
-
-    classAdReq.insertAttributeString("Owner", owner)
-    classAdReq.insertAttributeString("OwnerGroup", ownerGroup)
-    if vo:
-        classAdReq.insertAttributeString("VirtualOrganization", vo)
-
-    inputDataPolicy = Operations(vo=vo).getValue("InputDataPolicy/InputDataModule")
-    if inputDataPolicy and not classAdJob.lookupAttribute("InputDataModule"):
-        classAdJob.insertAttributeString("InputDataModule", inputDataPolicy)
-
-    softwareDistModule = Operations(vo=vo).getValue("SoftwareDistModule")
-    if softwareDistModule and not classAdJob.lookupAttribute("SoftwareDistModule"):
-        classAdJob.insertAttributeString("SoftwareDistModule", softwareDistModule)
-
-    # priority
-    priority = classAdJob.getAttributeInt("Priority")
-    if priority is None:
-        priority = 0
-    classAdReq.insertAttributeInt("UserPriority", priority)
-
-    # CPU time
-    cpuTime = classAdJob.getAttributeInt("CPUTime")
-    if cpuTime is None:
-        opsHelper = Operations(group=ownerGroup)
-        cpuTime = opsHelper.getValue("JobDescription/DefaultCPUTime", 86400)
-    classAdReq.insertAttributeInt("CPUTime", cpuTime)
-
-    # platform(s)
-    platformList = classAdJob.getListFromExpression("Platform")
-    if platformList:
-        result = getDIRACPlatform(platformList)
-        if not result["OK"]:
-            return result
-        if result["Value"]:
-            classAdReq.insertAttributeVectorString("Platforms", result["Value"])
-        else:
-            error = "OS compatibility info not found"
-    if error:
-        retVal = S_ERROR(EWMSSUBM, error)
-        retVal["JobId"] = jobID
-        retVal["Status"] = JobStatus.FAILED
-        retVal["MinorStatus"] = error
-
-        jobAttrs["Status"] = JobStatus.FAILED
-
-        jobAttrs["MinorStatus"] = error
-        return retVal
-    return S_OK()
-
-
-def createJDLWithInitialStatus(
-    classAdJob, classAdReq, jdl2DBParameters, jobAttrs, initialStatus, initialMinorStatus, *, modern=False
-):
-    """
-    :param modern: if True, store boolean instead of string for VerifiedFlag (used by diracx only)
-    """
-    priority = classAdJob.getAttributeInt("Priority")
-    if priority is None:
-        priority = 0
-    jobAttrs["UserPriority"] = priority
-
-    for jdlName in jdl2DBParameters:
-        # Defaults are set by the DB.
-        jdlValue = classAdJob.getAttributeString(jdlName)
-        if jdlValue:
-            jobAttrs[jdlName] = jdlValue
-
-    jdlValue = classAdJob.getAttributeString("Site")
-    if jdlValue:
-        if jdlValue.find(",") != -1:
-            jobAttrs["Site"] = "Multiple"
-        else:
-            jobAttrs["Site"] = jdlValue
-
-    jobAttrs["VerifiedFlag"] = True if modern else "True"
-
-    jobAttrs["Status"] = initialStatus
-
-    jobAttrs["MinorStatus"] = initialMinorStatus
-
-    reqJDL = classAdReq.asJDL()
-    classAdJob.insertAttributeInt("JobRequirements", reqJDL)
-
-    return classAdJob.asJDL()
+    return checkAndAddOwner(jdl, owner, ownerGroup, job_manifest_config=makeJobManifestConfig(ownerGroup))

DIRAC/WorkloadManagementSystem/DB/JobParametersDB.py

@@ -1,10 +1,10 @@
-""" Module containing a front-end to the OpenSearch-based JobParametersDB.
-    This is a drop-in replacement for MySQL-based table JobDB.JobParameters.
+"""Module containing a front-end to the OpenSearch-based JobParametersDB.
+This is a drop-in replacement for MySQL-based table JobDB.JobParameters.
 
-    The following class methods are provided for public usage
-      - getJobParameters()
-      - setJobParameter()
-      - deleteJobParameters()
+The following class methods are provided for public usage
+  - getJobParameters()
+  - setJobParameter()
+  - deleteJobParameters()
 """
 
 from DIRAC import S_ERROR, S_OK
@@ -37,11 +37,11 @@ class JobParametersDB(ElasticDB):
     def __init__(self, parentLogger=None):
         """Standard Constructor"""
 
-        self.fullname = "WorkloadManagement/ElasticJobParametersDB"
+        self.fullname = "WorkloadManagement/JobParametersDB"
         self.index_name = self.getCSOption("index_name", "job_parameters")
 
         try:
-            # Connecting to the ES cluster
+            # Connecting to the OpenSearch cluster
             super().__init__(self.fullname, self.index_name, parentLogger=parentLogger)
         except Exception:
             RuntimeError("Can't connect to JobParameters index")
@@ -69,7 +69,7 @@ class JobParametersDB(ElasticDB):
                 raise RuntimeError(result["Message"])
             self.log.always("Index created:", indexName)
 
-    def getJobParameters(self, jobIDs: int | list[int], vo: str, paramList=None) -> dict:
+    def getJobParameters(self, jobIDs: int | list[int], vo: str, paramList: list[str] | None = None) -> dict:
         """Get Job Parameters defined for jobID.
           Returns a dictionary with the Job Parameters.
           If paramList is empty - all the parameters are returned.

DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -447,9 +447,10 @@ AND SubmissionTime < DATE_SUB(UTC_TIMESTAMP(),INTERVAL %d DAY)"
         """Get IDs of Jobs that were executed by a pilot"""
         cmd = "SELECT pilotID,JobID FROM JobToPilotMapping "
         if isinstance(pilotID, list):
-            cmd = cmd + " WHERE pilotID IN (%s)" % ",".join(["%s" % x for x in pilotID])
+            pilotIDs_string = ",".join(str(int(x)) for x in pilotID)
+            cmd = f"{cmd} WHERE pilotID IN ({pilotIDs_string})"
         else:
-            cmd = cmd + f" WHERE pilotID = {pilotID}"
+            cmd = f"{cmd} WHERE pilotID = {pilotID}"
 
         result = self._query(cmd)
         if not result["OK"]:

DIRAC/WorkloadManagementSystem/DB/SandboxMetadataDB.py

@@ -1,6 +1,6 @@
 """ SandboxMetadataDB class is a front-end to the metadata for sandboxes
 """
-from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC import S_ERROR, S_OK
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Security import Properties
@@ -64,7 +64,7 @@ class SandboxMetadataDB(DB):
                 "Type": "VARCHAR(64) NOT NULL",
             },
             "Indexes": {"Entity": ["EntityId"], "SBIndex": ["SBId"]},
-            "UniqueIndexes": {"Mapping": ["SBId", "EntityId", "Type"]},
+            "PrimaryKey": ["SBId", "EntityId", "Type"],
         }
 
         for tableName in self.__tablesDesc:
@@ -220,48 +220,22 @@ class SandboxMetadataDB(DB):
             return result
         return S_OK(assigned)
 
-    def __entitiesByRequesterCond(self, requesterName, requesterGroup):
-        sqlCond = []
-        requesterProps = Registry.getPropertiesForEntity(requesterGroup, name=requesterName)
-        if Properties.JOB_ADMINISTRATOR in requesterProps:
-            # Do nothing, just ensure it doesn't fit in the other cases
-            pass
-        elif Properties.JOB_SHARING in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-        elif Properties.NORMAL_USER in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-            sqlCond.append(f"o.Owner='{requesterName}'")
-        else:
-            return S_ERROR("Not authorized to access sandbox")
-        return sqlCond
-
     @convertToReturnValue
-    def unassignEntities(self, entities, requesterName, requesterGroup):
+    def unassignEntities(self, entities: list):
         """
-        Unassign jobs to sandboxes
+        Unassign entities to sandboxes. Entities are a list of strings, e.g. ['job:1234', 'job:5678'].
 
         :param list entities: list of entities to unassign
         """
         if not entities:
             return None
-        conds = self.__entitiesByRequesterCond(requesterName, requesterGroup)
 
-        sqlCmd = "CREATE TEMPORARY TABLE to_delete_EntityId (EntityId VARCHAR(128)  NOT NULL, PRIMARY KEY (EntityId)) ENGINE=MEMORY;"
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_EntityId (EntityId VARCHAR(128) NOT NULL, PRIMARY KEY (EntityId)) ENGINE=MEMORY;"
         returnValueOrRaise(self._update(sqlCmd))
         try:
             sqlCmd = "INSERT INTO to_delete_EntityId (EntityId) VALUES ( %s )"
             returnValueOrRaise(self._updatemany(sqlCmd, [(e,) for e in entities]))
             sqlCmd = "DELETE m from `sb_EntityMapping` m JOIN to_delete_EntityId t USING (EntityId)"
-            if conds:
-                sqlCmd = " ".join(
-                    [
-                        sqlCmd,
-                        "JOIN `sb_SandBoxes` s ON s.SBId = m.SBId",
-                        "JOIN `sb_Owners` o ON s.OwnerId = o.OwnerId",
-                        "WHERE",
-                        " AND ".join(conds),
-                    ]
-                )
             returnValueOrRaise(self._update(sqlCmd))
         finally:
             sqlCmd = "DROP TEMPORARY TABLE to_delete_EntityId"
@@ -309,19 +283,34 @@ class SandboxMetadataDB(DB):
             "TIMESTAMPDIFF( DAY, LastAccessTime, UTC_TIMESTAMP() ) >= %d" % self.__assignedSBGraceDays,
             f"! Assigned AND TIMESTAMPDIFF( DAY, LastAccessTime, UTC_TIMESTAMP() ) >= {self.__unassignedSBGraceDays}",
         ]
-        sqlCmd = f"SELECT SBId, SEName, SEPFN FROM `sb_SandBoxes` WHERE ( {' ) OR ( '.join(sqlCond)} )"
+        # Exclude sandboxes that are in S3 as those are handled by DiracX
+        sqlCmd = f"SELECT SBId, SEName, SEPFN FROM `sb_SandBoxes` WHERE SEPFN not like '/S3/%' AND (( {' ) OR ( '.join(sqlCond)} ))"
         return self._query(sqlCmd)
 
+    @convertToReturnValue
     def deleteSandboxes(self, SBIdList):
         """
-        Delete sandboxes
+        Delete sandboxes using a temporary table for efficiency and consistency.
         """
-        sqlSBList = ", ".join([str(sbid) for sbid in SBIdList])
-        for table in ("sb_SandBoxes", "sb_EntityMapping"):
-            sqlCmd = f"DELETE FROM `{table}` WHERE SBId IN ( {sqlSBList} )"
-            result = self._update(sqlCmd)
-            if not result["OK"]:
-                return result
+        if not SBIdList:
+            return S_OK()
+        # Create temporary table
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_SBId (SBId INTEGER(10) UNSIGNED NOT NULL, PRIMARY KEY (SBId)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd))
+        try:
+            # Insert SBIds into temporary table
+            sqlCmd = "INSERT INTO to_delete_SBId (SBId) VALUES (%s)"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(sbid,) for sbid in SBIdList]))
+            # Delete from sb_EntityMapping first (to respect FK constraints if any)
+            sqlCmd = "DELETE FROM `sb_EntityMapping` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+            # Delete from sb_SandBoxes
+            sqlCmd = "DELETE FROM `sb_SandBoxes` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+        finally:
+            # Drop temporary table
+            sqlCmd = "DROP TEMPORARY TABLE to_delete_SBId"
+            returnValueOrRaise(self._update(sqlCmd))
         return S_OK()
 
     def getSandboxId(self, SEName, SEPFN, requesterName, requesterGroup, field="SBId"):

DIRAC/WorkloadManagementSystem/DB/StatusUtils.py

@@ -0,0 +1,125 @@
+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.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
+from DIRAC.WorkloadManagementSystem.Utilities.jobAdministration import _filterJobStateTransition
+
+
+def _deleteJob(jobID, force=False):
+    """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"]:
+        gLogger.warn("Failed to set job Deleted status", result["Message"])
+        return result
+
+    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)
+    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)
+        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)
+            if not result["OK"]:
+                gLogger.error("Failed to get pilot info", result["Message"])
+                return result
+            ret = PilotAgentsDB().deletePilot(result["Value"]["PilotJobReference"])
+            if not ret["OK"]:
+                gLogger.error("Failed to delete pilot from PilotAgentsDB", ret["Message"])
+                return ret
+
+    return S_OK()
+
+
+def _killJob(jobID, sendKillCommand=True, force=False):
+    """Kill one job
+
+    :param int jobID: job ID
+    :param bool sendKillCommand: send kill command
+
+    :return: S_OK()/S_ERROR()
+    """
+    if sendKillCommand:
+        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"]:
+        gLogger.warn("Failed to set job Killed status", result["Message"])
+    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):
+    """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()
+    """
+    badIDs = []
+
+    killJobList = []
+    deleteJobList = []
+    if validJobList:
+        result = JobDB().getJobsAttributes(validJobList, ["Status"])
+        if not result["OK"]:
+            return result
+        jobStates = result["Value"]
+
+        # Get the jobs allowed to transition to the Killed state
+        killJobList.extend(_filterJobStateTransition(jobStates, JobStatus.KILLED))
+
+        if right == RIGHT_DELETE:
+            # Get the jobs allowed to transition to the Deleted state
+            deleteJobList.extend(_filterJobStateTransition(jobStates, JobStatus.DELETED))
+
+        for jobID in killJobList:
+            result = _killJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        for jobID in deleteJobList:
+            result = _deleteJob(jobID, force=force)
+            if not result["OK"]:
+                badIDs.append(jobID)
+
+        # Look for jobs that are in the Staging state to send kill signal to the stager
+        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)
+            if not result["OK"]:
+                gLogger.warn("Failed to kill some Stager tasks", result["Message"])
+
+    if nonauthJobList or badIDs:
+        result = S_ERROR("Some jobs failed deletion")
+        if nonauthJobList:
+            gLogger.warn("Non-authorized JobIDs won't be deleted", str(nonauthJobList))
+            result["NonauthorizedJobIDs"] = nonauthJobList
+        if badIDs:
+            gLogger.warn("JobIDs failed to be deleted", str(badIDs))
+            result["FailedJobIDs"] = badIDs
+        return result
+
+    jobsList = killJobList if right == RIGHT_KILL else deleteJobList
+    return S_OK(jobsList)

DIRAC/WorkloadManagementSystem/DB/tests/Test_JobDB.py

@@ -28,7 +28,7 @@ def test_getInputData(jobDB: JobDB):
     """Test the getInputData method from JobDB"""
     # Arrange
     jobDB._escapeString = MagicMock(return_value=S_OK())
-    jobDB._query = MagicMock(return_value=S_OK((("/vo/user/lfn1",), ("LFN:/vo/user/lfn2",))))
+    jobDB._query = MagicMock(return_value=S_OK([(1234, "/vo/user/lfn1"), (1234, "LFN:/vo/user/lfn2")]))
 
     # Act
     res = jobDB.getInputData(1234)

DIRAC/WorkloadManagementSystem/DB/tests/Test_StatusUtils.py

@@ -0,0 +1,28 @@
+""" unit test (pytest) of JobAdministration module
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# sut
+from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
+
+
+@pytest.mark.parametrize(
+    "jobIDs_list, right",
+    [
+        ([], "Kill"),
+        ([], "Delete"),
+        (1, "Kill"),
+        ([1, 2], "Kill"),
+    ],
+)
+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())
+
+    res = kill_delete_jobs(right, jobIDs_list)
+    assert res["OK"]