From d5449c1445cb224f0191825922fcc47ca3ebd47b Mon Sep 17 00:00:00 2001 From: Christophe Haen Date: Tue, 8 Sep 2026 14:58:56 +0200 Subject: [PATCH 1/3] feat: use temporary table instead of large IN query in PilotAgentsDB --- .../DB/PilotAgentsDB.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index 2bd9f79e292..c67bcce8232 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -446,17 +446,24 @@ def setCurrentJobID(self, pilotRef, jobID): def getJobsForPilot(self, pilotID): """Get IDs of Jobs that were executed by a pilot""" if isinstance(pilotID, list): - cmd = "SELECT pilotID,JobID FROM JobToPilotMapping WHERE pilotID IN (" - cmd += ",".join(["%s"] * len(pilotID)) - cmd += ")" - args = [str(int(x)) for x in pilotID] + sqlCmd = "CREATE TEMPORARY TABLE to_select_JobToPilotMapping (PilotID INT UNSIGNED NOT NULL, PRIMARY KEY (PilotID)) ENGINE=MEMORY;" + returnValueOrRaise(self._update(sqlCmd)) + try: + sqlCmd = "INSERT INTO to_select_JobToPilotMapping (PilotID) VALUES ( %s )" + returnValueOrRaise(self._updatemany(sqlCmd, [(int(p),) for p in pilotID])) + sqlCmd = "SELECT pilotID, JobID FROM JobToPilotMapping JOIN to_select_JobToPilotMapping USING (PilotID)" + result = self._query(sqlCmd) + finally: + sqlCmd = "DROP TEMPORARY TABLE to_select_JobToPilotMapping" + returnValueOrRaise(self._update(sqlCmd)) + if not result["OK"]: + return result else: cmd = "SELECT pilotID,JobID FROM JobToPilotMapping WHERE pilotID = %s" args = [str(pilotID)] - - result = self._query(cmd, args=args) - if not result["OK"]: - return result + result = self._query(cmd, args=args) + if not result["OK"]: + return result resDict = {} for row in result["Value"]: From 1a65ad7e9b2b342de609a9a91923481e37e9bf8c Mon Sep 17 00:00:00 2001 From: Christophe Haen Date: Tue, 8 Sep 2026 15:28:14 +0200 Subject: [PATCH 2/3] feat: getPilotInfo use temporary table --- .../Agent/test/Test_Agent_PilotStatusAgent.py | 8 ++ .../DB/PilotAgentsDB.py | 86 +++++++++++++++---- 2 files changed, 77 insertions(+), 17 deletions(-) diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py index 8273ca12da4..c9589c9a41c 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py @@ -40,6 +40,14 @@ def psa(mocker): "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._escapeString", lambda s, c: S_OK(f'"{s}"'), ) # To bypass "connection.escape_string" + mocker.patch( + "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._update", + return_value=S_OK(), + ) + mocker.patch( + "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._updatemany", + return_value=S_OK(), + ) psa = PilotStatusAgent() psa._AgentModule__configDefaults = mockAM diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index c67bcce8232..6296b773cd1 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -41,6 +41,11 @@ class PilotAgentsDB(DB): "VO", } + # Number of values inserted into the in-memory temp tables per batch in + # getPilotInfo. Keeping batches bounded avoids overflowing MEMORY tables + # (ERROR 1114: The table is full) when querying large pilot reference lists. + PILOT_INFO_BATCH_SIZE = 10000 + def __init__(self, parentLogger=None): super().__init__("PilotAgentsDB", "WorkloadManagement/PilotAgentsDB", parentLogger=parentLogger) self._defaultLogger = self.log @@ -286,23 +291,70 @@ def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False) return S_ERROR(f"Unknown column: {col}") cmd = f"SELECT {', '.join(parameters)} FROM PilotAgents" # nosec - condSQL = [] - for key, value in [ - ("PilotJobReference", pilotRef), - ("PilotID", pilotID), - ]: - resList = [] - for v in value if isinstance(value, list) else [value] if value else []: - result = self._escapeString(v) - if not result["OK"]: - return result - resList.append(result["Value"]) - if resList: - condSQL.append(f"{key} IN ({','.join(resList)})") - if condSQL: - cmd = f"{cmd} WHERE {' AND '.join(condSQL)}" - - result = self._query(cmd, conn=conn) + tempTables = [] + activeKeys = [] + result = None + try: + for key, value in [("PilotJobReference", pilotRef), ("PilotID", pilotID)]: + values = [] + for v in value if isinstance(value, list) else [value] if value else []: + values.append(v) + if not values: + continue + tableName = f"to_select_PilotAgents_{key}" + if key == "PilotID": + sqlCmd = ( + f"CREATE TEMPORARY TABLE {tableName}" + f" (PilotID INT UNSIGNED NOT NULL, PRIMARY KEY (PilotID)) ENGINE=MEMORY;" + ) + insertValues = [(int(v),) for v in values] + else: + sqlCmd = ( + f"CREATE TEMPORARY TABLE {tableName}" + f" (PilotJobReference VARCHAR(255) NOT NULL, PRIMARY KEY (PilotJobReference)) ENGINE=MEMORY;" + ) + insertValues = [(v,) for v in values] + returnValueOrRaise(self._update(sqlCmd, conn=conn)) + tempTables.append(tableName) + activeKeys.append((key, tableName, insertValues)) + + joinSQL = [f"JOIN {tn} USING ({k})" for k, tn, _ in activeKeys] + selectCmd = f"{cmd} {' '.join(joinSQL)}" if joinSQL else cmd + + if not activeKeys: + # No conditions: plain SELECT over the whole table. + result = self._query(selectCmd, conn=conn) + else: + # Load every condition key fully except the last one, which is + # processed batch by batch (truncating between batches) so the + # in-memory temp table never holds all the rows at once. This + # preserves the AND semantics: the union of the per-batch + # results equals the full JOIN result. + for key, tableName, insertValues in activeKeys[:-1]: + insertCmd = f"INSERT INTO {tableName} ({key}) VALUES ( %s )" # nosec: B608 + for i in range(0, len(insertValues), self.PILOT_INFO_BATCH_SIZE): + returnValueOrRaise( + self._updatemany(insertCmd, insertValues[i : i + self.PILOT_INFO_BATCH_SIZE], conn=conn) + ) + batchKey, batchTable, batchValues = activeKeys[-1] + insertCmd = f"INSERT INTO {batchTable} ({batchKey}) VALUES ( %s )" # nosec: B608 + resultRows = [] + for i in range(0, len(batchValues), self.PILOT_INFO_BATCH_SIZE): + returnValueOrRaise(self._update(f"TRUNCATE TABLE {batchTable}", conn=conn)) + returnValueOrRaise( + self._updatemany(insertCmd, batchValues[i : i + self.PILOT_INFO_BATCH_SIZE], conn=conn) + ) + batchResult = self._query(selectCmd, conn=conn) + if not batchResult["OK"]: + result = batchResult + break + resultRows.extend(batchResult["Value"]) + else: + result = S_OK(resultRows) + finally: + for tableName in tempTables: + sqlCmd = f"DROP TEMPORARY TABLE {tableName}" + returnValueOrRaise(self._update(sqlCmd, conn=conn)) if not result["OK"]: return result if not result["Value"]: From 5a10eacbfe6b9128f88bfb94bf2cf68333196f8f Mon Sep 17 00:00:00 2001 From: aldbr Date: Thu, 15 Oct 2026 00:53:38 +0200 Subject: [PATCH 3/3] fix: select pilot info by condition instead of a large IN list Add PilotAgentsDB.getPilotInfoForSelection, which returns the pilot information for a selectPilots-style condition in a single query, and attaches the jobs through a subquery on the same condition. SiteDirector and PilotStatusAgent used selectPilots followed by getPilotInfo on the returned references: above a few thousand references MySQL gives up the index on the IN list and scans the whole PilotAgents table. This replaces the temporary-table approach: getPilotInfo and getJobsForPilot go back to their previous behaviour (S_OK/S_ERROR contract, duplicate references tolerated). Co-Authored-By: Claude Opus 5.5 --- .../Agent/PilotStatusAgent.py | 15 +- .../Agent/SiteDirector.py | 30 +-- .../Agent/test/Test_Agent_PilotStatusAgent.py | 20 +- .../DB/PilotAgentsDB.py | 178 +++++++----------- .../Test_PilotAgentsDB.py | 24 +++ 5 files changed, 109 insertions(+), 158 deletions(-) diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py index a32f60a2ae8..1ebb7669e94 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py @@ -82,22 +82,15 @@ def handleOldPilots(self, connection): """ pilotsToAccount = {} timeLimitToConsider = TimeUtilities.toString(DiracTime.utcnow() - TimeUtilities.day * self.pilotStalledDays) - result = self.pilotDB.selectPilots( + result = self.pilotDB.getPilotInfoForSelection( {"Status": PilotStatus.PILOT_TRANSIENT_STATES}, older=timeLimitToConsider, timeStamp="LastUpdateTime" ) if not result["OK"]: - self.log.error("Failed to get the Pilot Agents") + self.log.error("Failed to get Info for Pilot Agents", result["Message"]) return result - if not result["Value"]: - return S_OK() - - refList = result["Value"] - result = self.pilotDB.getPilotInfo(refList) - if not result["OK"]: - self.log.error("Failed to get Info for Pilot Agents") - return result - pilotsDict = result["Value"] + if not pilotsDict: + return S_OK() for pRef in pilotsDict: if pilotsDict[pRef].get("Jobs") and self._checkJobLastUpdateTime( diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py b/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py index ef4a14426f4..05a71ca5734 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py @@ -748,8 +748,8 @@ def _monitorPilotsPerQueue(self, queue: str): if callable(getattr(ce, "cleanupPilots", None)): ce.cleanupPilots() - # Select pilots in a transient states - result = self.pilotAgentsDB.selectPilots( + # Get the information of the pilots in a transient state + result = self.pilotAgentsDB.getPilotInfoForSelection( { "DestinationSite": ceName, "Queue": queueName, @@ -759,19 +759,13 @@ def _monitorPilotsPerQueue(self, queue: str): "VO": self.vo, } ) - if not result["OK"]: - self.log.error("Failed to select pilots", f": {result['Message']}") - return result - pilotRefs = result["Value"] - if not pilotRefs: - return S_OK() - - # Get their information - result = self.pilotAgentsDB.getPilotInfo(pilotRefs) if not result["OK"]: self.log.error("Failed to get pilots info from DB", result["Message"]) return result pilotDict = result["Value"] + if not pilotDict: + return S_OK() + pilotRefs = list(pilotDict) # The proxy is used for checking the pilot status and renewals # We really need at least a few hours otherwise the renewed @@ -799,7 +793,7 @@ def _monitorPilotsPerQueue(self, queue: str): # Check if the accounting is to be sent if self.sendAccounting: - result = self.pilotAgentsDB.selectPilots( + result = self.pilotAgentsDB.getPilotInfoForSelection( { "DestinationSite": ceName, "Queue": queueName, @@ -809,20 +803,14 @@ def _monitorPilotsPerQueue(self, queue: str): "Status": PilotStatus.PILOT_FINAL_STATES, } ) - if not result["OK"]: - self.log.error("Failed to select pilots", result["Message"]) - return result - - pilotRefs = result["Value"] - if not pilotRefs: - return S_OK() - - result = self.pilotAgentsDB.getPilotInfo(pilotRefs) if not result["OK"]: self.log.error("Failed to get pilots info from DB", result["Message"]) return result pilotDict = result["Value"] + if not pilotDict: + return S_OK() + result = self._sendPilotAccounting(pilotDict) if not result["OK"]: self.log.error("Failed to send pilot agent accounting") diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py index c9589c9a41c..feb0cc74c66 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py @@ -29,7 +29,7 @@ def psa(mocker): side_effect=lambda x, y=None: y, create=True, ) - module_str = "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.selectPilots" + module_str = "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.getPilotInfoForSelection" mocker.patch(module_str, side_effect=mockReply) mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.__init__", side_effect=mockNone) mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.JobDB.__init__", side_effect=mockNone) @@ -40,14 +40,6 @@ def psa(mocker): "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._escapeString", lambda s, c: S_OK(f'"{s}"'), ) # To bypass "connection.escape_string" - mocker.patch( - "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._update", - return_value=S_OK(), - ) - mocker.patch( - "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB._updatemany", - return_value=S_OK(), - ) psa = PilotStatusAgent() psa._AgentModule__configDefaults = mockAM @@ -63,16 +55,8 @@ def psa(mocker): @pytest.mark.parametrize( "mockReplyInput, expected", [ - ({"OK": True, "Value": False}, {"OK": True, "Value": None}), - ( - {"OK": True, "Value": ["Test"]}, - { - "OK": False, - "Message": "No pilots found for PilotJobReference(s): ['Test']", - }, - ), + ({"OK": True, "Value": {}}, {"OK": True, "Value": None}), ({"OK": False, "Message": "Test"}, {"OK": False, "Message": "Test"}), - ({"OK": True, "Value": []}, {"OK": True, "Value": None}), ], ) def test_handleOldPilots(psa, mockReplyInput, expected): diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index 6296b773cd1..5c4a598cf37 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -41,11 +41,6 @@ class PilotAgentsDB(DB): "VO", } - # Number of values inserted into the in-memory temp tables per batch in - # getPilotInfo. Keeping batches bounded avoids overflowing MEMORY tables - # (ERROR 1114: The table is full) when querying large pilot reference lists. - PILOT_INFO_BATCH_SIZE = 10000 - def __init__(self, parentLogger=None): super().__init__("PilotAgentsDB", "WorkloadManagement/PilotAgentsDB", parentLogger=parentLogger) self._defaultLogger = self.log @@ -265,8 +260,52 @@ def clearPilots(self, interval=30, aborted_interval=7): ########################################################################################## def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False): - """Get all the information for the pilot job reference or reference list""" + """Get all the information for the pilot job reference or reference list + + To get the information of pilots matching some criteria, prefer :meth:`getPilotInfoForSelection` + over :meth:`selectPilots` followed by this method: it avoids sending a potentially huge list of + references back to the DB in an ``IN`` clause. + """ + condSQL = [] + for key, value in [ + ("PilotJobReference", pilotRef), + ("PilotID", pilotID), + ]: + resList = [] + for v in value if isinstance(value, list) else [value] if value else []: + result = self._escapeString(v) + if not result["OK"]: + return result + resList.append(result["Value"]) + if resList: + condSQL.append(f"{key} IN ({','.join(resList)})") + condition = f"WHERE {' AND '.join(condSQL)}" if condSQL else "" + + result = self._getPilotInfoForCondition(condition, paramNames, conn=conn) + if not result["OK"]: + return result + if not result["Value"]: + msg = "No pilots found" + if pilotRef: + msg += f" for PilotJobReference(s): {pilotRef}" + return S_ERROR(DErrno.EWMSNOPILOT, msg) + return result + + def getPilotInfoForSelection( + self, condDict, older=None, newer=None, timeStamp="SubmissionTime", paramNames=None, conn=False + ): + """Get the information of the pilots matching the selection, in the format of :meth:`getPilotInfo` + + The selection arguments are the ones of :meth:`selectPilots`. + :return: S_OK with a dictionary keyed by the first parameter (PilotJobReference by default); + empty if no pilot matches + """ + condition = self.buildCondition(condDict, older, newer, timeStamp) + return self._getPilotInfoForCondition(condition, paramNames, conn=conn) + + def _getPilotInfoForCondition(self, condition, paramNames=None, conn=False): + """Get the pilot information for the pilots matching a SQL condition (e.g. ``WHERE ...``)""" parameters = ( [ "PilotJobReference", @@ -290,98 +329,28 @@ def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False) if not col in self.ALLOWED_COLUMNS: return S_ERROR(f"Unknown column: {col}") - cmd = f"SELECT {', '.join(parameters)} FROM PilotAgents" # nosec - tempTables = [] - activeKeys = [] - result = None - try: - for key, value in [("PilotJobReference", pilotRef), ("PilotID", pilotID)]: - values = [] - for v in value if isinstance(value, list) else [value] if value else []: - values.append(v) - if not values: - continue - tableName = f"to_select_PilotAgents_{key}" - if key == "PilotID": - sqlCmd = ( - f"CREATE TEMPORARY TABLE {tableName}" - f" (PilotID INT UNSIGNED NOT NULL, PRIMARY KEY (PilotID)) ENGINE=MEMORY;" - ) - insertValues = [(int(v),) for v in values] - else: - sqlCmd = ( - f"CREATE TEMPORARY TABLE {tableName}" - f" (PilotJobReference VARCHAR(255) NOT NULL, PRIMARY KEY (PilotJobReference)) ENGINE=MEMORY;" - ) - insertValues = [(v,) for v in values] - returnValueOrRaise(self._update(sqlCmd, conn=conn)) - tempTables.append(tableName) - activeKeys.append((key, tableName, insertValues)) - - joinSQL = [f"JOIN {tn} USING ({k})" for k, tn, _ in activeKeys] - selectCmd = f"{cmd} {' '.join(joinSQL)}" if joinSQL else cmd - - if not activeKeys: - # No conditions: plain SELECT over the whole table. - result = self._query(selectCmd, conn=conn) - else: - # Load every condition key fully except the last one, which is - # processed batch by batch (truncating between batches) so the - # in-memory temp table never holds all the rows at once. This - # preserves the AND semantics: the union of the per-batch - # results equals the full JOIN result. - for key, tableName, insertValues in activeKeys[:-1]: - insertCmd = f"INSERT INTO {tableName} ({key}) VALUES ( %s )" # nosec: B608 - for i in range(0, len(insertValues), self.PILOT_INFO_BATCH_SIZE): - returnValueOrRaise( - self._updatemany(insertCmd, insertValues[i : i + self.PILOT_INFO_BATCH_SIZE], conn=conn) - ) - batchKey, batchTable, batchValues = activeKeys[-1] - insertCmd = f"INSERT INTO {batchTable} ({batchKey}) VALUES ( %s )" # nosec: B608 - resultRows = [] - for i in range(0, len(batchValues), self.PILOT_INFO_BATCH_SIZE): - returnValueOrRaise(self._update(f"TRUNCATE TABLE {batchTable}", conn=conn)) - returnValueOrRaise( - self._updatemany(insertCmd, batchValues[i : i + self.PILOT_INFO_BATCH_SIZE], conn=conn) - ) - batchResult = self._query(selectCmd, conn=conn) - if not batchResult["OK"]: - result = batchResult - break - resultRows.extend(batchResult["Value"]) - else: - result = S_OK(resultRows) - finally: - for tableName in tempTables: - sqlCmd = f"DROP TEMPORARY TABLE {tableName}" - returnValueOrRaise(self._update(sqlCmd, conn=conn)) + result = self._query(f"SELECT {', '.join(parameters)} FROM PilotAgents {condition}", conn=conn) # nosec if not result["OK"]: return result - if not result["Value"]: - msg = "No pilots found" - if pilotRef: - msg += f" for PilotJobReference(s): {pilotRef}" - return S_ERROR(DErrno.EWMSNOPILOT, msg) - resDict = {} - pilotIDs = [] - for row in result["Value"]: - pilotDict = {} - for i, par in enumerate(parameters): - pilotDict[par] = row[i] - if par == "PilotID": - pilotIDs.append(row[i]) - resDict[row[0]] = pilotDict - - result = self.getJobsForPilot(pilotIDs) - if not result["OK"]: + resDict = {row[0]: dict(zip(parameters, row)) for row in result["Value"]} + if not resDict or "PilotID" not in parameters: return S_OK(resDict) - jobsDict = result["Value"] - for pilotRef, pilotInfo in resDict.items(): - pilotID = pilotInfo["PilotID"] - if pilotID in jobsDict: - pilotInfo["Jobs"] = jobsDict[pilotID] + # Reuse the condition in a subquery rather than sending the list of PilotIDs back to the DB + cmd = ( + "SELECT PilotID, JobID FROM JobToPilotMapping" + f" WHERE PilotID IN (SELECT PilotID FROM PilotAgents {condition})" # nosec + ) + result = self._query(cmd, conn=conn) + if not result["OK"]: + return S_OK(resDict) + jobsDict = {} + for pID, jobID in result["Value"]: + jobsDict.setdefault(pID, []).append(jobID) + for pilotInfo in resDict.values(): + if pilotInfo["PilotID"] in jobsDict: + pilotInfo["Jobs"] = jobsDict[pilotInfo["PilotID"]] return S_OK(resDict) @@ -498,24 +467,17 @@ def setCurrentJobID(self, pilotRef, jobID): def getJobsForPilot(self, pilotID): """Get IDs of Jobs that were executed by a pilot""" if isinstance(pilotID, list): - sqlCmd = "CREATE TEMPORARY TABLE to_select_JobToPilotMapping (PilotID INT UNSIGNED NOT NULL, PRIMARY KEY (PilotID)) ENGINE=MEMORY;" - returnValueOrRaise(self._update(sqlCmd)) - try: - sqlCmd = "INSERT INTO to_select_JobToPilotMapping (PilotID) VALUES ( %s )" - returnValueOrRaise(self._updatemany(sqlCmd, [(int(p),) for p in pilotID])) - sqlCmd = "SELECT pilotID, JobID FROM JobToPilotMapping JOIN to_select_JobToPilotMapping USING (PilotID)" - result = self._query(sqlCmd) - finally: - sqlCmd = "DROP TEMPORARY TABLE to_select_JobToPilotMapping" - returnValueOrRaise(self._update(sqlCmd)) - if not result["OK"]: - return result + cmd = "SELECT pilotID,JobID FROM JobToPilotMapping WHERE pilotID IN (" + cmd += ",".join(["%s"] * len(pilotID)) + cmd += ")" + args = [str(int(x)) for x in pilotID] else: cmd = "SELECT pilotID,JobID FROM JobToPilotMapping WHERE pilotID = %s" args = [str(pilotID)] - result = self._query(cmd, args=args) - if not result["OK"]: - return result + + result = self._query(cmd, args=args) + if not result["OK"]: + return result resDict = {} for row in result["Value"]: diff --git a/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py b/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py index 3243ae9ad54..eed90f22ee5 100644 --- a/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py +++ b/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py @@ -186,3 +186,27 @@ def test_PivotedPilotSummaryTable(): assert "Total" in columns cleanUpPilots(pilotRef) + + +def test_getPilotInfoForSelection(): + """Selecting pilot info by criteria gives the same result as selectPilots + getPilotInfo""" + pilotRef = preparePilots([2, 3], "TestSite", "TestCE", "VO") + try: + res = paDB.setJobForPilot(123, "pilotRef_0", updateStatus=False) + assert res["OK"] is True, res["Message"] + + condDict = {"DestinationSite": "TestCE", "Status": "Submitted"} + res = paDB.getPilotInfoForSelection(condDict) + assert res["OK"] is True, res["Message"] + assert sorted(res["Value"]) == ["pilotRef_0", "pilotRef_1"] + assert res["Value"]["pilotRef_0"]["Jobs"] == [123] + assert "Jobs" not in res["Value"]["pilotRef_1"] + + refs = paDB.selectPilots(condDict)["Value"] + assert paDB.getPilotInfo(refs)["Value"] == res["Value"] + + res = paDB.getPilotInfoForSelection({"DestinationSite": "TestCE", "Status": "Unknown"}) + assert res["OK"] is True, res["Message"] + assert res["Value"] == {} + finally: + cleanUpPilots(pilotRef)