diracx-db 0.0.1a16__py3-none-any.whl → 0.0.1a18__py3-none-any.whl

diracx/db/sql/{jobs → job}/db.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import time
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
 
@@ -15,34 +14,20 @@ from diracx.core.exceptions import InvalidQueryError, JobNotFound
 from diracx.core.models import (
     JobMinorStatus,
     JobStatus,
-    JobStatusReturn,
     LimitedJobStatusReturn,
     ScalarSearchOperator,
     ScalarSearchSpec,
     SearchSpec,
-    SortDirection,
     SortSpec,
 )
-from diracx.core.properties import JOB_SHARING, SecurityProperty
 
-from ..utils import BaseSQLDB, apply_search_filters
+from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
 from .schema import (
-    BannedSitesQueue,
-    GridCEsQueue,
     InputData,
     JobCommands,
     JobDBBase,
     JobJDLs,
-    JobLoggingDBBase,
     Jobs,
-    JobsQueue,
-    JobTypesQueue,
-    LoggingInfo,
-    PlatformsQueue,
-    SitesQueue,
-    TagsQueue,
-    TaskQueueDBBase,
-    TaskQueues,
 )
 
 
@@ -74,7 +59,7 @@ class JobDB(BaseSQLDB):
         columns = _get_columns(Jobs.__table__, group_by)
 
         stmt = select(*columns, func.count(Jobs.JobID).label("count"))
-        stmt = apply_search_filters(Jobs.__table__, stmt, search)
+        stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
         stmt = stmt.group_by(*columns)
 
         # Execute the query
@@ -98,27 +83,8 @@ class JobDB(BaseSQLDB):
         columns = _get_columns(Jobs.__table__, parameters)
         stmt = select(*columns)
 
-        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"]]
-            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)
+        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()
@@ -160,9 +126,7 @@ class JobDB(BaseSQLDB):
         await self.conn.execute(stmt)
 
     async def setJobAttributes(self, job_id, jobData):
-        """
-        TODO: add myDate and force parameters
-        """
+        """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)
@@ -178,9 +142,8 @@ class JobDB(BaseSQLDB):
         job_attrs,
         vo,
     ):
-        """
-        Check Consistency of Submitted JDL and set some defaults
-        Prepare subJDL with Job Requirements
+        """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
@@ -330,7 +293,7 @@ class JobDB(BaseSQLDB):
         }
 
     async def rescheduleJob(self, job_id) -> dict[str, Any]:
-        """Reschedule given job"""
+        """Reschedule given job."""
         from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd
         from DIRAC.Core.Utilities.ReturnValues import SErrorException
 
@@ -476,7 +439,7 @@ class JobDB(BaseSQLDB):
             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"""
+        """Store a command to be passed to the job together with the next heart beat."""
         try:
             stmt = insert(JobCommands).values(
                 JobID=job_id,
@@ -489,9 +452,7 @@ class JobDB(BaseSQLDB):
             raise JobNotFound(job_id) from e
 
     async def delete_jobs(self, job_ids: list[int]):
-        """
-        Delete jobs from the database
-        """
+        """Delete jobs from the database."""
         stmt = delete(JobJDLs).where(JobJDLs.JobID.in_(job_ids))
         await self.conn.execute(stmt)
 
@@ -499,7 +460,7 @@ class JobDB(BaseSQLDB):
         self, properties: dict[int, dict[str, Any]], update_timestamp: bool = False
     ) -> int:
         """Update the job parameters
-        All the jobs must update the same properties
+        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
@@ -507,11 +468,8 @@ class JobDB(BaseSQLDB):
         :return rowcount
 
         """
-
         # Check that all we always update the same set of properties
