feat: enable remote pilot logging system

Author: martynia

diracx-core/src/diracx/core/exceptions.py

@@ -64,6 +64,14 @@ def __init__(self, job_id: int, detail: str | None = None):
         super().__init__(f"Job {job_id} not found" + (" ({detail})" if detail else ""))
 
 
+class PilotNotFoundError(Exception):
+    def __init__(self, pilot_stamp: str, detail: str | None = None):
+        self.pilot_stamp: str = pilot_stamp
+        super().__init__(
+            f"Pilot (stamp) {pilot_stamp} not found" + (" ({detail})" if detail else "")
+        )
+
+
 class SandboxNotFoundError(Exception):
     def __init__(self, pfn: str, se_name: str, detail: str | None = None):
         self.pfn: str = pfn

diracx-db/pyproject.toml

@@ -35,6 +35,7 @@ TaskQueueDB = "diracx.db.sql:TaskQueueDB"
 
 [project.entry-points."diracx.dbs.os"]
 JobParametersDB = "diracx.db.os:JobParametersDB"
+PilotLogsDB = "diracx.db.os:PilotLogsDB"
 
 [tool.setuptools.packages.find]
 where = ["src"]

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

@@ -1,5 +1,9 @@
 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"},
+        "SubmissionTime": {"type": "date"},
+        "LineNumber": {"type": "long"},
+        "Message": {"type": "text"},
+        "VO": {"type": "keyword"},
+        "timestamp": {"type": "date"},
+    }
+    index_prefix = "pilot_logs"
+
+    def index_name(self, doc_id: int) -> str:
+        # TODO decide how to define the index name
+        # use pilot ID
+        return f"{self.index_prefix}_{doc_id // 1e6:.0f}"

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

@@ -13,6 +13,7 @@
 from typing import Any, Self
 
 from opensearchpy import AsyncOpenSearch
+from opensearchpy.helpers import async_bulk
 
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.extensions import select_from_extension
@@ -192,6 +193,13 @@ async def upsert(self, doc_id, document) -> None:
         )
         print(f"{response=}")
 
+    async def bulk_insert(self, index_name: str, docs: list[dict[str, Any]]) -> None:
+        """Bulk inserting to database."""
+        n_inserted = await async_bulk(
+            self.client, actions=[doc | {"_index": index_name} for doc in docs]
+        )
+        logger.info("Inserted %d documents to %r", n_inserted, index_name)
+
     async def search(
         self, parameters, search, sorts, *, per_page: int = 100, page: int | None = None
     ) -> list[dict[str, Any]]:
@@ -233,6 +241,17 @@ async def search(
 
         return hits
 
+    async def delete(self, query: list[dict[str, Any]]) -> dict:
+        """Delete multiple documents by query."""
+        body = {}
+        res = {}
+        if query:
+            body["query"] = apply_search_filters(self.fields, query)
+            res = await self.client.delete_by_query(
+                body=body, index=f"{self.index_prefix}*"
+            )
+        return res
+
 
 def require_type(operator, field_name, field_type, allowed_types):
     if field_type not in allowed_types:

diracx-db/src/diracx/db/sql/job/db.py

@@ -14,7 +14,7 @@
     SortSpec,
 )
 
-from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
+from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints, get_columns
 from .schema import (
     InputData,
     JobCommands,
@@ -24,17 +24,6 @@
 )
 
 
-def _get_columns(table, parameters):
-    columns = [x for x in table.columns]
-    if parameters:
-        if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
-            raise InvalidQueryError(
-                f"Unrecognised parameters requested {unrecognised_parameters}"
-            )
-        columns = [c for c in columns if c.name in parameters]
-    return columns
-
-
 class JobDB(BaseSQLDB):
     metadata = JobDBBase.metadata
 
@@ -70,7 +59,7 @@ async def search(
     ) -> tuple[int, list[dict[Any, Any]]]:
         """Search for jobs in the database."""
         # Find which columns to select
-        columns = _get_columns(Jobs.__table__, parameters)
+        columns = get_columns(Jobs.__table__, parameters)
 
         stmt = select(*columns)
 
