feat: Add pilot logging

Author: pre-commit-ci[bot]

diracx-core/src/diracx/core/models.py

@@ -34,7 +34,7 @@ class VectorSearchOperator(StrEnum):
 class ScalarSearchSpec(TypedDict):
     parameter: str
     operator: ScalarSearchOperator
-    value: str | int | datetime
+    value: str | int | datetime | bytes
 
 
 class VectorSearchSpec(TypedDict):
@@ -392,3 +392,10 @@ class VacuumPilotAuth(PilotAuthCredentials):
     grid_type: str
     grid_site: str
     status: str
+
+
+class LogLine(BaseModel):
+    timestamp: str
+    severity: str
+    message: str
+    scope: str

diracx-db/pyproject.toml

@@ -34,6 +34,7 @@ TaskQueueDB = "diracx.db.sql:TaskQueueDB"
 
 [project.entry-points."diracx.dbs.os"]
 JobParametersDB = "diracx.db.os:JobParametersDB"
+PilotLogsDB = "diracx.db.os:PilotLogsDB"
 
 [build-system]
 requires = ["hatchling", "hatch-vcs"]

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
-__all__ = ("JobParametersDB",)
+__all__ = ("JobParametersDB", "PilotLogsDB")
 
 from .job_parameters import JobParametersDB
+from .pilot_logs import PilotLogsDB

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

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from diracx.db.os.utils import BaseOSDB
+
+
+class PilotLogsDB(BaseOSDB):
+    fields = {
+        "PilotStamp": {"type": "keyword"},
+        "PilotID": {"type": "long"},
+        "Severity": {"type": "keyword"},
+        "Message": {"type": "text"},
+        "VO": {"type": "keyword"},
+        "TimeStamp": {"type": "date_nanos"},
+        "Scope": {"type": "keyword"},
+    }
+    index_prefix = "pilot_logs"
+
+    def index_name(self, vo: str, doc_id: int) -> str:
+        split = int(int(doc_id) // 1e6)
+        # We split docs into smaller one (grouped by 1 million pilot)
+        # Ex: pilot_logs_dteam_1030m
+        return f"{self.index_prefix}_{vo.lower()}_{split}m"

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

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from opensearchpy.helpers import async_bulk
+
 __all__ = ("BaseOSDB",)
 
 import contextlib
@@ -197,13 +199,35 @@ async def upsert(self, vo: str, doc_id: int, document: Any) -> None:
             response,
         )
 
+    async def bulk_insert(self, index_name: str, docs: list[dict[str, Any]]) -> None:
+        """Bulk inserting to database."""
+        n_inserted, failed = await async_bulk(
+            self.client, actions=[doc | {"_index": index_name} for doc in docs]
+        )
+        logger.info("Inserted %d documents to %s", n_inserted, index_name)
+
+        if failed:
+            logger.error("Fail to insert %d documents to %s", failed, index_name)
+
     async def search(
-        self, parameters, search, sorts, *, per_page: int = 100, page: int | None = None
-    ) -> list[dict[str, Any]]:
+        self,
+        parameters,
+        search,
+        sorts,
+        *,
+        per_page: int = 10000,
+        page: int | None = None,
+    ) -> tuple[int, list[dict[str, Any]]]:
         """Search the database for matching results.
 
         See the DiracX search API documentation for details.
         """
+        if page:
+            if page < 1:
+                raise InvalidQueryError("Page must be a positive integer")
+            if per_page < 1:
+                raise InvalidQueryError("Per page must be a positive integer")
+
         body = {}
         if parameters:
             body["_source"] = parameters
@@ -213,7 +237,12 @@ async def search(
         for sort in sorts:
             field_name = sort["parameter"]
             field_type = self.fields.get(field_name, {}).get("type")
-            require_type("sort", field_name, field_type, {"keyword", "long", "date"})
+            require_type(
+                "sort",
+                field_name,
+                field_type,
+                {"keyword", "long", "date", "date_nanos"},
+            )
             body["sort"].append({field_name: {"order": sort["direction"]}})
 
         params = {}
@@ -226,17 +255,19 @@ async def search(
         )
         hits = [hit["_source"] for hit in response["hits"]["hits"]]
 
+        total_hits = response["hits"]["total"]["value"]
+
         # Dates are returned as strings, convert them to Python datetimes
         for hit in hits:
             for field_name in hit:
                 if field_name not in self.fields:
                     continue
-                if self.fields[field_name]["type"] == "date":
+                if self.fields[field_name]["type"] in ["date", "date_nanos"]:
                     hit[field_name] = datetime.strptime(
                         hit[field_name], "%Y-%m-%dT%H:%M:%S.%f%z"
                     )
 
-        return hits
+        return total_hits, hits
 
 
 def require_type(operator, field_name, field_type, allowed_types):

diracx-db/tests/auth/test_authorization_flow.py

@@ -4,7 +4,7 @@
 from sqlalchemy.exc import NoResultFound
 
 from diracx.core.exceptions import AuthorizationError
-from diracx.db.sql import AuthDB
+from diracx.db.sql.auth.db import AuthDB
 
 MAX_VALIDITY = 2
 EXPIRED = 0

diracx-db/tests/auth/test_device_flow.py

@@ -7,7 +7,7 @@
 from sqlalchemy.exc import NoResultFound
 
 from diracx.core.exceptions import AuthorizationError
-from diracx.db.sql import AuthDB
+from diracx.db.sql.auth.db import AuthDB
 from diracx.db.sql.auth.schema import USER_CODE_LENGTH
 from diracx.db.sql.utils.functions import substract_date
 

diracx-db/tests/auth/test_refresh_token.py

@@ -3,7 +3,7 @@
 import pytest
 from uuid_utils import UUID, uuid7
 
-from diracx.db.sql import AuthDB
+from diracx.db.sql.auth.db import AuthDB
 from diracx.db.sql.auth.schema import RefreshTokenStatus