feat: use temporary table instead of large IN query in PilotAgentsDB

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

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"]:
@@ -446,17 +498,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"]: