fix(api): add bulk_upsert method and unit tests for opensearch BaseOSDB class

diracx-db/pyproject.toml

@@ -59,7 +59,7 @@ testpaths = ["tests"]
 addopts = [
     "-v",
     "--cov=diracx.db", "--cov-report=term-missing",
-    "-pdiracx.testing", "-pdiracx.testing.osdb",
+    "-pdiracx.testing", "-pdiracx.testing.osdb", "-pdiracx.testing.mock_osdb",
     "--import-mode=importlib",
 ]
 asyncio_mode = "auto"

diracx-db/src/diracx/db/os/job_parameters.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import UTC, datetime
+from typing import Any, Iterable
 
 from diracx.db.os.utils import BaseOSDB
 
@@ -41,3 +42,15 @@ def upsert(self, vo, doc_id, document):
             **document,
         }
         return super().upsert(vo, doc_id, document)
+
+    async def bulk_upsert(
+        self,
+        documents: Iterable[tuple[str, int, dict[str, Any]]],
+    ) -> tuple[int, list[Any]]:
+        """bulk_upsert API implementation."""
+        transformed = []
+        for vo, doc_id, document in documents:
+            timestamp = int(datetime.now(tz=UTC).timestamp() * 1000)
+            document = {"JobID": doc_id, "timestamp": timestamp, **document}
+            transformed.append((vo, doc_id, document))
+        return await super().bulk_upsert(transformed)

diracx-db/src/diracx/db/os/utils.py

@@ -4,13 +4,14 @@
 import json
 import logging
 from abc import ABCMeta, abstractmethod
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterable
 from contextvars import ContextVar
 from datetime import datetime
 from typing import Any, Self
 
 from opensearchpy import AsyncOpenSearch
 from opensearchpy.exceptions import RequestError
+from opensearchpy.helpers import async_bulk
 
 from diracx.core.exceptions import DocumentUpsertError, InvalidQueryError
 from diracx.core.extensions import DiracEntryPoint, select_from_extension
@@ -213,6 +214,35 @@ async def upsert(self, vo: str, doc_id: int, document: Any) -> None:
             response,
         )
 
+    async def bulk_upsert(
+        self,
+        documents: Iterable[tuple[str, int, dict[str, Any]]],
+    ) -> tuple[int, list[Any]]:
+        """Bulk upsert documents."""
+        actions = (
+            {
+                "_op_type": "update",
+                "_index": self.index_name(vo, doc_id),
+                "_id": doc_id,
+                "doc": document,
+                "doc_as_upsert": True,
+                "retry_on_conflict": 10,
+            }
+            for vo, doc_id, document in documents
+        )
+
+        success, errors = await async_bulk(
+            self.client,
+            actions,
+            raise_on_error=False,
+            raise_on_exception=False,
+        )
+
+        if errors:
+            logger.warning("Bulk upsert completed with %d errors", len(errors))
+
+        return success, errors
+
     async def search(
         self, parameters, search, sorts, *, per_page: int = 100, page: int | None = None
     ) -> list[dict[str, Any]]:

diracx-db/tests/conftest.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+pytest_plugins = ["diracx.testing.mock_osdb"]