feat: Fixes and now use of search engine instead of DIY fetch records

Author: Robin-Van-de-Merghel

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

@@ -13,13 +13,7 @@
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.models import JobCommand, SearchSpec, SortSpec
 
-from ..utils import (
-    BaseSQLDB,
-    _get_columns,
-    apply_search_filters,
-    apply_sort_constraints,
-    utcnow,
-)
+from ..utils import BaseSQLDB, _get_columns, apply_search_filters, utcnow
 from .schema import (
     HeartBeatLoggingInfo,
     InputData,
@@ -63,7 +57,7 @@ async def summary(self, group_by, search) -> list[dict[str, str | int]]:
             if row.count > 0  # type: ignore
         ]
 
-    async def search(
+    async def search_jobs(
         self,
         parameters: list[str] | None,
         search: list[SearchSpec],
@@ -74,34 +68,15 @@ async def search(
         page: int | None = None,
     ) -> tuple[int, list[dict[Any, Any]]]:
         """Search for jobs in the database."""
-        # Find which columns to select
-        columns = _get_columns(Jobs.__table__, parameters)
-
-        stmt = select(*columns)
-
-        stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
-        stmt = apply_sort_constraints(Jobs.__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))
-        ]
+        return await self.search(
+            model=Jobs,
+            parameters=parameters,
+            search=search,
+            sorts=sorts,
+            distinct=distinct,
+            per_page=per_page,
+            page=page,
+        )
 
     async def create_job(self, compressed_original_jdl: str):
         """Used to insert a new job with original JDL. Returns inserted job id."""

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

@@ -1,16 +1,14 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
-from typing import Any, Sequence
+from typing import Any
 
