feat: add VO to sb_Owners table

fstagni · Jan 24, 2024 · 799d302 · 799d302
1 parent 3b4fde8
commit 799d302
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/src/DIRAC/WorkloadManagementSystem/DB/SandboxMetadataDB.py b/src/DIRAC/WorkloadManagementSystem/DB/SandboxMetadataDB.py
@@ -33,6 +33,7 @@ def __initializeDB(self):
                 "OwnerId": "INTEGER(10) UNSIGNED AUTO_INCREMENT NOT NULL",
                 "Owner": "VARCHAR(32) NOT NULL",
                 "OwnerGroup": "VARCHAR(32) NOT NULL",
+                "VO": "VARCHAR(64) NOT NULL",
             },
             "PrimaryKey": "OwnerId",
         }
@@ -71,13 +72,14 @@ def __initializeDB(self):
 
         return self._createTables(tablesToCreate)
 
-    def __registerAndGetOwnerId(self, owner, ownerGroup):
+    def __registerAndGetOwnerId(self, owner, ownerGroup, VO):
         """
         Get the owner ID and register it if it's not there
         """
         ownerEscaped = self._escapeString(owner)["Value"]
         ownerGroupEscaped = self._escapeString(ownerGroup)["Value"]
-        sqlCmd = f"SELECT OwnerId FROM `sb_Owners` WHERE Owner = {ownerEscaped} AND OwnerGroup = {ownerGroupEscaped}"
+        VOEscaped = self._escapeString(VO)["Value"]
+        sqlCmd = f"SELECT OwnerId FROM `sb_Owners` WHERE Owner = {ownerEscaped} AND OwnerGroup = {ownerGroupEscaped} AND VO = {VOEscaped}"
         result = self._query(sqlCmd)
         if not result["OK"]:
             return result
@@ -86,7 +88,7 @@ def __registerAndGetOwnerId(self, owner, ownerGroup):
             return S_OK(data[0][0])
         # Its not there, insert it
         sqlCmd = (
-            f"INSERT INTO `sb_Owners` ( OwnerId, Owner, OwnerGroup ) VALUES ( 0, {ownerEscaped}, {ownerGroupEscaped} )"
+            f"INSERT INTO `sb_Owners` ( OwnerId, Owner, OwnerGroup, VO ) VALUES ( 0, {ownerEscaped}, {ownerGroupEscaped}, {VOEscaped} )"
         )
         result = self._update(sqlCmd)
         if not result["OK"]:
@@ -275,7 +277,7 @@ def unassignEntities(self, entities, requesterName, requesterGroup):
             updated += 1
         return S_OK(updated)
 
-    def getSandboxesAssignedToEntity(self, entityId, requesterName, requesterGroup):
+    def getSandboxesAssignedToEntity(self, entityId, requesterName, requesterGroup, requestedVO):
         """
         Get the sandboxes and the type of assignation to the jobId
         """
@@ -292,11 +294,13 @@ def getSandboxesAssignedToEntity(self, entityId, requesterName, requesterGroup):
             sqlTables.append("`sb_Owners` o")
             sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
             sqlCond.append("s.OwnerId=o.OwnerId")
+            sqlCond.append(f"o.VO='{requestedVO}'")
         elif Properties.NORMAL_USER in requesterProps:
             sqlTables.append("`sb_Owners` o")
             sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
             sqlCond.append(f"o.Owner='{requesterName}'")
             sqlCond.append("s.OwnerId=o.OwnerId")
+            sqlCond.append(f"o.VO='{requestedVO}'")
         else:
             return S_ERROR("Not authorized to access sandbox")
         sqlCmd = "SELECT DISTINCT s.SEName, s.SEPFN, e.Type FROM  {} WHERE {}".format(
@@ -378,13 +382,13 @@ def getSandboxOwner(self, SEName, SEPFN, requesterDN, requesterGroup):
         :param requesterDN: host DN used as credentials
         :param requesterGroup: group used to use as credentials (should be 'hosts')
 
-        :returns: S_OK with tuple (owner, ownerGroup)
+        :returns: S_OK with tuple (owner, ownerGroup, VO)
         """
         res = self.getSandboxId(SEName, SEPFN, None, requesterGroup, "OwnerId", requesterDN=requesterDN)
         if not res["OK"]:
             return res
 
-        sqlCmd = "SELECT `Owner`, `OwnerGroup` FROM `sb_Owners` WHERE `OwnerId` = %d" % res["Value"]
+        sqlCmd = "SELECT `Owner`, `OwnerGroup`, `VO` FROM `sb_Owners` WHERE `OwnerId` = %d" % res["Value"]
         res = self._query(sqlCmd)
         if not res["OK"]:
             return res

diff --git a/src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py
@@ -441,7 +441,7 @@ def export_getSandboxesAssignedToEntity(self, entityId):
         Get the sandboxes associated to a job and the association type
         """
         credDict = self.getRemoteCredentials()
-        result = self.sandboxDB.getSandboxesAssignedToEntity(entityId, credDict["username"], credDict["group"])
+        result = self.sandboxDB.getSandboxesAssignedToEntity(entityId, credDict["username"], credDict["group"], credDict["VO"])
         if not result["OK"]:
             return result
         sbDict = {}
@@ -616,7 +616,7 @@ def __deleteSandboxFromExternalBackend(self, SEName, SEPFN):
                 result = self.sandboxDB.getSandboxOwner(SEName, SEPFN, hostDN, "hosts")
                 if not result["OK"]:
                     return result
-                owner, _ownerDN, ownerGroup = result["Value"]
+                owner, ownerGroup, _VO = result["Value"]
 
                 request = Request()
                 request.RequestName = f"RemoteSBDeletion:{SEName}|{SEPFN}:{time.time()}"

diff --git a/tests/Integration/WorkloadManagementSystem/Test_SandboxMetadataDB.py b/tests/Integration/WorkloadManagementSystem/Test_SandboxMetadataDB.py
@@ -17,6 +17,7 @@ def test_SandboxMetadataDB():
     owner = "adminusername"
     ownerDN = "/C=ch/O=DIRAC/OU=DIRAC CI/CN=ciuser"
     ownerGroup = "dirac_admin"
+    VO = "vo"
 
     sbSE = "ProductionSandboxSE"
     sbPFN = "/sb/pfn/1.tar.bz2"
@@ -34,7 +35,7 @@ def test_SandboxMetadataDB():
 
     res = smDB.getSandboxOwner(sbSE, sbPFN, ownerDN, ownerGroup)
     assert res["OK"], res["Message"]
-    assert res["Value"] == (owner, ownerGroup)
+    assert res["Value"] == (owner, ownerGroup, VO)
 
     res = smDB.getSandboxId(sbSE, sbPFN, owner, ownerGroup)
     assert res["OK"], res["Message"]