feat: removed SandboxStore.assignSandboxesToEntities

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -154,10 +154,9 @@ def removeDeletedJobs(self):
         self.log.info("Unassigning sandboxes from soon to be deleted jobs", f"({len(jobList)})")
 
         entitiesList = [f"Job:{jobId}" for jobId in jobList]
-        res = SandboxMetadataDB().unassignEntities(entitiesList)
-        if not res["OK"]:
-            self.log.error("Cannot unassign jobs to sandboxes", res["Message"])
-            return res
+        if not (result := SandboxMetadataDB().unassignEntities(entitiesList))["OK"]:
+            self.log.error("Cannot unassign jobs to sandboxes", result["Message"])
+            return result
 
         self.log.info("Attempting to remove deleted jobs", f"({len(jobList)})")
 

src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py

@@ -20,7 +20,6 @@
 
 class SandboxStoreClient:
     __validSandboxTypes = ("Input", "Output")
-    __smdb = None
 
     def __init__(self, rpcClient=None, transferClient=None, smdb=False, **kwargs):
         """Constructor
@@ -37,21 +36,8 @@ def __init__(self, rpcClient=None, transferClient=None, smdb=False, **kwargs):
         self.__transferClient = transferClient
         self.__kwargs = kwargs
         self.__vo = None
-        SandboxStoreClient.__smdb = smdb
         if "delegatedGroup" in kwargs:
             self.__vo = getVOForGroup(kwargs["delegatedGroup"])
-        if SandboxStoreClient.__smdb is True:
-            try:
-                from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
-
-                SandboxStoreClient.__smdb = SandboxMetadataDB()
-                result = SandboxStoreClient.__smdb._getConnection()  # pylint: disable=protected-access
-                if not result["OK"]:
-                    SandboxStoreClient.__smdb = False
-                else:
-                    result["Value"].close()
-            except (ImportError, RuntimeError, AttributeError):
-                SandboxStoreClient.__smdb = False
 
     def __getRPCClient(self):
         """Get an RPC client for SB service"""
@@ -227,20 +213,6 @@ def downloadSandbox(self, sbLocation, destinationDir="", inMemory=False, unpack=
     ##############
     # Jobs
 
-    def assignSandboxesToJob(self, jobId, sbList, ownerName="", ownerGroup=""):
-        """
-        Assign sandboxes to a job.
-        sbList must be a list of sandboxes and relation types
-        sbList = [ ( "SB:SEName|SEPFN", "Input" ), ( "SB:SEName|SEPFN", "Output" ) ]
-        """
-        eId = f"Job:{jobId}"
-        for sbT in sbList:
-            if sbT[1] not in self.__validSandboxTypes:
-                return S_ERROR(f"Invalid Sandbox type {sbT[1]}")
-        if SandboxStoreClient.__smdb and ownerName and ownerGroup:
-            return SandboxStoreClient.__smdb.assignSandboxesToEntities({eId: sbList}, ownerName, ownerGroup)
-        return self.__getRPCClient().assignSandboxesToEntities({eId: sbList}, ownerName, ownerGroup)
-
     def downloadSandboxForJob(self, jobId, sbType, destinationPath="", inMemory=False, unpack=True):
         """Download SB for a job"""
         result = self.__getRPCClient().getSandboxesAssignedToEntity(f"Job:{jobId}")

src/DIRAC/WorkloadManagementSystem/Executor/JobSanity.py

@@ -1,6 +1,6 @@
 """ The Job Sanity executor assigns sandboxes to the job """
 from DIRAC import S_OK
-from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
+from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.Executor.Base.OptimizerExecutor import OptimizerExecutor
 
 
@@ -15,7 +15,6 @@ class JobSanity(OptimizerExecutor):
     @classmethod
     def initializeOptimizer(cls):
         """Initialize specific parameters for JobSanityAgent."""
-        cls.sandboxClient = SandboxStoreClient(useCertificates=True, smdb=True)
         return S_OK()
 
     def optimizeJob(self, jid, jobState):
@@ -57,7 +56,8 @@ def checkInputSandbox(self, jobState, manifest):
         if not numSBsToAssign:
             return S_OK(0)
         self.jobLog.info("Assigning sandboxes", f"({numSBsToAssign} on behalf of {ownerName}@{ownerGroup}@{vo})")
-        result = self.sandboxClient.assignSandboxesToJob(jobState.jid, sbsToAssign, ownerName, ownerGroup)
+        eId = f"Job:{jobState.jid}"
+        result = SandboxMetadataDB().assignSandboxesToEntities({eId: sbsToAssign}, ownerName, ownerGroup)
         if not result["OK"]:
             self.jobLog.error("Could not assign sandboxes in the SandboxStore")
             return result

src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py

@@ -139,7 +139,7 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
                     gLogger.debug("Sandbox already exists in storage backend", res.pfn)
 
                 assignTo = {key: [(res.pfn, assignTo[key])] for key in assignTo}
-                result = self.export_assignSandboxesToEntities(assignTo)
+                result = self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"])
                 if not result["OK"]:
                     return result
                 return S_OK(res.pfn)
@@ -153,7 +153,7 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
                 fileHelper.markAsTransferred()
             sbURL = f"SB:{self.__localSEName}|{sbPath}"
             assignTo = {key: [(sbURL, assignTo[key])] for key in assignTo}
-            result = self.export_assignSandboxesToEntities(assignTo)
+            result = self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"])
             if not result["OK"]:
                 return result
             return S_OK(sbURL)
@@ -210,7 +210,7 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
 
         sbURL = f"SB:{self.__localSEName}|{sbPath}"
         assignTo = {key: [(sbURL, assignTo[key])] for key in assignTo}
-        if not (result := self.export_assignSandboxesToEntities(assignTo))["OK"]:
+        if not (result := self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"]))["OK"]:
             return result
         return S_OK(sbURL)
 
@@ -328,21 +328,6 @@ def __moveToFinalLocation(self, localFilePath, sbPath):
 
         return result
 
-    ##################
-    # Assigning sbs to jobs
-
-    types_assignSandboxesToEntities = [dict]
-
-    def export_assignSandboxesToEntities(self, enDict, ownerName="", ownerGroup=""):
-        """
-        Assign sandboxes to jobs.
-        Expects a dict of { entityId : [ ( SB, SBType ), ... ] }
-        """
-        credDict = self.getRemoteCredentials()
-        return self.sandboxDB.assignSandboxesToEntities(
-            enDict, credDict["username"], credDict["group"], ownerName, ownerGroup
-        )
-
     ##################
     # Getting assigned sandboxes