diracx-db 0.0.1a15__tar.gz → 0.0.1a17__tar.gz

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a15
+Version: 0.0.1a17
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -14,7 +14,7 @@ Requires-Dist: dirac
 Requires-Dist: diracx-core
 Requires-Dist: fastapi
 Requires-Dist: opensearch-py[async]
-Requires-Dist: pydantic
-Requires-Dist: sqlalchemy[aiomysql,aiosqlite]
+Requires-Dist: pydantic>=2.4
+Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Provides-Extra: testing
 Requires-Dist: diracx-testing; extra == "testing"

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/pyproject.toml

@@ -17,8 +17,8 @@ dependencies = [
     "diracx-core",
     "fastapi",
     "opensearch-py[async]",
-    "pydantic",
-    "sqlalchemy[aiomysql,aiosqlite]",
+    "pydantic >=2.4",
+    "sqlalchemy[aiomysql,aiosqlite] >= 2",
 ]
 dynamic = ["version"]
 

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/src/diracx/db/sql/jobs/db.py

@@ -19,6 +19,9 @@ from diracx.core.models import (
     LimitedJobStatusReturn,
     ScalarSearchOperator,
     ScalarSearchSpec,
+    SearchSpec,
+    SortDirection,
+    SortSpec,
 )
 from diracx.core.properties import JOB_SHARING, SecurityProperty
 
@@ -83,14 +86,14 @@ class JobDB(BaseSQLDB):
 
     async def search(
         self,
-        parameters,
-        search,
-        sorts,
+        parameters: list[str] | None,
+        search: list[SearchSpec],
+        sorts: list[SortSpec],
         *,
         distinct: bool = False,
         per_page: int = 100,
         page: int | None = None,
-    ) -> list[dict[str, Any]]:
+    ) -> tuple[int, list[dict[Any, Any]]]:
         # Find which columns to select
         columns = _get_columns(Jobs.__table__, parameters)
         stmt = select(*columns)
@@ -98,28 +101,45 @@ class JobDB(BaseSQLDB):
         stmt = apply_search_filters(Jobs.__table__, stmt, search)
 
         # Apply any sort constraints
+        sort_columns = []
         for sort in sorts:
             if sort["parameter"] not in Jobs.__table__.columns:
                 raise InvalidQueryError(
                     f"Cannot sort by {sort['parameter']}: unknown column"
                 )
             column = Jobs.__table__.columns[sort["parameter"]]
-            if sort["direction"] == "asc":
-                column = column.asc()
-            elif sort["direction"] == "desc":
-                column = column.desc()
+            sorted_column = None
+            if sort["direction"] == SortDirection.ASC:
+                sorted_column = column.asc()
+            elif sort["direction"] == SortDirection.DESC:
+                sorted_column = column.desc()
             else:
                 raise InvalidQueryError(f"Unknown sort {sort['direction']=}")
+            sort_columns.append(sorted_column)
+
+        if sort_columns:
+            stmt = stmt.order_by(*sort_columns)
 
         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:
-            raise NotImplementedError("TODO Not yet implemented")
+        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 [dict(row._mapping) async for row in (await self.conn.stream(stmt))]
+        return total, [
+            dict(row._mapping) async for row in (await self.conn.stream(stmt))
+        ]
 
     async def _insertNewJDL(self, jdl) -> int:
         from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
@@ -314,7 +334,7 @@ class JobDB(BaseSQLDB):
         from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
         from DIRAC.Core.Utilities.ReturnValues import SErrorException
 