-        required_parameters_set = set(
-            [tuple(sorted(k.keys())) for k in properties.values()]
-        )
+        required_parameters_set = {tuple(sorted(k.keys())) for k in properties.values()}
 
         if len(required_parameters_set) != 1:
             raise NotImplementedError(
@@ -532,401 +490,3 @@ class JobDB(BaseSQLDB):
         rows = await self.conn.execute(stmt, update_parameters)
 
         return rows.rowcount
-
-
-MAGIC_EPOC_NUMBER = 1270000000
-
-
-class JobLoggingDB(BaseSQLDB):
-    """Frontend for the JobLoggingDB. Provides the ability to store changes with timestamps"""
-
-    metadata = JobLoggingDBBase.metadata
-
-    async def insert_record(
-        self,
-        job_id: int,
-        status: JobStatus,
-        minor_status: str,
-        application_status: str,
-        date: datetime,
-        source: str,
-    ):
-        """
-        Add a new entry to the JobLoggingDB table. One, two or all the three status
-        components (status, minorStatus, applicationStatus) can be specified.
-        Optionally the time stamp of the status can
-        be provided in a form of a string in a format '%Y-%m-%d %H:%M:%S' or
-        as datetime.datetime object. If the time stamp is not provided the current
-        UTC time is used.
-        """
-
-        # First, fetch the maximum SeqNum for the given job_id
-        seqnum_stmt = select(func.coalesce(func.max(LoggingInfo.SeqNum) + 1, 1)).where(
-            LoggingInfo.JobID == job_id
-        )
-        seqnum = await self.conn.scalar(seqnum_stmt)
-
-        epoc = (
-            time.mktime(date.timetuple())
-            + date.microsecond / 1000000.0
-            - MAGIC_EPOC_NUMBER
-        )
-
-        stmt = insert(LoggingInfo).values(
-            JobID=int(job_id),
-            SeqNum=seqnum,
-            Status=status,
-            MinorStatus=minor_status,
-            ApplicationStatus=application_status[:255],
-            StatusTime=date,
-            StatusTimeOrder=epoc,
-            Source=source[:32],
-        )
-        await self.conn.execute(stmt)
-
-    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
-        """
-
-        stmt = (
-            select(
-                LoggingInfo.Status,
-                LoggingInfo.MinorStatus,
-                LoggingInfo.ApplicationStatus,
-                LoggingInfo.StatusTime,
-                LoggingInfo.Source,
-            )
-            .where(LoggingInfo.JobID == int(job_id))
-            .order_by(LoggingInfo.StatusTimeOrder, LoggingInfo.StatusTime)
-        )
-        rows = await self.conn.execute(stmt)
-
-        values = []
-        for (
-            status,
-            minor_status,
-            application_status,
-            status_time,
-            status_source,
-        ) in rows:
-            values.append(
-                [
-                    status,
-                    minor_status,
-                    application_status,
-                    status_time.replace(tzinfo=timezone.utc),
-                    status_source,
-                ]
-            )
-
-        # If no value has been set for the application status in the first place,
-        # We put this status to unknown
-        res = []
-        if values:
-            if values[0][2] == "idem":
-                values[0][2] = "Unknown"
-
-            # We replace "idem" values by the value previously stated
-            for i in range(1, len(values)):
-                for j in range(3):
-                    if values[i][j] == "idem":
-                        values[i][j] = values[i - 1][j]
-
-            # And we replace arrays with tuples
-            for (
-                status,
-                minor_status,
-                application_status,
-                status_time,
-                status_source,
-            ) in values:
-                res.append(
-                    JobStatusReturn(
-                        Status=status,
-                        MinorStatus=minor_status,
-                        ApplicationStatus=application_status,
-                        StatusTime=status_time,
-                        Source=status_source,
-                    )
-                )
-
-        return res
-
-    async def delete_records(self, job_ids: list[int]):
-        """Delete logging records for given jobs"""
-        stmt = delete(LoggingInfo).where(LoggingInfo.JobID.in_(job_ids))
-        await self.conn.execute(stmt)
-
-    async def get_wms_time_stamps(self, job_id):
-        """Get TimeStamps for job MajorState transitions
-        return a {State:timestamp} dictionary
-        """
-
-        result = {}
-        stmt = select(
-            LoggingInfo.Status,
-            LoggingInfo.StatusTimeOrder,
-        ).where(LoggingInfo.JobID == job_id)
-        rows = await self.conn.execute(stmt)
-        if not rows.rowcount:
-            raise JobNotFound(job_id) from None
-
-        for event, etime in rows:
-            result[event] = str(etime + MAGIC_EPOC_NUMBER)
-
-        return result
-
-
-class TaskQueueDB(BaseSQLDB):
-    metadata = TaskQueueDBBase.metadata
-
-    async def get_tq_infos_for_jobs(
-        self, job_ids: list[int]
-    ) -> set[tuple[int, str, str, str]]:
-        """
-        Get the task queue info for given jobs
-        """
-        stmt = (
-            select(
-                TaskQueues.TQId, TaskQueues.Owner, TaskQueues.OwnerGroup, TaskQueues.VO
-            )
-            .join(JobsQueue, TaskQueues.TQId == JobsQueue.TQId)
-            .where(JobsQueue.JobId.in_(job_ids))
-        )
-        return set(
-            (int(row[0]), str(row[1]), str(row[2]), str(row[3]))
-            for row in (await self.conn.execute(stmt)).all()
-        )
-
-    async def get_owner_for_task_queue(self, tq_id: int) -> dict[str, str]:
-        """
-        Get the owner and owner group for a task queue
-        """
-        stmt = select(TaskQueues.Owner, TaskQueues.OwnerGroup, TaskQueues.VO).where(
-            TaskQueues.TQId == tq_id
-        )
-        return dict((await self.conn.execute(stmt)).one()._mapping)
-
-    async def remove_job(self, job_id: int):
-        """
-        Remove a job from the task queues
-        """
-        stmt = delete(JobsQueue).where(JobsQueue.JobId == job_id)
-        await self.conn.execute(stmt)
-
-    async def remove_jobs(self, job_ids: list[int]):
-        """
-        Remove jobs from the task queues
-        """
-        stmt = delete(JobsQueue).where(JobsQueue.JobId.in_(job_ids))
-        await self.conn.execute(stmt)
-
-    async def delete_task_queue_if_empty(
-        self,
-        tq_id: int,
-        tq_owner: str,
-        tq_group: str,
-        job_share: int,
-        group_properties: set[SecurityProperty],
-        enable_shares_correction: bool,
-        allow_background_tqs: bool,
-    ):
-        """
-        Try to delete a task queue if it's empty
-        """
-        # Check if the task queue is empty
-        stmt = (
-            select(TaskQueues.TQId)
-            .where(TaskQueues.Enabled >= 1)
-            .where(TaskQueues.TQId == tq_id)
-            .where(~TaskQueues.TQId.in_(select(JobsQueue.TQId)))
-        )
-        rows = await self.conn.execute(stmt)
-        if not rows.rowcount:
-            return
-
-        # Deleting the task queue (the other tables will be deleted in cascade)
-        stmt = delete(TaskQueues).where(TaskQueues.TQId == tq_id)
-        await self.conn.execute(stmt)
-
-        await self.recalculate_tq_shares_for_entity(
-            tq_owner,
-            tq_group,
-            job_share,
-            group_properties,
-            enable_shares_correction,
-            allow_background_tqs,
-        )
-
-    async def recalculate_tq_shares_for_entity(
-        self,
-        owner: str,
-        group: str,
-        job_share: int,
-        group_properties: set[SecurityProperty],
-        enable_shares_correction: bool,
-        allow_background_tqs: bool,
-    ):
-        """
-        Recalculate the shares for a user/userGroup combo
-        """
-        if JOB_SHARING in group_properties:
-            # If group has JobSharing just set prio for that entry, user is irrelevant
-            return await self.__set_priorities_for_entity(
-                owner, group, job_share, group_properties, allow_background_tqs
-            )
-
-        stmt = (
-            select(TaskQueues.Owner, func.count(TaskQueues.Owner))
-            .where(TaskQueues.OwnerGroup == group)
-            .group_by(TaskQueues.Owner)
-        )
-        rows = await self.conn.execute(stmt)
-        # make the rows a list of tuples
-        # Get owners in this group and the amount of times they appear
-        # TODO: I guess the rows are already a list of tupes
-        # maybe refactor
-        data = [(r[0], r[1]) for r in rows if r]
-        numOwners = len(data)
-        # If there are no owners do now
-        if numOwners == 0:
-            return
-        # Split the share amongst the number of owners
-        entities_shares = {row[0]: job_share / numOwners for row in data}
-
-        # TODO: implement the following
-        # If corrector is enabled let it work it's magic
-        # if enable_shares_correction:
-        #     entities_shares = await self.__shares_corrector.correct_shares(
-        #         entitiesShares, group=group
-        #     )
-
-        # Keep updating
-        owners = dict(data)
-        # IF the user is already known and has more than 1 tq, the rest of the users don't need to be modified
-        # (The number of owners didn't change)
-        if owner in owners and owners[owner] > 1:
-            await self.__set_priorities_for_entity(
-                owner,
-                group,
-                entities_shares[owner],
-                group_properties,
-                allow_background_tqs,
-            )
-            return
-        # Oops the number of owners may have changed so we recalculate the prio for all owners in the group
-        for owner in owners:
-            await self.__set_priorities_for_entity(
-                owner,
-                group,
-                entities_shares[owner],
-                group_properties,
-                allow_background_tqs,
-            )
-
-    async def __set_priorities_for_entity(
-        self,
-        owner: str,
-        group: str,
-        share,
-        properties: set[SecurityProperty],
-        allow_background_tqs: bool,
-    ):
-        """
-        Set the priority for a user/userGroup combo given a splitted share
-        """
-        from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import calculate_priority
-
-        stmt = (
-            select(
-                TaskQueues.TQId,
-                func.sum(JobsQueue.RealPriority) / func.count(JobsQueue.RealPriority),
-            )
-            .join(JobsQueue, TaskQueues.TQId == JobsQueue.TQId)
-            .where(TaskQueues.OwnerGroup == group)
-            .group_by(TaskQueues.TQId)
-        )
-        if JOB_SHARING not in properties:
-            stmt = stmt.where(TaskQueues.Owner == owner)
-        rows = await self.conn.execute(stmt)
-        tq_dict: dict[int, float] = {tq_id: priority for tq_id, priority in rows}
-
-        if not tq_dict:
-            return
-
-        rows = await self.retrieve_task_queues(list(tq_dict))
-
-        prio_dict = calculate_priority(tq_dict, rows, share, allow_background_tqs)
-
-        # Execute updates
-        for prio, tqs in prio_dict.items():
-            update_stmt = (
-                update(TaskQueues).where(TaskQueues.TQId.in_(tqs)).values(Priority=prio)
-            )
-            await self.conn.execute(update_stmt)
-
-    async def retrieve_task_queues(self, tq_id_list=None):
-        """
-        Get all the task queues
-        """
-        if tq_id_list is not None and not tq_id_list:
-            # Empty list => Fast-track no matches
-            return {}
-
-        stmt = (
-            select(
-                TaskQueues.TQId,
-                TaskQueues.Priority,
-                func.count(JobsQueue.TQId).label("Jobs"),
-                TaskQueues.Owner,
-                TaskQueues.OwnerGroup,
-                TaskQueues.VO,
-                TaskQueues.CPUTime,
-            )
-            .join(JobsQueue, TaskQueues.TQId == JobsQueue.TQId)
-            .join(SitesQueue, TaskQueues.TQId == SitesQueue.TQId)
-            .join(GridCEsQueue, TaskQueues.TQId == GridCEsQueue.TQId)
-            .group_by(
-                TaskQueues.TQId,
-                TaskQueues.Priority,
-                TaskQueues.Owner,
-                TaskQueues.OwnerGroup,
-                TaskQueues.VO,
-                TaskQueues.CPUTime,
-            )
-        )
-        if tq_id_list is not None:
-            stmt = stmt.where(TaskQueues.TQId.in_(tq_id_list))
-
-        tq_data: dict[int, dict[str, list[str]]] = dict(
-            dict(row._mapping) for row in await self.conn.execute(stmt)
-        )
-        # TODO: the line above should be equivalent to the following commented code, check this is the case
-        # for record in rows:
-        #     tqId = record[0]
-        #     tqData[tqId] = {
-        #         "Priority": record[1],
-        #         "Jobs": record[2],
-        #         "Owner": record[3],
-        #         "OwnerGroup": record[4],
-        #         "VO": record[5],
-        #         "CPUTime": record[6],
-        #     }
-
-        for tq_id in tq_data:
-            # TODO: maybe factorize this handy tuple list
-            for table, field in {
-                (SitesQueue, "Sites"),
-                (GridCEsQueue, "GridCEs"),
-                (BannedSitesQueue, "BannedSites"),
-                (PlatformsQueue, "Platforms"),
-                (JobTypesQueue, "JobTypes"),
-                (TagsQueue, "Tags"),
-            }:
-                stmt = select(table.Value).where(table.TQId == tq_id)
-                tq_data[tq_id][field] = list(
-                    row[0] for row in await self.conn.execute(stmt)
-                )
-
-        return tq_data

diracx/db/sql/{jobs → job}/schema.py

@@ -1,29 +1,22 @@
 import sqlalchemy.types as types
 from sqlalchemy import (
-    BigInteger,
-    Boolean,
     DateTime,
     Enum,
-    Float,
     ForeignKey,
     Index,
     Integer,
-    Numeric,
-    PrimaryKeyConstraint,
     String,
     Text,
 )
 from sqlalchemy.orm import declarative_base
 
-from ..utils import Column, DateNowColumn, NullColumn
+from ..utils import Column, NullColumn
 
 JobDBBase = declarative_base()
-JobLoggingDBBase = declarative_base()
-TaskQueueDBBase = declarative_base()
 
 
 class EnumBackedBool(types.TypeDecorator):
-    """Maps a ``EnumBackedBool()`` column to True/False in Python"""
+    """Maps a ``EnumBackedBool()`` column to True/False in Python."""
 
     impl = types.Enum
     cache_ok: bool = True
@@ -179,112 +172,3 @@ class JobCommands(JobDBBase):
     Status = Column(String(64), default="Received")
     ReceptionTime = Column(DateTime, primary_key=True)
     ExecutionTime = NullColumn(DateTime)
-
-
-class LoggingInfo(JobLoggingDBBase):
-    __tablename__ = "LoggingInfo"
-    JobID = Column(Integer)
-    SeqNum = Column(Integer)
-    Status = Column(String(32), default="")
-    MinorStatus = Column(String(128), default="")
-    ApplicationStatus = Column(String(255), default="")
-    StatusTime = DateNowColumn()
-    # TODO: Check that this corresponds to the DOUBLE(12,3) type in MySQL
-    StatusTimeOrder = Column(Numeric(precision=12, scale=3), default=0)
-    Source = Column(String(32), default="Unknown", name="StatusSource")
-    __table_args__ = (PrimaryKeyConstraint("JobID", "SeqNum"),)
-
-
-class TaskQueues(TaskQueueDBBase):
-    __tablename__ = "tq_TaskQueues"
-    TQId = Column(Integer, primary_key=True)
-    Owner = Column(String(255), nullable=False)
-    OwnerGroup = Column(String(32), nullable=False)
-    VO = Column(String(32), nullable=False)
-    CPUTime = Column(BigInteger, nullable=False)
-    Priority = Column(Float, nullable=False)
-    Enabled = Column(Boolean, nullable=False, default=0)
-    __table_args__ = (Index("TQOwner", "Owner", "OwnerGroup", "CPUTime"),)
-
-
-class JobsQueue(TaskQueueDBBase):
-    __tablename__ = "tq_Jobs"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    JobId = Column(Integer, primary_key=True)
-    Priority = Column(Integer, nullable=False)
-    RealPriority = Column(Float, nullable=False)
-    __table_args__ = (Index("TaskIndex", "TQId"),)
-
-
-class SitesQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToSites"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("SitesTaskIndex", "TQId"),
-        Index("SitesIndex", "Value"),
-    )
-
-
-class GridCEsQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToGridCEs"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("GridCEsTaskIndex", "TQId"),
-        Index("GridCEsValueIndex", "Value"),
-    )
-
-
-class BannedSitesQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToBannedSites"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("BannedSitesTaskIndex", "TQId"),
-        Index("BannedSitesValueIndex", "Value"),
-    )
-
-
-class PlatformsQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToPlatforms"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("PlatformsTaskIndex", "TQId"),
-        Index("PlatformsValueIndex", "Value"),
-    )
-
-
-class JobTypesQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToJobTypes"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("JobTypesTaskIndex", "TQId"),
-        Index("JobTypesValueIndex", "Value"),
-    )
-
-
-class TagsQueue(TaskQueueDBBase):
-    __tablename__ = "tq_TQToTags"
-    TQId = Column(
-        Integer, ForeignKey("tq_TaskQueues.TQId", ondelete="CASCADE"), primary_key=True
-    )
-    Value = Column(String(64), primary_key=True)
-    __table_args__ = (
-        Index("TagsTaskIndex", "TQId"),
-        Index("TagsValueIndex", "Value"),
-    )