@@ -235,7 +224,7 @@ async def set_properties(
         required_parameters = list(required_parameters_set)[0]
         update_parameters = [{"job_id": k, **v} for k, v in properties.items()]
 
-        columns = _get_columns(Jobs.__table__, required_parameters)
+        columns = get_columns(Jobs.__table__, required_parameters)
         values: dict[str, BindParameter[Any] | datetime] = {
             c.name: bindparam(c.name) for c in columns
         }

diracx-db/src/diracx/db/sql/pilot_agents/db.py

@@ -1,10 +1,17 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
+from typing import Any
 
-from sqlalchemy import insert
+from sqlalchemy import func, insert, select
 
-from ..utils import BaseSQLDB
+from diracx.core.exceptions import InvalidQueryError
+from diracx.core.models import (
+    SearchSpec,
+    SortSpec,
+)
+
+from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints, get_columns
 from .schema import PilotAgents, PilotAgentsDBBase
 
 
@@ -44,3 +51,46 @@ async def add_pilot_references(
         stmt = insert(PilotAgents).values(values)
         await self.conn.execute(stmt)
         return
+
+    async def search(
+        self,
+        parameters: list[str] | None,
+        search: list[SearchSpec],
+        sorts: list[SortSpec],
+        *,
+        distinct: bool = False,
+        per_page: int = 100,
+        page: int | None = None,
+    ) -> tuple[int, list[dict[Any, Any]]]:
+        # Find which columns to select
+        columns = get_columns(PilotAgents.__table__, parameters)
+
+        stmt = select(*columns)
+
+        stmt = apply_search_filters(
+            PilotAgents.__table__.columns.__getitem__, stmt, search
+        )
+        stmt = apply_sort_constraints(
+            PilotAgents.__table__.columns.__getitem__, stmt, sorts
+        )
+
+        if distinct:
+            stmt = stmt.distinct()
+
+        # Calculate total count before applying pagination
+        total_count_subquery = stmt.alias()
+        total_count_stmt = select(func.count()).select_from(total_count_subquery)
+        total = (await self.conn.execute(total_count_stmt)).scalar_one()
+
+        # Apply pagination
+        if page is not None:
+            if page < 1:
+                raise InvalidQueryError("Page must be a positive integer")
+            if per_page < 1:
+                raise InvalidQueryError("Per page must be a positive integer")
+            stmt = stmt.offset((page - 1) * per_page).limit(per_page)
+
+        # Execute the query
+        return total, [
+            dict(row._mapping) async for row in (await self.conn.stream(stmt))
+        ]

diracx-db/src/diracx/db/sql/utils/__init__.py

@@ -5,6 +5,7 @@
     SQLDBUnavailableError,
     apply_search_filters,
     apply_sort_constraints,
+    get_columns,
 )
 from .functions import hash, substract_date, utcnow
 from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn
@@ -19,6 +20,7 @@
     "EnumColumn",
     "apply_search_filters",
     "apply_sort_constraints",
+    "get_columns",
     "substract_date",
     "hash",
     "SQLDBUnavailableError",

diracx-db/src/diracx/db/sql/utils/base.py

@@ -258,6 +258,17 @@ def find_time_resolution(value):
     raise InvalidQueryError(f"Cannot parse {value=}")
 
 
+def get_columns(table, parameters):
+    columns = [x for x in table.columns]
+    if parameters:
+        if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
+            raise InvalidQueryError(
+                f"Unrecognised parameters requested {unrecognised_parameters}"
+            )
+        columns = [c for c in columns if c.name in parameters]
+    return columns
+
+
 def apply_search_filters(column_mapping, stmt, search):
     for query in search:
         try:

diracx-routers/pyproject.toml

@@ -41,6 +41,7 @@ types = [
 ]
 
 [project.entry-points."diracx.services"]
+pilots = "diracx.routers.pilots:router"
 jobs = "diracx.routers.jobs:router"
 config = "diracx.routers.configuration:router"
 auth = "diracx.routers.auth:router"
@@ -49,6 +50,7 @@ auth = "diracx.routers.auth:router"
 [project.entry-points."diracx.access_policies"]
 WMSAccessPolicy = "diracx.routers.jobs.access_policies:WMSAccessPolicy"
 SandboxAccessPolicy = "diracx.routers.jobs.access_policies:SandboxAccessPolicy"
+PilotLogsAccessPolicy = "diracx.routers.pilots.access_policies:PilotLogsAccessPolicy"
 
 # Minimum version of the client supported
 [project.entry-points."diracx.min_client_version"]