-        result = await self.search(
+        _, result = await self.search(
             parameters=[
                 "Status",
                 "MinorStatus",

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/src/diracx/db/sql/jobs/status_utility.py

@@ -41,9 +41,9 @@ async def set_job_status(
     # transform JobStateUpdate objects into dicts
     statusDict = {}
     for key, value in status.items():
-        statusDict[key] = {k: v for k, v in value.dict().items() if v is not None}
+        statusDict[key] = {k: v for k, v in value.model_dump().items() if v is not None}
 
-    res = await job_db.search(
+    _, res = await job_db.search(
         parameters=["Status", "StartExecTime", "EndExecTime"],
         search=[
             {

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/src/diracx/db/sql/utils.py

@@ -11,7 +11,7 @@ from datetime import datetime, timedelta, timezone
 from functools import partial
 from typing import TYPE_CHECKING, AsyncIterator, Self, cast
 
-from pydantic import parse_obj_as
+from pydantic import TypeAdapter
 from sqlalchemy import Column as RawColumn
 from sqlalchemy import DateTime, Enum, MetaData, select
 from sqlalchemy.exc import OperationalError
@@ -123,7 +123,9 @@ class BaseSQLDB(metaclass=ABCMeta):
                     if db_url == "sqlite+aiosqlite:///:memory:":
                         db_urls[db_name] = db_url
                     else:
-                        db_urls[db_name] = parse_obj_as(SqlalchemyDsn, db_url)
+                        db_urls[db_name] = str(
+                            TypeAdapter(SqlalchemyDsn).validate_python(db_url)
+                        )
                 except Exception:
                     logger.error("Error loading URL for %s", db_name)
                     raise

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/src/diracx_db.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a15
+Version: 0.0.1a17
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -14,7 +14,7 @@ Requires-Dist: dirac
 Requires-Dist: diracx-core
 Requires-Dist: fastapi
 Requires-Dist: opensearch-py[async]
-Requires-Dist: pydantic
-Requires-Dist: sqlalchemy[aiomysql,aiosqlite]
+Requires-Dist: pydantic>=2.4
+Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Provides-Extra: testing
 Requires-Dist: diracx-testing; extra == "testing"

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/src/diracx_db.egg-info/requires.txt

@@ -2,8 +2,8 @@ dirac
 diracx-core
 fastapi
 opensearch-py[async]
-pydantic
-sqlalchemy[aiomysql,aiosqlite]
+pydantic>=2.4
+sqlalchemy[aiomysql,aiosqlite]>=2
 
 [testing]
 diracx-testing

diracx_db-0.0.1a17/tests/jobs/test_jobDB.py

@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from diracx.core.exceptions import InvalidQueryError, JobNotFound
+from diracx.core.models import (
+    ScalarSearchOperator,
+    ScalarSearchSpec,
+    SortDirection,
+    SortSpec,
+    VectorSearchOperator,
+    VectorSearchSpec,
+)
+from diracx.db.sql.jobs.db import JobDB
+
+
+@pytest.fixture
+async def job_db(tmp_path):
+    job_db = JobDB("sqlite+aiosqlite:///:memory:")
+    async with job_db.engine_context():
+        async with job_db.engine.begin() as conn:
+            # set PRAGMA foreign_keys=ON if sqlite
+            if job_db._db_url.startswith("sqlite"):
+                await conn.exec_driver_sql("PRAGMA foreign_keys=ON")
+            await conn.run_sync(job_db.metadata.create_all)
+        yield job_db
+
+
+async def test_search_parameters(job_db):
+    """Test that we can search specific parameters for jobs in the database."""
+    async with job_db as job_db:
+        total, result = await job_db.search(["JobID"], [], [])
+        assert total == 0
+        assert not result
+
+        result = await asyncio.gather(
+            *(
+                job_db.insert(
+                    f"JDL{i}",
+                    "owner",
+                    "owner_group",
+                    "New",
+                    "dfdfds",
+                    "lhcb",
+                )
+                for i in range(100)
+            )
+        )
+
+    async with job_db as job_db:
+        # Search a specific parameter: JobID
+        total, result = await job_db.search(["JobID"], [], [])
+        assert total == 100
+        assert result
+        for r in result:
+            assert r.keys() == {"JobID"}
+
+        # Search a specific parameter: Status
+        total, result = await job_db.search(["Status"], [], [])
+        assert total == 100
+        assert result
+        for r in result:
+            assert r.keys() == {"Status"}
+
+        # Search for multiple parameters: JobID, Status
+        total, result = await job_db.search(["JobID", "Status"], [], [])
+        assert total == 100
+        assert result
+        for r in result:
+            assert r.keys() == {"JobID", "Status"}
+
+        # Search for a specific parameter but use distinct: Status
+        total, result = await job_db.search(["Status"], [], [], distinct=True)
+        assert total == 1
+        assert result
+
+        # Search for a non-existent parameter: Dummy
+        with pytest.raises(InvalidQueryError):
+            total, result = await job_db.search(["Dummy"], [], [])
+
+
+async def test_search_conditions(job_db):
+    """Test that we can search for specific jobs in the database."""
+    async with job_db as job_db:
+        result = await asyncio.gather(
+            *(
+                job_db.insert(
+                    f"JDL{i}",
+                    f"owner{i}",
+                    "owner_group",
+                    "New",
+                    "dfdfds",
+                    "lhcb",
+                )
+                for i in range(100)
+            )
+        )
+
+    async with job_db as job_db:
+        # Search a specific scalar condition: JobID eq 3
+        condition = ScalarSearchSpec(
+            parameter="JobID", operator=ScalarSearchOperator.EQUAL, value=3
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 1
+        assert result
+        assert len(result) == 1
+        assert result[0]["JobID"] == 3
+
+        # Search a specific scalar condition: JobID lt 3
+        condition = ScalarSearchSpec(
+            parameter="JobID", operator=ScalarSearchOperator.LESS_THAN, value=3
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 2
+        assert result
+        assert len(result) == 2
+        assert result[0]["JobID"] == 1
+        assert result[1]["JobID"] == 2
+
+        # Search a specific scalar condition: JobID neq 3
+        condition = ScalarSearchSpec(
+            parameter="JobID", operator=ScalarSearchOperator.NOT_EQUAL, value=3
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 99
+        assert result
+        assert len(result) == 99
+        assert all(r["JobID"] != 3 for r in result)
+
+        # Search a specific scalar condition: JobID eq 5873 (does not exist)
+        condition = ScalarSearchSpec(
+            parameter="JobID", operator=ScalarSearchOperator.EQUAL, value=5873
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert not result
+
+        # Search a specific vector condition: JobID in 1,2,3
+        condition = VectorSearchSpec(
+            parameter="JobID", operator=VectorSearchOperator.IN, values=[1, 2, 3]
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 3
+        assert result
+        assert len(result) == 3
+        assert all(r["JobID"] in [1, 2, 3] for r in result)
+
+        # Search a specific vector condition: JobID in 1,2,5873 (one of them does not exist)
+        condition = VectorSearchSpec(
+            parameter="JobID", operator=VectorSearchOperator.IN, values=[1, 2, 5873]
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 2
+        assert result
+        assert len(result) == 2
+        assert all(r["JobID"] in [1, 2] for r in result)
+
+        # Search for multiple conditions based on different parameters: JobID eq 70, JobID in 4,5,6
+        condition1 = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.EQUAL, value="owner4"
+        )
+        condition2 = VectorSearchSpec(
+            parameter="JobID", operator=VectorSearchOperator.IN, values=[4, 5, 6]
+        )
+        total, result = await job_db.search([], [condition1, condition2], [])
+        assert total == 1
+        assert result
+        assert len(result) == 1
+        assert result[0]["JobID"] == 5
+        assert result[0]["Owner"] == "owner4"
+
+        # Search for multiple conditions based on the same parameter: JobID eq 70, JobID in 4,5,6
+        condition1 = ScalarSearchSpec(
+            parameter="JobID", operator=ScalarSearchOperator.EQUAL, value=70
+        )
+        condition2 = VectorSearchSpec(
+            parameter="JobID", operator=VectorSearchOperator.IN, values=[4, 5, 6]
+        )
+        total, result = await job_db.search([], [condition1, condition2], [])
+        assert total == 0
+        assert not result
+
+
+async def test_search_sorts(job_db):
+    """Test that we can search for jobs in the database and sort the results."""
+    async with job_db as job_db:
+        result = await asyncio.gather(
+            *(
+                job_db.insert(
+                    f"JDL{i}",
+                    f"owner{i}",
+                    "owner_group1" if i < 50 else "owner_group2",
+                    "New",
+                    "dfdfds",
+                    "lhcb",
+                )
+                for i in range(100)
+            )
+        )
+
+    async with job_db as job_db:
+        # Search and sort by JobID in ascending order
+        sort = SortSpec(parameter="JobID", direction=SortDirection.ASC)
+        total, result = await job_db.search([], [], [sort])
+        assert total == 100
+        assert result
+        for i, r in enumerate(result):
+            assert r["JobID"] == i + 1
+
+        # Search and sort by JobID in descending order
+        sort = SortSpec(parameter="JobID", direction=SortDirection.DESC)
+        total, result = await job_db.search([], [], [sort])
+        assert total == 100
+        assert result
+        for i, r in enumerate(result):
+            assert r["JobID"] == 100 - i
+
+        # Search and sort by Owner in ascending order
+        sort = SortSpec(parameter="Owner", direction=SortDirection.ASC)
+        total, result = await job_db.search([], [], [sort])
+        assert total == 100
+        assert result
+        # Assert that owner10 is before owner2 because of the lexicographical order
+        assert result[2]["Owner"] == "owner10"
+        assert result[12]["Owner"] == "owner2"
+
+        # Search and sort by Owner in descending order
+        sort = SortSpec(parameter="Owner", direction=SortDirection.DESC)
+        total, result = await job_db.search([], [], [sort])
+        assert total == 100
+        assert result
+        # Assert that owner10 is before owner2 because of the lexicographical order
+        assert result[97]["Owner"] == "owner10"
+        assert result[87]["Owner"] == "owner2"
+
+        # Search and sort by OwnerGroup in ascending order and JobID in descending order
+        sort1 = SortSpec(parameter="OwnerGroup", direction=SortDirection.ASC)
+        sort2 = SortSpec(parameter="JobID", direction=SortDirection.DESC)
+        total, result = await job_db.search([], [], [sort1, sort2])
+        assert total == 100
+        assert result
+        assert result[0]["OwnerGroup"] == "owner_group1"
+        assert result[0]["JobID"] == 50
+        assert result[99]["OwnerGroup"] == "owner_group2"
+        assert result[99]["JobID"] == 51
+
+
+async def test_search_pagination(job_db):
+    """Test that we can search for jobs in the database."""
+    async with job_db as job_db:
+        result = await asyncio.gather(
+            *(
+                job_db.insert(
+                    f"JDL{i}",
+                    f"owner{i}",
+                    "owner_group1" if i < 50 else "owner_group2",
+                    "New",
+                    "dfdfds",
+                    "lhcb",
+                )
+                for i in range(100)
+            )
+        )
+
+    async with job_db as job_db:
+        # Search for the first 10 jobs
+        total, result = await job_db.search([], [], [], per_page=10, page=1)
+        assert total == 100
+        assert result
+        assert len(result) == 10
+        assert result[0]["JobID"] == 1
+
+        # Search for the second 10 jobs
+        total, result = await job_db.search([], [], [], per_page=10, page=2)
+        assert total == 100
+        assert result
+        assert len(result) == 10
+        assert result[0]["JobID"] == 11
+
+        # Search for the last 10 jobs
+        total, result = await job_db.search([], [], [], per_page=10, page=10)
+        assert total == 100
+        assert result
+        assert len(result) == 10
+        assert result[0]["JobID"] == 91
+
+        # Search for the second 50 jobs
+        total, result = await job_db.search([], [], [], per_page=50, page=2)
+        assert total == 100
+        assert result
+        assert len(result) == 50
+        assert result[0]["JobID"] == 51
+
+        # Invalid page number
+        total, result = await job_db.search([], [], [], per_page=10, page=11)
+        assert total == 100
+        assert not result
+
+        # Invalid page number
+        with pytest.raises(InvalidQueryError):
+            result = await job_db.search([], [], [], per_page=10, page=0)
+
+        # Invalid per_page number
+        with pytest.raises(InvalidQueryError):
+            result = await job_db.search([], [], [], per_page=0, page=1)
+
+
+async def test_set_job_command_invalid_job_id(job_db: JobDB):
+    """Test that setting a command for a non-existent job raises JobNotFound."""
+    async with job_db as job_db:
+        with pytest.raises(JobNotFound):
+            await job_db.set_job_command(123456, "test_command")

{diracx_db-0.0.1a15 → diracx_db-0.0.1a17}/tests/test_sandbox_metadata.py

@@ -26,13 +26,16 @@ def test_get_pfn(sandbox_metadata_db: SandboxMetadataDB):
         sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
     )
     sandbox_info = SandboxInfo(
-        checksum="checksum",
+        checksum="90e0ba6763c91a905bb9fd6e025aac1952ae742e6d756a31a0963aa7df7cd7b1",
         checksum_algorithm="sha256",
         format="tar.bz2",
         size=100,
     )
     pfn = sandbox_metadata_db.get_pfn("bucket1", user_info, sandbox_info)
-    assert pfn == "/S3/bucket1/vo/group1/user1/sha256:checksum.tar.bz2"
+    assert pfn == (
+        "/S3/bucket1/vo/group1/user1/"
+        "sha256:90e0ba6763c91a905bb9fd6e025aac1952ae742e6d756a31a0963aa7df7cd7b1.tar.bz2"
+    )
 
 
 async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):

diracx_db-0.0.1a15/tests/jobs/test_jobDB.py

@@ -1,50 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-
-import pytest
-
-from diracx.core.exceptions import JobNotFound
-from diracx.db.sql.jobs.db import JobDB
-
-
-@pytest.fixture
-async def job_db(tmp_path):
-    job_db = JobDB("sqlite+aiosqlite:///:memory:")
-    async with job_db.engine_context():
-        async with job_db.engine.begin() as conn:
-            # set PRAGMA foreign_keys=ON if sqlite
-            if job_db._db_url.startswith("sqlite"):
-                await conn.exec_driver_sql("PRAGMA foreign_keys=ON")
-            await conn.run_sync(job_db.metadata.create_all)
-        yield job_db
-
-
-async def test_some_asyncio_code(job_db):
-    async with job_db as job_db:
-        result = await job_db.search(["JobID"], [], [])
-        assert not result
-
-        result = await asyncio.gather(
-            *(
-                job_db.insert(
-                    f"JDL{i}",
-                    "owner",
-                    "owner_group",
-                    "New",
-                    "dfdfds",
-                    "lhcb",
-                )
-                for i in range(100)
-            )
-        )
-
-    async with job_db as job_db:
-        result = await job_db.search(["JobID"], [], [])
-        assert result
-
-
-async def test_set_job_command_invalid_job_id(job_db: JobDB):
-    async with job_db as job_db:
-        with pytest.raises(JobNotFound):
-            await job_db.set_job_command(123456, "test_command")