From d5449c1445cb224f0191825922fcc47ca3ebd47b Mon Sep 17 00:00:00 2001 From: Christophe Haen Date: Tue, 8 Sep 2026 14:58:56 +0200 Subject: [PATCH 1/2] 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/2] 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"]: