diracx-db 0.0.1a17__py3-none-any.whl → 0.0.1a19__py3-none-any.whl

diracx/db/sql/job_logging/db.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import time
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, func, insert, select
+
+if TYPE_CHECKING:
+    pass
+
+from diracx.core.exceptions import JobNotFound
+from diracx.core.models import (
+    JobStatus,
+    JobStatusReturn,
+)
+
+from ..utils import BaseSQLDB
+from .schema import (
+    JobLoggingDBBase,
+    LoggingInfo,
+)
+
+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

diracx/db/sql/job_logging/schema.py

@@ -0,0 +1,25 @@
+from sqlalchemy import (
+    Integer,
+    Numeric,
+    PrimaryKeyConstraint,
+    String,
+)
+from sqlalchemy.orm import declarative_base
+
+from ..utils import Column, DateNowColumn
+
+JobLoggingDBBase = declarative_base()
+
+
+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"),)

diracx/db/sql/pilot_agents/db.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from sqlalchemy import insert
+
+from ..utils import BaseSQLDB
+from .schema import PilotAgents, PilotAgentsDBBase
+
+
+class PilotAgentsDB(BaseSQLDB):
+    """PilotAgentsDB class is a front-end to the PilotAgents Database."""
+
+    metadata = PilotAgentsDBBase.metadata
+
+    async def add_pilot_references(
+        self,
+        pilot_ref: list[str],
+        vo: str,
+        grid_type: str = "DIRAC",
+        pilot_stamps: dict | None = None,
+    ) -> None:
+
+        if pilot_stamps is None:
+            pilot_stamps = {}
+
+        now = datetime.now(tz=timezone.utc)
+
+        # Prepare the list of dictionaries for bulk insertion
+        values = [
+            {
+                "PilotJobReference": ref,
+                "VO": vo,
+                "GridType": grid_type,
+                "SubmissionTime": now,
+                "LastUpdateTime": now,
+                "Status": "Submitted",
+                "PilotStamp": pilot_stamps.get(ref, ""),
+            }
+            for ref in pilot_ref
+        ]
+
+        # Insert multiple rows in a single execute call
+        stmt = insert(PilotAgents).values(values)
+        await self.conn.execute(stmt)
+        return

diracx/db/sql/pilot_agents/schema.py

@@ -0,0 +1,58 @@
+from sqlalchemy import (
+    DateTime,
+    Double,
+    Index,
+    Integer,
+    String,
+    Text,
+)
+from sqlalchemy.orm import declarative_base
+
+from ..utils import Column, EnumBackedBool, NullColumn
+
+PilotAgentsDBBase = declarative_base()
+
+
+class PilotAgents(PilotAgentsDBBase):
+    __tablename__ = "PilotAgents"
+
+    PilotID = Column("PilotID", Integer, autoincrement=True, primary_key=True)
+    InitialJobID = Column("InitialJobID", Integer, default=0)
+    CurrentJobID = Column("CurrentJobID", Integer, default=0)
+    PilotJobReference = Column("PilotJobReference", String(255), default="Unknown")
+    PilotStamp = Column("PilotStamp", String(32), default="")
+    DestinationSite = Column("DestinationSite", String(128), default="NotAssigned")
+    Queue = Column("Queue", String(128), default="Unknown")
+    GridSite = Column("GridSite", String(128), default="Unknown")
+    VO = Column("VO", String(128))
+    GridType = Column("GridType", String(32), default="LCG")
+    BenchMark = Column("BenchMark", Double, default=0.0)
+    SubmissionTime = NullColumn("SubmissionTime", DateTime)
+    LastUpdateTime = NullColumn("LastUpdateTime", DateTime)
+    Status = Column("Status", String(32), default="Unknown")
+    StatusReason = Column("StatusReason", String(255), default="Unknown")
+    AccountingSent = Column("AccountingSent", EnumBackedBool(), default=False)
+
+    __table_args__ = (
+        Index("PilotJobReference", "PilotJobReference"),
+        Index("Status", "Status"),
+        Index("Statuskey", "GridSite", "DestinationSite", "Status"),
+    )
+
+
+class JobToPilotMapping(PilotAgentsDBBase):
+    __tablename__ = "JobToPilotMapping"
+
+    PilotID = Column("PilotID", Integer, primary_key=True)
+    JobID = Column("JobID", Integer, primary_key=True)
+    StartTime = Column("StartTime", DateTime)
+
+    __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
+
+
+class PilotOutput(PilotAgentsDBBase):
+    __tablename__ = "PilotOutput"
+
+    PilotID = Column("PilotID", Integer, primary_key=True)
+    StdOutput = Column("StdOutput", Text)
+    StdError = Column("StdError", Text)

diracx/db/sql/sandbox_metadata/db.py

@@ -15,7 +15,7 @@ class SandboxMetadataDB(BaseSQLDB):
     metadata = SandboxMetadataDBBase.metadata
 
     async def upsert_owner(self, user: UserInfo) -> int:
-        """Get the id of the owner from the database"""
+        """Get the id of the owner from the database."""
         # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
         stmt = sqlalchemy.select(sb_Owners.OwnerID).where(
             sb_Owners.Owner == user.preferred_username,
@@ -36,7 +36,7 @@ class SandboxMetadataDB(BaseSQLDB):
 
     @staticmethod
     def get_pfn(bucket_name: str, user: UserInfo, sandbox_info: SandboxInfo) -> str:
-        """Get the sandbox's user namespaced and content addressed PFN"""
+        """Get the sandbox's user namespaced and content addressed PFN."""
         parts = [
             "S3",
             bucket_name,
@@ -50,7 +50,7 @@ class SandboxMetadataDB(BaseSQLDB):
     async def insert_sandbox(
         self, se_name: str, user: UserInfo, pfn: str, size: int
     ) -> None:
-        """Add a new sandbox in SandboxMetadataDB"""
+        """Add a new sandbox in SandboxMetadataDB."""
         # TODO: Follow https://github.com/DIRACGrid/diracx/issues/49
         owner_id = await self.upsert_owner(user)
         stmt = sqlalchemy.insert(sb_SandBoxes).values(
@@ -88,13 +88,13 @@ class SandboxMetadataDB(BaseSQLDB):
 
     @staticmethod
     def jobid_to_entity_id(job_id: int) -> str:
-        """Define the entity id as 'Entity:entity_id' due to the DB definition"""
+        """Define the entity id as 'Entity:entity_id' due to the DB definition."""
         return f"Job:{job_id}"
 
     async def get_sandbox_assigned_to_job(
         self, job_id: int, sb_type: SandboxType
     ) -> list[Any]:
-        """Get the sandbox assign to job"""
+        """Get the sandbox assign to job."""
         entity_id = self.jobid_to_entity_id(job_id)
         stmt = (
             sqlalchemy.select(sb_SandBoxes.SEPFN)
@@ -114,7 +114,7 @@ class SandboxMetadataDB(BaseSQLDB):
         sb_type: SandboxType,
         se_name: str,
     ) -> None:
-        """Mapp sandbox and jobs"""
+        """Mapp sandbox and jobs."""
         for job_id in jobs_ids:
             # Define the entity id as 'Entity:entity_id' due to the DB definition:
             entity_id = self.jobid_to_entity_id(job_id)
@@ -140,12 +140,14 @@ class SandboxMetadataDB(BaseSQLDB):
             assert result.rowcount == 1
 
     async def unassign_sandboxes_to_jobs(self, jobs_ids: list[int]) -> None:
-        """Delete mapping between jobs and sandboxes"""
+        """Delete mapping between jobs and sandboxes."""
         for job_id in jobs_ids:
             entity_id = self.jobid_to_entity_id(job_id)
-            sb_sel_stmt = sqlalchemy.select(
-                sb_SandBoxes.SBId,
-            ).where(sb_EntityMapping.EntityId == entity_id)
+            sb_sel_stmt = sqlalchemy.select(sb_SandBoxes.SBId)
+            sb_sel_stmt = sb_sel_stmt.join(
+                sb_EntityMapping, sb_EntityMapping.SBId == sb_SandBoxes.SBId
+            )
+            sb_sel_stmt = sb_sel_stmt.where(sb_EntityMapping.EntityId == entity_id)
 
             result = await self.conn.execute(sb_sel_stmt)
             sb_ids = [row.SBId for row in result]

diracx/db/sql/task_queue/db.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, func, select, update
+
+if TYPE_CHECKING:
+    pass
+
+from diracx.core.properties import JOB_SHARING, SecurityProperty
+
+from ..utils import BaseSQLDB
+from .schema import (
+    BannedSitesQueue,
+    GridCEsQueue,
+    JobsQueue,
+    JobTypesQueue,
+    PlatformsQueue,
+    SitesQueue,
+    TagsQueue,
+    TaskQueueDBBase,
+    TaskQueues,
+)
+
+
+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