feat: refactor pilot logging to include a business logic layer

Author: martynia

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

@@ -253,3 +253,14 @@ class VOInfo(TypedDict):
 
 class Metadata(TypedDict):
     virtual_organizations: dict[str, VOInfo]
+
+
+class LogLine(BaseModel):
+    line_no: int
+    line: str
+
+
+class LogMessage(BaseModel):
+    pilot_stamp: str
+    lines: list[LogLine]
+    vo: str

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

@@ -19,3 +19,17 @@ 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}"
+
+
+async def search_message(db: PilotLogsDB, search_params: list[dict]):
+
+    return await db.search(
+        ["Message"],
+        search_params,
+        [{"parameter": "LineNumber", "direction": "asc"}],
+    )
+
+
+async def bulk_insert(db: PilotLogsDB, docs: list[dict], pilot_id: int):
+
+    await db.bulk_insert(db.index_name(pilot_id), docs)

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

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+
+from diracx.core.models import LogMessage, ScalarSearchOperator, ScalarSearchSpec
+from diracx.db.os.pilot_logs import PilotLogsDB, bulk_insert, search_message
+from diracx.db.sql.pilot_agents.db import PilotAgentsDB
+
+logger = logging.getLogger(__name__)
+
+
+async def send_message(
+    data: LogMessage,
+    pilot_logs_db: PilotLogsDB,
+    pilot_agents_db: PilotAgentsDB,
+) -> int:
+
+    # get the pilot ID corresponding to a given pilot stamp, expecting exactly one row:
+    search_params = ScalarSearchSpec(
+        parameter="PilotStamp",
+        operator=ScalarSearchOperator.EQUAL,
+        value=data.pilot_stamp,
+    )
+
+    total, result = await pilot_agents_db.search(
+        ["PilotID", "VO", "SubmissionTime"], [search_params], []
+    )
+    if total != 1:
+        logger.error(
+            "Cannot determine PilotID for requested PilotStamp: %r, (%d candidates)",
+            data.pilot_stamp,
+            total,
+        )
+        raise Exception(f"Number of rows !=1 {total}")
+
+    pilot_id, vo, submission_time = (
+        result[0]["PilotID"],
+        result[0]["VO"],
+        result[0]["SubmissionTime"],
+    )
+    docs = []
+    for line in data.lines:
+        docs.append(
+            {
+                "PilotStamp": data.pilot_stamp,
+                "PilotID": pilot_id,
+                "SubmissionTime": submission_time,
+                "VO": vo,
+                "LineNumber": line.line_no,
+                "Message": line.line,
+            }
+        )
+    # bulk insert pilot logs to OpenSearch DB:
+    await bulk_insert(pilot_logs_db, docs, pilot_id)
+    return pilot_id
+
+
+async def get_logs(
+    pilot_id: int,
+    db: PilotLogsDB,
+) -> list[dict]:
+
+    search_params = [{"parameter": "PilotID", "operator": "eq", "value": pilot_id}]
+
+    result = await search_message(db, search_params)
+
+    if not result:
+        return [{"Message": f"No logs for pilot ID = {pilot_id}"}]
+    return result

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

@@ -3,9 +3,10 @@
 import logging
 
 from fastapi import HTTPException, status
-from pydantic import BaseModel
 
-from diracx.core.models import ScalarSearchOperator, ScalarSearchSpec
+from diracx.core.models import LogMessage
+from diracx.logic.pilots.logging import get_logs as get_logs_bl
+from diracx.logic.pilots.logging import send_message as send_message_bl
 
 from ..access_policies import open_access
 from ..dependencies import PilotAgentsDB, PilotLogsDB
@@ -16,22 +17,6 @@
 router = DiracxRouter()
 
 
-class LogLine(BaseModel):
-    line_no: int
-    line: str
-
-
-class LogMessage(BaseModel):
-    pilot_stamp: str
-    lines: list[LogLine]
-    vo: str
-
-
-class DateRange(BaseModel):
-    min: str | None = None  # expects a string in ISO 8601 ("%Y-%m-%dT%H:%M:%S.%f%z")
-    max: str | None = None  # expects a string in ISO 8601 ("%Y-%m-%dT%H:%M:%S.%f%z")
-
-
 @open_access
 @router.post("/")
 async def send_message(
@@ -41,55 +26,11 @@ async def send_message(
     # user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)],
 ) -> int:
 
-    # expecting exactly one row:
-    search_params = ScalarSearchSpec(
-        parameter="PilotStamp",
-        operator=ScalarSearchOperator.EQUAL,
-        value=data.pilot_stamp,
-    )
-
-    total, result = await pilot_agents_db.search(
-        ["PilotID", "VO", "SubmissionTime"], [search_params], []
-    )
-    if total != 1:
-        logger.error(
-            "Cannot determine PilotID for requested PilotStamp: %r, (%d candidates)",
-            data.pilot_stamp,
-            total,
-        )
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST, detail=f"Number of rows !=1: {total}"
-        )
-    pilot_id, vo, submission_time = (
-        result[0]["PilotID"],
-        result[0]["VO"],
-        result[0]["SubmissionTime"],
-    )
-
-    # await check_permissions(action=ActionType.CREATE, pilot_agent_db, pilot_id),
-
-    docs = []
-    for line in data.lines:
-        docs.append(
-            {
-                "PilotStamp": data.pilot_stamp,
-                "PilotID": pilot_id,
-                "SubmissionTime": submission_time,
-                "VO": vo,
-                "LineNumber": line.line_no,
-                "Message": line.line,
-            }
-        )
-    await pilot_logs_db.bulk_insert(pilot_logs_db.index_name(pilot_id), docs)
-    """
-    search_params = [{"parameter": "PilotID", "operator": "eq", "value": pilot_id}]
-
-    result = await pilot_logs_db.search(
-        ["Message"],
-        search_params,
-        [{"parameter": "LineNumber", "direction": "asc"}],
-    )
-    """
+    # await check_permissions(action=ActionType.CREATE, pilot_agent_db, pilot_id)
+    try:
+        pilot_id = await send_message_bl(data, pilot_logs_db, pilot_agents_db)
+    except Exception as exc:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
     return pilot_id
 
 
@@ -107,13 +48,4 @@ async def get_logs(
         action=ActionType.QUERY, pilot_agents_db=pilot_agents_db, pilot_id=pilot_id
     )
 
-    search_params = [{"parameter": "PilotID", "operator": "eq", "value": pilot_id}]
-
-    result = await db.search(
-        ["Message"],
-        search_params,
-        [{"parameter": "LineNumber", "direction": "asc"}],
-    )
-    if not result:
-        return [{"Message": f"No logs for pilot ID = {pilot_id}"}]
-    return result
+    return await get_logs_bl(pilot_id, db)