diracx-db 0.0.1a45__tar.gz → 0.0.1a47__tar.gz

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-db
-Version: 0.0.1a45
+Version: 0.0.1a47
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/__main__.py

@@ -31,7 +31,6 @@ async def init_sql():
     from diracx.db.sql.utils import BaseSQLDB
 
     for db_name, db_url in BaseSQLDB.available_urls().items():
-
         logger.info("Initialising %s", db_name)
         db = BaseSQLDB.available_implementations(db_name)[0](db_url)
         async with db.engine_context():

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/os/job_parameters.py

@@ -9,13 +9,19 @@ class JobParametersDB(BaseOSDB):
     fields = {
         "JobID": {"type": "long"},
         "timestamp": {"type": "date"},
+        "PilotAgent": {"type": "keyword"},
+        "Pilot_Reference": {"type": "keyword"},
+        "JobGroup": {"type": "keyword"},
         "CPUNormalizationFactor": {"type": "long"},
         "NormCPUTime(s)": {"type": "long"},
-        "Memory(kB)": {"type": "long"},
+        "Memory(MB)": {"type": "long"},
+        "LocalAccount": {"type": "keyword"},
         "TotalCPUTime(s)": {"type": "long"},
-        "MemoryUsed(kb)": {"type": "long"},
-        "HostName": {"type": "keyword"},
+        "PayloadPID": {"type": "long"},
+        "HostName": {"type": "text"},
         "GridCE": {"type": "keyword"},
+        "CEQueue": {"type": "keyword"},
+        "BatchSystem": {"type": "keyword"},
         "ModelName": {"type": "keyword"},
         "Status": {"type": "keyword"},
         "JobType": {"type": "keyword"},

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/dummy/db.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from sqlalchemy import func, insert, select
+from sqlalchemy import insert
 from uuid_utils import UUID
 
-from diracx.db.sql.utils import BaseSQLDB, apply_search_filters
+from diracx.db.sql.utils import BaseSQLDB
 
 from .schema import Base as DummyDBBase
 from .schema import Cars, Owners
@@ -22,18 +22,7 @@ class DummyDB(BaseSQLDB):
     metadata = DummyDBBase.metadata
 
     async def summary(self, group_by, search) -> list[dict[str, str | int]]:
-        columns = [Cars.__table__.columns[x] for x in group_by]
-
-        stmt = select(*columns, func.count(Cars.license_plate).label("count"))
-        stmt = apply_search_filters(Cars.__table__.columns.__getitem__, stmt, search)
-        stmt = stmt.group_by(*columns)
-
-        # Execute the query
-        return [
-            dict(row._mapping)
-            async for row in (await self.conn.stream(stmt))
-            if row.count > 0  # type: ignore
-        ]
+        return await self._summary(Cars, group_by, search)
 
     async def insert_owner(self, name: str) -> int:
         stmt = insert(Owners).values(name=name)

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/job/db.py

@@ -5,7 +5,7 @@ __all__ = ["JobDB"]
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Iterable
 
-from sqlalchemy import bindparam, case, delete, func, insert, select, update
+from sqlalchemy import bindparam, case, delete, insert, select, update
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BindParameter
@@ -13,7 +13,7 @@ if TYPE_CHECKING:
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.models import JobCommand, SearchSpec, SortSpec
 
-from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
+from ..utils import BaseSQLDB, _get_columns
 from ..utils.functions import utcnow
 from .schema import (
     HeartBeatLoggingInfo,
@@ -25,17 +25,6 @@ from .schema import (
 )
 
 
-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
 
@@ -54,20 +43,11 @@ class JobDB(BaseSQLDB):
     # to find a way to make it dynamic
     jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
 
-    async def summary(self, group_by, search) -> list[dict[str, str | int]]:
+    async def summary(
+        self, group_by: list[str], search: list[SearchSpec]
+    ) -> list[dict[str, str | int]]:
         """Get a summary of the jobs."""
-        columns = _get_columns(Jobs.__table__, group_by)
-
-        stmt = select(*columns, func.count(Jobs.job_id).label("count"))
-        stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
-        stmt = stmt.group_by(*columns)
-
-        # Execute the query
-        return [
-            dict(row._mapping)
-            async for row in (await self.conn.stream(stmt))
-            if row.count > 0  # type: ignore
-        ]
+        return await self._summary(table=Jobs, group_by=group_by, search=search)
 
     async def search(
         self,
@@ -80,34 +60,15 @@ class JobDB(BaseSQLDB):
         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(
+            table=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-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/job_logging/db.py

@@ -89,7 +89,6 @@ class JobLoggingDB(BaseSQLDB):
             status_time,
             status_source,
         ) in rows:
-
             values[job_id].append(
                 [
                     status,

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/pilot_agents/db.py

@@ -20,7 +20,6 @@ class PilotAgentsDB(BaseSQLDB):
         grid_type: str = "DIRAC",
         pilot_stamps: dict | None = None,
     ) -> None:
-
         if pilot_stamps is None:
             pilot_stamps = {}
 

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/sandbox_metadata/db.py

@@ -13,6 +13,7 @@ from sqlalchemy import (
     Table,
     and_,
     delete,
+    exists,
     insert,
     literal,
     or_,
@@ -236,7 +237,10 @@ class SandboxMetadataDB(BaseSQLDB):
         """
         conditions = [
             # If it has assigned to a job but is no longer mapped it can be removed
-            # and_(SandBoxes.Assigned, ~exists(SandBoxes.SBId == SBEntityMapping.SBId)),
+            and_(
+                SandBoxes.Assigned,
+                ~exists().where(SBEntityMapping.SBId == SandBoxes.SBId),
+            ),
             # If the sandbox is still unassigned after 15 days, remove it
             and_(~SandBoxes.Assigned, days_since(SandBoxes.LastAccessTime) >= 15),
         ]

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/utils/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from .base import (
     BaseSQLDB,
     SQLDBUnavailableError,
+    _get_columns,
     apply_search_filters,
     apply_sort_constraints,
 )
@@ -10,6 +11,7 @@ from .functions import hash, substract_date, utcnow
 from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn
 
 __all__ = (
+    "_get_columns",
     "utcnow",
     "Column",
     "NullColumn",

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/src/diracx/db/sql/utils/base.py

@@ -8,16 +8,20 @@ from abc import ABCMeta
 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 +231,71 @@ class BaseSQLDB(metaclass=ABCMeta):
         except OperationalError as e:
             raise SQLDBUnavailableError("Cannot ping the DB") from e
 
+    async def _search(
+        self,
+        table: 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[str, Any]]]:
+        """Search for elements in a table."""
+        # Find which columns to select
+        columns = _get_columns(table.__table__, parameters)
+
+        stmt = select(*columns)
+
+        stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
+        stmt = apply_sort_constraints(table.__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 _summary(
+        self, table: Any, group_by: list[str], search: list[SearchSpec]
+    ) -> list[dict[str, str | int]]:
+        """Get a summary of the elements of a table."""
+        columns = _get_columns(table.__table__, group_by)
+
+        pk_columns = list(table.__table__.primary_key.columns)
+        if not pk_columns:
+            raise ValueError(
+                "Model has no primary key and no count_column was provided."
+            )
+        count_col = pk_columns[0]
+
+        stmt = select(*columns, func.count(count_col).label("count"))
+        stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
+        stmt = stmt.group_by(*columns)
+
+        # Execute the query
+        return [
+            dict(row._mapping)
+            async for row in (await self.conn.stream(stmt))
+            if row.count > 0  # type: ignore
+        ]
+
 
 def find_time_resolution(value):
     if isinstance(value, datetime):
@@ -258,6 +327,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:
@@ -300,6 +380,15 @@ def apply_search_filters(column_mapping, stmt, search):
             expr = column.like(query["value"])
         elif query["operator"] in "ilike":
             expr = column.ilike(query["value"])
+        elif query["operator"] == "not like":
+            expr = column.not_like(query["value"])
+        elif query["operator"] == "regex":
+            # We check the regex validity here
+            try:
+                re.compile(query["value"])
+            except re.error as e:
+                raise InvalidQueryError(f"Invalid regex {query['value']}") from e
+            expr = column.regexp_match(query["value"])
         else:
             raise InvalidQueryError(f"Unknown filter {query=}")
         stmt = stmt.where(expr)

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/tests/jobs/test_job_db.py

@@ -187,6 +187,55 @@ async def test_search_conditions(populated_job_db):
         assert total == 0
         assert not result
 
+        # Search for a specific scalar condition: Owner not like 'owner1%'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.NOT_LIKE, value="owner1%"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100 - 11
+        assert result
+        assert len(result) == 100 - 11
+        assert all(not r["Owner"].startswith("owner1") for r in result)
+
+        # Search for a specific scalar condition: OwnerGroup not like 'owner_group2'
+        condition = ScalarSearchSpec(
+            parameter="OwnerGroup",
+            operator=ScalarSearchOperator.NOT_LIKE,
+            value="owner_group2",
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100 - 50
+        assert result
+        assert len(result) == 100 - 50
+        assert all(not r["OwnerGroup"] == "owner_group2" for r in result)
+
+        # Search for a specific scalar condition: Owner regex '^owner\d+$'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="^owner\\d+$"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100
+        assert result
+        assert len(result) == 100
+
+        # Search for a specific scalar condition: JobID regex 'owner[0-3]+'
+        # owner0, owner1, owner2, owner3 (4 jobs)
+        # owner11 -> owner39 (30 jobs)
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="owner[0-3]+"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 34
+        assert result
+        assert len(result) == 34
+
+        # Search for a specific scalar condition: JobID regex 'owner[1-'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="owner[1-"
+        )
+        with pytest.raises(InvalidQueryError):
+            await job_db.search([], [condition], [])
+
 
 async def test_search_sorts(populated_job_db):
     """Test that we can search for jobs in the database and sort the results."""

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/tests/opensearch/test_index_template.py

@@ -37,13 +37,15 @@ async def _get_test_index_mappings(dummy_opensearch_db: DummyOSDB):
 
     # At this point the index should not exist yet
     with pytest.raises(opensearchpy.exceptions.NotFoundError):
-        await dummy_opensearch_db.client.indices.get_mapping(index_name)
+        await dummy_opensearch_db.client.indices.get_mapping(index=index_name)
 
     # Insert document which will automatically create the index based on the template
     await dummy_opensearch_db.upsert(vo, document_id, DUMMY_DOCUMENT)
 
     # Ensure the result looks as expected and return the mappings
-    index_mapping = await dummy_opensearch_db.client.indices.get_mapping(index_name)
+    index_mapping = await dummy_opensearch_db.client.indices.get_mapping(
+        index=index_name
+    )
     assert list(index_mapping) == [index_name]
     assert list(index_mapping[index_name]) == ["mappings"]
     return index_mapping[index_name]["mappings"]

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/tests/pilot_agents/test_pilot_agents_db.py

@@ -15,7 +15,6 @@ async def pilot_agents_db(tmp_path) -> PilotAgentsDB:
 
 
 async def test_insert_and_select(pilot_agents_db: PilotAgentsDB):
-
     async with pilot_agents_db as pilot_agents_db:
         # Add a pilot reference
         refs = [f"ref_{i}" for i in range(10)]

{diracx_db-0.0.1a45 → diracx_db-0.0.1a47}/tests/test_dummy_db.py

@@ -129,7 +129,7 @@ async def test_failed_transaction(dummy_db):
 
     # The connection is created when the context manager is entered
     # This is our transaction
-    with pytest.raises(KeyError):
+    with pytest.raises(InvalidQueryError):
         async with dummy_db as dummy_db:
             assert dummy_db.conn