cont

Author: chaen

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,8 +291,9 @@ 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
-        joinSQL = []
         tempTables = []
+        activeKeys = []
+        result = None
         try:
             for key, value in [("PilotJobReference", pilotRef), ("PilotID", pilotID)]:
                 values = []
@@ -310,12 +316,41 @@ def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False)
                     insertValues = [(v,) for v in values]
                 returnValueOrRaise(self._update(sqlCmd, conn=conn))
                 tempTables.append(tableName)
-                sqlCmd = f"INSERT INTO {tableName} ({key}) VALUES ( %s )"  # nosec: B608
-                returnValueOrRaise(self._updatemany(sqlCmd, insertValues, conn=conn))
-                joinSQL.append(f"JOIN {tableName} USING ({key})")
-            if joinSQL:
-                cmd = f"{cmd} {' '.join(joinSQL)}"
-            result = self._query(cmd, conn=conn)
+                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}"