-from sqlalchemy import RowMapping, bindparam, func
+from sqlalchemy import bindparam
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.sql import delete, insert, select, update
+from sqlalchemy.sql import delete, insert, update
 
 from diracx.core.exceptions import (
-    InvalidQueryError,
     PilotAlreadyAssociatedWithJobError,
-    PilotJobsNotFoundError,
     PilotNotFoundError,
 )
 from diracx.core.models import (
@@ -21,10 +19,6 @@
 
 from ..utils import (
     BaseSQLDB,
-    _get_columns,
-    apply_search_filters,
-    apply_sort_constraints,
-    fetch_records_bulk_or_raises,
 )
 from .schema import (
     JobToPilotMapping,
@@ -43,7 +37,7 @@ async def add_pilots_bulk(
         pilot_stamps: list[str],
         vo: str,
         grid_type: str = "DIRAC",
-        pilot_references: dict | None = None,
+        pilot_references: dict[str, str] | None = None,
     ):
         """Bulk add pilots in the DB.
 
@@ -85,7 +79,9 @@ async def delete_pilots_by_stamps_bulk(self, pilot_stamps: list[str]):
         if res.rowcount != len(pilot_stamps):
             raise PilotNotFoundError(data={"pilot_stamps": str(pilot_stamps)})
 
-    async def associate_pilot_with_jobs(self, job_to_pilot_mapping: list[dict]):
+    async def associate_pilot_with_jobs(
+        self, job_to_pilot_mapping: list[dict[str, Any]]
+    ):
         """Associate a pilot with jobs.
 
         job_to_pilot_mapping format:
@@ -182,61 +178,28 @@ async def update_pilot_fields_bulk(
                 data={"mapping": str(pilot_stamps_to_fields_mapping)}
             )
 
-    async def get_pilots_by_stamp_bulk(
-        self, pilot_stamps: list[str]
-    ) -> Sequence[RowMapping]:
-        """Bulk fetch pilots.
-
-        Raises PilotNotFoundError if one of the stamp is not associated with a pilot.
-
-        """
-        results = await fetch_records_bulk_or_raises(
-            self.conn,
-            PilotAgents,
-            PilotNotFoundError,
-            "pilot_stamp",
-            "PilotStamp",
-            pilot_stamps,
-            allow_no_result=True,
-        )
-
-        # Custom handling, to see which pilot_stamp does not exist (if so, say which one)
-        found_keys = {row["PilotStamp"] for row in results}
-        missing = set(pilot_stamps) - found_keys
-
-        if missing:
-            raise PilotNotFoundError(
-                data={"pilot_stamp": str(missing)},
-                detail=str(missing),
-                non_existing_pilots=missing,
-            )
-
-        return results
-
-    async def get_pilot_jobs_ids_by_pilot_id(self, pilot_id: int) -> list[int]:
-        """Fetch pilot jobs."""
-        job_to_pilot_mapping = await fetch_records_bulk_or_raises(
-            self.conn,
-            JobToPilotMapping,
-            PilotJobsNotFoundError,
-            "pilot_id",
-            "PilotID",
-            [pilot_id],
-            allow_more_than_one_result_per_input=True,
-            allow_no_result=True,
+    async def search_pilots(
+        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]]]:
+        """Search for pilots in the database."""
+        return await self.search(
+            model=PilotAgents,
+            parameters=parameters,
+            search=search,
+            sorts=sorts,
+            distinct=distinct,
+            per_page=per_page,
+            page=page,
         )
 
-        return [mapping["JobID"] for mapping in job_to_pilot_mapping]
-
-    async def get_pilot_ids_by_stamps(self, pilot_stamps: list[str]) -> list[int]:
-        """Get pilot ids."""
-        # This function is currently needed while we are relying on pilot_ids instead of pilot_stamps
-        # (Ex: JobToPilotMapping)
-        pilots = await self.get_pilots_by_stamp_bulk(pilot_stamps)
-
-        return [pilot["PilotID"] for pilot in pilots]
-
-    async def search(
+    async def search_pilot_to_job_mapping(
         self,
         parameters: list[str] | None,
         search: list[SearchSpec],
@@ -247,39 +210,15 @@ async def search(
         page: int | None = None,
     ) -> tuple[int, list[dict[Any, Any]]]:
         """Search for pilots in the database."""
-        # TODO: Refactorize with the search function for jobs.
-        # Find which columns to select
-        columns = _get_columns(PilotAgents.__table__, parameters)
-
-        stmt = select(*columns)
-
-        stmt = apply_search_filters(
-            PilotAgents.__table__.columns.__getitem__, stmt, search
+        return await self.search(
+            model=JobToPilotMapping,
+            parameters=parameters,
+            search=search,
+            sorts=sorts,
+            distinct=distinct,
+            per_page=per_page,
+            page=page,
         )
-        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))
-        ]
 
     async def clear_pilots_bulk(
         self, cutoff_date: datetime, delete_only_aborted: bool

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

@@ -3,12 +3,11 @@
 from .base import (
     BaseSQLDB,
     SQLDBUnavailableError,
+    _get_columns,
     apply_search_filters,
     apply_sort_constraints,
 )
 from .functions import (
-    _get_columns,
-    fetch_records_bulk_or_raises,
     hash,
     substract_date,
     utcnow,
@@ -24,7 +23,6 @@
     "DateNowColumn",
     "EnumBackedBool",
     "EnumColumn",
-    "fetch_records_bulk_or_raises",
     "hash",
     "NullColumn",
     "substract_date",

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

@@ -8,16 +8,16 @@
 from collections.abc import AsyncIterator
 from contextvars import ContextVar
 from datetime import datetime
-from typing import Self, cast
+from typing import Any, Self, cast
 
 from pydantic import TypeAdapter
-from sqlalchemy import DateTime, MetaData, select
+from sqlalchemy import DateTime, MetaData, func, select
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.extensions import select_from_extension
-from diracx.core.models import SortDirection
+from diracx.core.models import SearchSpec, SortDirection, SortSpec
 from diracx.core.settings import SqlalchemyDsn
 from diracx.db.exceptions import DBUnavailableError
 
@@ -227,6 +227,47 @@ async def ping(self):
         except OperationalError as e:
             raise SQLDBUnavailableError("Cannot ping the DB") from e
 
+    async def search(
+        self,
+        model: Any,
+        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]]]:
+        """Search for pilots in the database."""
+        # Find which columns to select
+        columns = _get_columns(model.__table__, parameters)
+
+        stmt = select(*columns)
+
+        stmt = apply_search_filters(model.__table__.columns.__getitem__, stmt, search)
+        stmt = apply_sort_constraints(model.__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))
+        ]
+
 
 def find_time_resolution(value):
     if isinstance(value, datetime):
@@ -258,6 +299,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: