feat: use bulk insert

DIRACGrid · Oct 10, 2024 · 4d4fbe2 · 4d4fbe2
1 parent 69cf73a
commit 4d4fbe2
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 23 deletions.
diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/db.py b/diracx-db/src/diracx/db/sql/pilot_agents/db.py
@@ -17,24 +17,28 @@ async def addPilotReferences(
         vo: str,
         grid_type: str = "DIRAC",
         pilot_stamps: dict | None = None,
-    ) -> list[int]:
+    ) -> None:
 
         if pilot_stamps is None:
             pilot_stamps = {}
-        row_ids = []
-        for ref in pilot_ref:
-            stamp = pilot_stamps.get(ref, "")
-            now = datetime.now(tz=timezone.utc)
-            stmt = insert(PilotAgents).values(
-                PilotJobReference=ref,
-                VO=vo,
-                GridType=grid_type,
-                SubmissionTime=now,
-                LastUpdateTime=now,
-                Status="Submitted",
-                PilotStamp=stamp,
-            )
-            result = await self.conn.execute(stmt)
-            row_ids.append(result.lastrowid)
-
-        return row_ids
+
+        now = datetime.now(tz=timezone.utc)
+
+        # Prepare the list of dictionaries for bulk insertion
+        values = [
+            {
+                "PilotJobReference": ref,
+                "VO": vo,
+                "GridType": grid_type,
+                "SubmissionTime": now,
+                "LastUpdateTime": now,
+                "Status": "Submitted",
+                "PilotStamp": pilot_stamps.get(ref, ""),
+            }
+            for ref in pilot_ref
+        ]
+
+        # Insert multiple rows in a single execute call
+        stmt = insert(PilotAgents)
+        await self.conn.execute(stmt, values)
+        return
diff --git a/diracx-db/tests/pilot_agents/test_pilotAgentsDB.py b/diracx-db/tests/pilot_agents/test_pilotAgentsDB.py
@@ -22,13 +22,10 @@ async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
         stamps = [f"stamp_{i}" for i in range(10)]
         stamp_dict = dict(zip(refs, stamps))
 
-        pilot_id = await pilot_agents_db.addPilotReferences(
+        await pilot_agents_db.addPilotReferences(
             refs, "test_vo", grid_type="DIRAC", pilot_stamps=stamp_dict
         )
-        assert pilot_id
-        assert pilot_id == list(range(1, 11))
 
-        pilot_id = await pilot_agents_db.addPilotReferences(
+        await pilot_agents_db.addPilotReferences(
             refs, "test_vo", grid_type="DIRAC", pilot_stamps=None
         )
-        assert pilot_id