feat: Add pilot logging

Author: Robin-Van-de-Merghel

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

@@ -352,3 +352,15 @@ class PilotRefreshTokenPayload(BaseTokenPayload):
 
 class PilotCredentialsInfo(PilotSecretsInfo):
     pilot_stamp: str
+
+
+class LogLine(BaseModel):
+    timestamp: str
+    severity: str
+    message: str
+    scope: str
+
+
+class LogMessage(BaseModel):
+    pilot_stamp: str
+    lines: list[LogLine]

diracx-db/pyproject.toml

@@ -33,6 +33,7 @@ SandboxMetadataDB = "diracx.db.sql:SandboxMetadataDB"
 TaskQueueDB = "diracx.db.sql:TaskQueueDB"
 
 [project.entry-points."diracx.dbs.os"]
+PilotLogsDB = "diracx.db.os:PilotLogsDB"
 JobParametersDB = "diracx.db.os:JobParametersDB"
 
 [build-system]

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,21 @@
+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:
+        # TODO decide how to define the index name
+        # use pilot ID
+        return f"{self.index_prefix}_{doc_id}"

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-logic/src/diracx/logic/pilots/query.py

@@ -7,9 +7,11 @@
     ScalarSearchOperator,
     ScalarSearchSpec,
     SearchParams,
+    SortDirection,
     VectorSearchOperator,
     VectorSearchSpec,
 )
+from diracx.db.os.pilot_logs import PilotLogsDB
 from diracx.db.sql import PilotAgentsDB
 
 MAX_PER_PAGE = 10000
@@ -176,3 +178,40 @@ async def get_secrets_by_uuid_bulk(
         )
 
     return secrets
+
+
+async def search_logs(
+    vo: str,
+    body: SearchParams | None,
+    per_page: int,
+    page: int,
+    pilot_logs_db: PilotLogsDB,
+) -> tuple[int, list[dict]]:
+    """Retrieve logs from OpenSearch for a given PilotStamp."""
+    # Apply a limit to per_page to prevent abuse of the API
+    if per_page > MAX_PER_PAGE:
+        per_page = MAX_PER_PAGE
+
+    if body is None:
+        body = SearchParams()
+
+    search = body.search
+    parameters = body.parameters
+    sorts = body.sort
+
+    # Add the vo to make sure that we filter for pilots we can see
+    # TODO: Test it
+    search = search + [
+        {
+            "parameter": "VO",
+            "operator": "eq",
+            "value": vo,
+        }
+    ]
+
+    if not sorts:
+        sorts = [{"parameter": "TimeStamp", "direction": SortDirection("asc")}]
+
+    return await pilot_logs_db.search(
+        parameters=parameters, search=search, sorts=sorts, per_page=per_page, page=page
+    )

diracx-logic/src/diracx/logic/pilots/resources.py

@@ -0,0 +1,52 @@
+"""File dedicated to logic for pilot only resources (logs, jobs, etc.)."""
+
+from __future__ import annotations
+
+import logging
+
+from diracx.core.exceptions import PilotNotFoundError
+from diracx.core.models import (
+    LogMessage,
+)
+from diracx.db.os.pilot_logs import PilotLogsDB
+from diracx.db.sql.pilots.db import PilotAgentsDB
+
+from .query import get_pilot_ids_by_stamps
+
+logger = logging.getLogger(__name__)
+
+
+MAX_PER_PAGE = 10000
+
+
+async def send_message(
+    data: LogMessage,
+    pilot_logs_db: PilotLogsDB,
+    pilot_db: PilotAgentsDB,
+    vo: str,
+):
+    try:
+        pilot_ids = await get_pilot_ids_by_stamps(
+            pilot_db=pilot_db, pilot_stamps=[data.pilot_stamp]
+        )
+        pilot_id = pilot_ids[0]  # Semantic
+    except PilotNotFoundError:
+        # If a pilot is not found, then we still store the data (to not lost it)
+        # We log it as it's not supposed to happen
+        pilot_id = -1  # To detect
+
+    docs = []
+    for line in data.lines:
+        docs.append(
+            {
+                "PilotStamp": data.pilot_stamp,
+                "PilotID": pilot_id,
+                "VO": vo,
+                "Severity": line.severity,
+                "Message": line.message,
+                "TimeStamp": line.timestamp,
+                "Scope": line.scope,
+            }
+        )
+    # bulk insert pilot logs to OpenSearch DB:
+    await pilot_logs_db.bulk_insert(pilot_logs_db.index_name(vo, pilot_id), docs)

diracx-routers/src/diracx/routers/dependencies.py

@@ -23,6 +23,7 @@
 from diracx.core.settings import DevelopmentSettings as _DevelopmentSettings
 from diracx.core.settings import SandboxStoreSettings as _SandboxStoreSettings
 from diracx.db.os import JobParametersDB as _JobParametersDB
+from diracx.db.os import PilotLogsDB as _PilotLogsDB
 from diracx.db.sql import AuthDB as _AuthDB
 from diracx.db.sql import JobDB as _JobDB
 from diracx.db.sql import JobLoggingDB as _JobLoggingDB
@@ -50,7 +51,7 @@ def add_settings_annotation(cls: T) -> T:
 
 # Opensearch databases
 JobParametersDB = Annotated[_JobParametersDB, Depends(_JobParametersDB.session)]
-
+PilotLogsDB = Annotated[_PilotLogsDB, Depends(_PilotLogsDB.session)]
 
 # Miscellaneous
 Config = Annotated[_Config, Depends(ConfigSource.create)]

diracx-routers/src/diracx/routers/pilots/__init__.py

@@ -6,10 +6,12 @@
 from .auth import router as auth_router
 from .management import router as management_router
 from .query import router as query_router
+from .resources import router as resources_router
 
 logger = logging.getLogger(__name__)
 
 router = DiracxRouter(require_auth=False)
 router.include_router(management_router)
 router.include_router(query_router)
 router.include_router(auth_router)
+router.include_router(resources_router)

diracx-routers/src/diracx/routers/pilots/access_policies.py

@@ -23,6 +23,8 @@ class ActionType(StrEnum):
     CHANGE_PILOT_FIELD = auto()
     # Read some pilot info
     READ_PILOT_FIELDS = auto()
+    # Read pilot logs
+    READ_PILOT_LOGS = auto()
 
 
 class PilotManagementAccessPolicy(BaseAccessPolicy):
@@ -44,7 +46,7 @@ async def policy(
     ):
         assert action, "action is a mandatory parameter"
 
-        if action == ActionType.READ_PILOT_FIELDS:
+        if action in {ActionType.READ_PILOT_FIELDS, ActionType.READ_PILOT_LOGS}:
             if NORMAL_USER in user_info.properties:
                 return