diracx-db 0.0.1a19__tar.gz → 0.0.1a21__tar.gz

{diracx_db-0.0.1a19 → diracx_db-0.0.1a21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a19
+Version: 0.0.1a21
 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>=2.4
+Requires-Dist: pydantic>=2.10
 Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Provides-Extra: testing
 Requires-Dist: diracx-testing; extra == "testing"

{diracx_db-0.0.1a19 → diracx_db-0.0.1a21}/pyproject.toml

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

diracx_db-0.0.1a21/src/diracx/db/sql/job/db.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import bindparam, delete, func, insert, select, update
+from sqlalchemy.exc import IntegrityError, NoResultFound
+
+if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import BindParameter
+from diracx.core.exceptions import InvalidQueryError, JobNotFound
+from diracx.core.models import (
+    LimitedJobStatusReturn,
+    SearchSpec,
+    SortSpec,
+)
+
+from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
+from .schema import (
+    InputData,
+    JobCommands,
+    JobDBBase,
+    JobJDLs,
+    Jobs,
+)
+
+
+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
+
+    # TODO: this is copied from the DIRAC JobDB
+    # but is overwriten in LHCbDIRAC, so we need
+    # to find a way to make it dynamic
+    jdl2DBParameters = ["JobName", "JobType", "JobGroup"]
+
+    async def summary(self, group_by, search) -> list[dict[str, str | int]]:
+        columns = _get_columns(Jobs.__table__, group_by)
+
+        stmt = select(*columns, func.count(Jobs.JobID).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
+        ]
+
+    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(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))
+        ]
+
+    async def insert_input_data(self, lfns: dict[int, list[str]]):
+        await self.conn.execute(
+            InputData.__table__.insert(),
+            [
+                {
+                    "JobID": job_id,
+                    "LFN": lfn,
+                }
+                for job_id, lfns_ in lfns.items()
+                for lfn in lfns_
+            ],
+        )
+
+    async def setJobAttributes(self, job_id, jobData):
+        """TODO: add myDate and force parameters."""
+        if "Status" in jobData:
+            jobData = jobData | {"LastUpdateTime": datetime.now(tz=timezone.utc)}
+        stmt = update(Jobs).where(Jobs.JobID == job_id).values(jobData)
+        await self.conn.execute(stmt)
+
+    async def create_job(self, original_jdl):
+        """Used to insert a new job with original JDL. Returns inserted job id."""
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
+
+        result = await self.conn.execute(
+            JobJDLs.__table__.insert().values(
+                JDL="",
+                JobRequirements="",
+                OriginalJDL=compressJDL(original_jdl),
+            )
+        )
+        return result.lastrowid
+
+    async def insert_job_attributes(self, jobs_to_update: dict[int, dict]):
+        await self.conn.execute(
+            Jobs.__table__.insert(),
+            [
+                {
+                    "JobID": job_id,
+                    **attrs,
+                }
+                for job_id, attrs in jobs_to_update.items()
+            ],
+        )
+
+    async def update_job_jdls(self, jdls_to_update: dict[int, str]):
+        """Used to update the JDL, typically just after inserting the original JDL, or rescheduling, for example."""
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
+
+        await self.conn.execute(
+            JobJDLs.__table__.update().where(
+                JobJDLs.__table__.c.JobID == bindparam("b_JobID")
+            ),
+            [
+                {
+                    "b_JobID": job_id,
+                    "JDL": compressJDL(jdl),
+                }
+                for job_id, jdl in jdls_to_update.items()
+            ],
+        )
+
+    async def checkAndPrepareJob(
+        self,
+        jobID,
+        class_ad_job,
+        class_ad_req,
+        owner,
+        owner_group,
+        job_attrs,
+        vo,
+    ):
+        """Check Consistency of Submitted JDL and set some defaults
+        Prepare subJDL with Job Requirements.
+        """
+        from DIRAC.Core.Utilities.DErrno import EWMSSUBM, cmpError
+        from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import (
+            checkAndPrepareJob,
+        )
+
+        retVal = checkAndPrepareJob(
+            jobID,
+            class_ad_job,
+            class_ad_req,
+            owner,
+            owner_group,
+            job_attrs,
+            vo,
+        )
+
+        if not retVal["OK"]:
+            if cmpError(retVal, EWMSSUBM):
+                await self.setJobAttributes(jobID, job_attrs)
+
+            returnValueOrRaise(retVal)
+
+    async def setJobJDL(self, job_id, jdl):
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
+
+        stmt = (
+            update(JobJDLs).where(JobJDLs.JobID == job_id).values(JDL=compressJDL(jdl))
+        )
+        await self.conn.execute(stmt)
+
+    async def setJobJDLsBulk(self, jdls):
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import compressJDL
+
+        await self.conn.execute(
+            JobJDLs.__table__.update().where(
+                JobJDLs.__table__.c.JobID == bindparam("b_JobID")
+            ),
+            [{"b_JobID": jid, "JDL": compressJDL(jdl)} for jid, jdl in jdls.items()],
+        )
+
+    async def setJobAttributesBulk(self, jobData):
+        """TODO: add myDate and force parameters."""
+        for job_id in jobData.keys():
+            if "Status" in jobData[job_id]:
+                jobData[job_id].update(
+                    {"LastUpdateTime": datetime.now(tz=timezone.utc)}
+                )
+
+        await self.conn.execute(
+            Jobs.__table__.update().where(
+                Jobs.__table__.c.JobID == bindparam("b_JobID")
+            ),
+            [{"b_JobID": job_id, **attrs} for job_id, attrs in jobData.items()],
+        )
+
+    async def getJobJDL(self, job_id: int, original: bool = False) -> str:
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import extractJDL
+
+        if original:
+            stmt = select(JobJDLs.OriginalJDL).where(JobJDLs.JobID == job_id)
+        else:
+            stmt = select(JobJDLs.JDL).where(JobJDLs.JobID == job_id)
+
+        jdl = (await self.conn.execute(stmt)).scalar_one()
+        if jdl:
+            jdl = extractJDL(jdl)
+
+        return jdl
+
+    async def getJobJDLs(self, job_ids, original: bool = False) -> dict[int | str, str]:
+        from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import extractJDL
+
+        if original:
+            stmt = select(JobJDLs.JobID, JobJDLs.OriginalJDL).where(
+                JobJDLs.JobID.in_(job_ids)
+            )
+        else:
+            stmt = select(JobJDLs.JobID, JobJDLs.JDL).where(JobJDLs.JobID.in_(job_ids))
+
+        return {
+            jobid: extractJDL(jdl)
+            for jobid, jdl in (await self.conn.execute(stmt))
+            if jdl
+        }
+
+    async def get_job_status(self, job_id: int) -> LimitedJobStatusReturn:
+        try:
+            stmt = select(Jobs.Status, Jobs.MinorStatus, Jobs.ApplicationStatus).where(
+                Jobs.JobID == job_id
+            )
+            return LimitedJobStatusReturn(
+                **dict((await self.conn.execute(stmt)).one()._mapping)
+            )
+        except NoResultFound as e:
+            raise JobNotFound(job_id) from e
+
+    async def set_job_command(self, job_id: int, command: str, arguments: str = ""):
+        """Store a command to be passed to the job together with the next heart beat."""
+        try:
+            stmt = insert(JobCommands).values(
+                JobID=job_id,
+                Command=command,
+                Arguments=arguments,
+                ReceptionTime=datetime.now(tz=timezone.utc),
+            )
+            await self.conn.execute(stmt)
+        except IntegrityError as e:
+            raise JobNotFound(job_id) from e
+
+    async def set_job_command_bulk(self, commands):
+        """Store a command to be passed to the job together with the next heart beat."""
+        self.conn.execute(
+            insert(JobCommands),
+            [
+                {
+                    "JobID": job_id,
+                    "Command": command,
+                    "Arguments": arguments,
+                    "ReceptionTime": datetime.now(tz=timezone.utc),
+                }
+                for job_id, command, arguments in commands
+            ],
+        )
+        # FIXME handle IntegrityError
+
+    async def delete_jobs(self, job_ids: list[int]):
+        """Delete jobs from the database."""
+        stmt = delete(JobJDLs).where(JobJDLs.JobID.in_(job_ids))
+        await self.conn.execute(stmt)
+
+    async def set_properties(
+        self, properties: dict[int, dict[str, Any]], update_timestamp: bool = False
+    ) -> int:
+        """Update the job parameters
+        All the jobs must update the same properties.
+
+        :param properties: {job_id : {prop1: val1, prop2:val2}
+        :param update_timestamp: if True, update the LastUpdate to now
+
+        :return rowcount
+
+        """
+        # Check that all we always update the same set of properties
+        required_parameters_set = {tuple(sorted(k.keys())) for k in properties.values()}
+
+        if len(required_parameters_set) != 1:
+            raise NotImplementedError(
+                "All the jobs should update the same set of 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)
+        values: dict[str, BindParameter[Any] | datetime] = {
+            c.name: bindparam(c.name) for c in columns
+        }
+        if update_timestamp:
+            values["LastUpdateTime"] = datetime.now(tz=timezone.utc)
+
+        stmt = update(Jobs).where(Jobs.JobID == bindparam("job_id")).values(**values)
+        rows = await self.conn.execute(stmt, update_parameters)
+
+        return rows.rowcount

{diracx_db-0.0.1a19 → diracx_db-0.0.1a21}/src/diracx/db/sql/job_logging/db.py

@@ -4,11 +4,14 @@ import time
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING
 
+from pydantic import BaseModel
 from sqlalchemy import delete, func, insert, select
 
 if TYPE_CHECKING:
     pass
 
+from collections import defaultdict
+
 from diracx.core.exceptions import JobNotFound
 from diracx.core.models import (
     JobStatus,
@@ -24,6 +27,15 @@ from .schema import (
 MAGIC_EPOC_NUMBER = 1270000000
 
 
+class JobLoggingRecord(BaseModel):
+    job_id: int
+    status: JobStatus
+    minor_status: str
+    application_status: str
+    date: datetime
+    source: str
+
+
 class JobLoggingDB(BaseSQLDB):
     """Frontend for the JobLoggingDB. Provides the ability to store changes with timestamps."""
 
@@ -69,6 +81,49 @@ class JobLoggingDB(BaseSQLDB):
         )
         await self.conn.execute(stmt)
 
+    async def bulk_insert_record(
+        self,
+        records: list[JobLoggingRecord],
+    ):
+        """Bulk insert entries to the JobLoggingDB table."""
+
+        def get_epoc(date):
+            return (
+                time.mktime(date.timetuple())
+                + date.microsecond / 1000000.0
+                - MAGIC_EPOC_NUMBER
+            )
+
+        # First, fetch the maximum SeqNums for the given job_ids
+        seqnum_stmt = (
+            select(
+                LoggingInfo.JobID, func.coalesce(func.max(LoggingInfo.SeqNum) + 1, 1)
+            )
+            .where(LoggingInfo.JobID.in_([record.job_id for record in records]))
+            .group_by(LoggingInfo.JobID)
+        )
+
+        seqnum = {jid: seqnum for jid, seqnum in (await self.conn.execute(seqnum_stmt))}
+        # IF a seqnum is not found, then assume it does not exist and the first sequence number is 1.
+
+        # https://docs.sqlalchemy.org/en/20/orm/queryguide/dml.html#orm-bulk-insert-statements
+        await self.conn.execute(
+            insert(LoggingInfo),
+            [
+                {
+                    "JobID": record.job_id,
+                    "SeqNum": seqnum.get(record.job_id, 1),
+                    "Status": record.status,
+                    "MinorStatus": record.minor_status,
+                    "ApplicationStatus": record.application_status[:255],
+                    "StatusTime": record.date,
+                    "StatusTimeOrder": get_epoc(record.date),
+                    "Source": record.source[:32],
+                }
+                for record in records
+            ],
+        )
+
     async def get_records(self, job_id: int) -> list[JobStatusReturn]:
         """Returns a Status,MinorStatus,ApplicationStatus,StatusTime,Source tuple
         for each record found for job specified by its jobID in historical order.
@@ -159,3 +214,22 @@ class JobLoggingDB(BaseSQLDB):
             result[event] = str(etime + MAGIC_EPOC_NUMBER)
 
         return result
+
+    async def get_wms_time_stamps_bulk(self, job_ids):
+        """Get TimeStamps for job MajorState transitions for multiple jobs at once
+        return a {JobID: {State:timestamp}} dictionary.
+        """
+        result = defaultdict(dict)
+        stmt = select(
+            LoggingInfo.JobID,
+            LoggingInfo.Status,
+            LoggingInfo.StatusTimeOrder,
+        ).where(LoggingInfo.JobID.in_(job_ids))
+        rows = await self.conn.execute(stmt)
+        if not rows.rowcount:
+            return {}
+
+        for job_id, event, etime in rows:
+            result[job_id][event] = str(etime + MAGIC_EPOC_NUMBER)
+
+        return result

{diracx_db-0.0.1a19 → diracx_db-0.0.1a21}/src/diracx/db/sql/utils/__init__.py

@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Self, cast
 import sqlalchemy.types as types
 from pydantic import TypeAdapter
 from sqlalchemy import Column as RawColumn
-from sqlalchemy import DateTime, Enum, MetaData, select
+from sqlalchemy import DateTime, Enum, MetaData, func, select
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 from sqlalchemy.ext.compiler import compiles
@@ -100,7 +100,9 @@ def mysql_date_trunc(element, compiler, **kw):
         "MONTH": "%Y-%m",
         "YEAR": "%Y",
     }[element._time_resolution]
-    return f"DATE_FORMAT({compiler.process(element.clauses)}, '{pattern}')"
+
+    (dt_col,) = list(element.clauses)
+    return compiler.process(func.date_format(dt_col, pattern))
 
 
 @compiles(date_trunc, "sqlite")
@@ -113,7 +115,13 @@ def sqlite_date_trunc(element, compiler, **kw):
         "MONTH": "%Y-%m",
         "YEAR": "%Y",
     }[element._time_resolution]
-    return f"strftime('{pattern}', {compiler.process(element.clauses)})"
+    (dt_col,) = list(element.clauses)
+    return compiler.process(
+        func.strftime(
+            pattern,
+            dt_col,
+        )
+    )
 
 
 def substract_date(**kwargs: float) -> datetime: