diracx-db 0.0.1a22__py3-none-any.whl → 0.0.1a24__py3-none-any.whl

diracx/db/sql/job_logging/db.py

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
 
 from collections import defaultdict
 
-from diracx.core.exceptions import JobNotFound
+from diracx.core.exceptions import JobNotFoundError
 from diracx.core.models import (
     JobStatus,
     JobStatusReturn,
@@ -57,9 +57,9 @@ class JobLoggingDB(BaseSQLDB):
         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
+        # First, fetch the maximum seq_num for the given job_id
+        seqnum_stmt = select(func.coalesce(func.max(LoggingInfo.seq_num) + 1, 1)).where(
+            LoggingInfo.job_id == job_id
         )
         seqnum = await self.conn.scalar(seqnum_stmt)
 
@@ -70,14 +70,14 @@ class JobLoggingDB(BaseSQLDB):
         )
 
         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],
+            job_id=int(job_id),
+            seq_num=seqnum,
+            status=status,
+            minor_status=minor_status,
+            application_status=application_status[:255],
+            status_time=date,
+            status_time_order=epoc,
+            source=source[:32],
         )
         await self.conn.execute(stmt)
 
@@ -97,18 +97,17 @@ class JobLoggingDB(BaseSQLDB):
         # First, fetch the maximum SeqNums for the given job_ids
         seqnum_stmt = (
             select(
-                LoggingInfo.JobID, func.coalesce(func.max(LoggingInfo.SeqNum) + 1, 1)
+                LoggingInfo.job_id, func.coalesce(func.max(LoggingInfo.seq_num) + 1, 1)
             )
-            .where(LoggingInfo.JobID.in_([record.job_id for record in records]))
-            .group_by(LoggingInfo.JobID)
+            .where(LoggingInfo.job_id.in_([record.job_id for record in records]))
+            .group_by(LoggingInfo.job_id)
         )
 
         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),
+            LoggingInfo.__table__.insert(),
             [
                 {
                     "JobID": record.job_id,
@@ -118,38 +117,43 @@ class JobLoggingDB(BaseSQLDB):
                     "ApplicationStatus": record.application_status[:255],
                     "StatusTime": record.date,
                     "StatusTimeOrder": get_epoc(record.date),
-                    "Source": record.source[:32],
+                    "StatusSource": record.source[:32],
                 }
                 for record in records
             ],
         )
 
-    async def get_records(self, job_id: int) -> list[JobStatusReturn]:
+    async def get_records(self, job_ids: list[int]) -> dict[int, JobStatusReturn]:
         """Returns a Status,MinorStatus,ApplicationStatus,StatusTime,Source tuple
         for each record found for job specified by its jobID in historical order.
         """
+        # We could potentially use a group_by here, but we need to post-process the
+        # results later.
         stmt = (
             select(
-                LoggingInfo.Status,
-                LoggingInfo.MinorStatus,
-                LoggingInfo.ApplicationStatus,
-                LoggingInfo.StatusTime,
-                LoggingInfo.Source,
+                LoggingInfo.job_id,
+                LoggingInfo.status,
+                LoggingInfo.minor_status,
+                LoggingInfo.application_status,
+                LoggingInfo.status_time,
+                LoggingInfo.source,
             )
-            .where(LoggingInfo.JobID == int(job_id))
-            .order_by(LoggingInfo.StatusTimeOrder, LoggingInfo.StatusTime)
+            .where(LoggingInfo.job_id.in_(job_ids))
+            .order_by(LoggingInfo.status_time_order, LoggingInfo.status_time)
         )
         rows = await self.conn.execute(stmt)
 
-        values = []
+        values = defaultdict(list)
         for (
+            job_id,
             status,
             minor_status,
             application_status,
             status_time,
             status_source,
         ) in rows:
-            values.append(
+
+            values[job_id].append(
                 [
                     status,
                     minor_status,
@@ -161,16 +165,16 @@ class JobLoggingDB(BaseSQLDB):
 
         # 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"
+        res: dict = defaultdict(list)
+        for job_id, history in values.items():
+            if history[0][2] == "idem":
+                history[0][2] = "Unknown"
 
             # We replace "idem" values by the value previously stated
-            for i in range(1, len(values)):
+            for i in range(1, len(history)):
                 for j in range(3):
-                    if values[i][j] == "idem":
-                        values[i][j] = values[i - 1][j]
+                    if history[i][j] == "idem":
+                        history[i][j] = history[i - 1][j]
 
             # And we replace arrays with tuples
             for (
@@ -179,8 +183,8 @@ class JobLoggingDB(BaseSQLDB):
                 application_status,
                 status_time,
                 status_source,
-            ) in values:
-                res.append(
+            ) in history:
+                res[job_id].append(
                     JobStatusReturn(
                         Status=status,
                         MinorStatus=minor_status,
@@ -194,7 +198,7 @@ class JobLoggingDB(BaseSQLDB):
 
     async def delete_records(self, job_ids: list[int]):
         """Delete logging records for given jobs."""
-        stmt = delete(LoggingInfo).where(LoggingInfo.JobID.in_(job_ids))
+        stmt = delete(LoggingInfo).where(LoggingInfo.job_id.in_(job_ids))
         await self.conn.execute(stmt)
 
     async def get_wms_time_stamps(self, job_id):
@@ -203,12 +207,12 @@ class JobLoggingDB(BaseSQLDB):
         """
         result = {}
         stmt = select(
-            LoggingInfo.Status,
-            LoggingInfo.StatusTimeOrder,
-        ).where(LoggingInfo.JobID == job_id)
+            LoggingInfo.status,
+            LoggingInfo.status_time_order,
+        ).where(LoggingInfo.job_id == job_id)
         rows = await self.conn.execute(stmt)
         if not rows.rowcount:
-            raise JobNotFound(job_id) from None
+            raise JobNotFoundError(job_id) from None
 
         for event, etime in rows:
             result[event] = str(etime + MAGIC_EPOC_NUMBER)
@@ -221,10 +225,10 @@ class JobLoggingDB(BaseSQLDB):
         """
         result = defaultdict(dict)
         stmt = select(
-            LoggingInfo.JobID,
-            LoggingInfo.Status,
-            LoggingInfo.StatusTimeOrder,
-        ).where(LoggingInfo.JobID.in_(job_ids))
+            LoggingInfo.job_id,
+            LoggingInfo.status,
+            LoggingInfo.status_time_order,
+        ).where(LoggingInfo.job_id.in_(job_ids))
         rows = await self.conn.execute(stmt)
         if not rows.rowcount:
             return {}

diracx/db/sql/job_logging/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     Integer,
     Numeric,
@@ -13,13 +15,15 @@ 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()
+    job_id = Column("JobID", Integer)
+    seq_num = Column("SeqNum", Integer)
+    status = Column("Status", String(32), default="")
+    minor_status = Column("MinorStatus", String(128), default="")
+    application_status = Column("ApplicationStatus", String(255), default="")
+    status_time = DateNowColumn("StatusTime")
     # 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")
+    status_time_order = Column(
+        "StatusTimeOrder", Numeric(precision=12, scale=3), default=0
+    )
+    source = Column("StatusSource", String(32), default="Unknown")
     __table_args__ = (PrimaryKeyConstraint("JobID", "SeqNum"),)

diracx/db/sql/pilot_agents/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     DateTime,
     Double,
@@ -16,22 +18,22 @@ 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)
+    pilot_id = Column("PilotID", Integer, autoincrement=True, primary_key=True)
+    initial_job_id = Column("InitialJobID", Integer, default=0)
+    current_job_id = Column("CurrentJobID", Integer, default=0)
+    pilot_job_reference = Column("PilotJobReference", String(255), default="Unknown")
+    pilot_stamp = Column("PilotStamp", String(32), default="")
+    destination_site = Column("DestinationSite", String(128), default="NotAssigned")
+    queue = Column("Queue", String(128), default="Unknown")
+    grid_site = Column("GridSite", String(128), default="Unknown")
+    vo = Column("VO", String(128))
+    grid_type = Column("GridType", String(32), default="LCG")
+    benchmark = Column("BenchMark", Double, default=0.0)
+    submission_time = NullColumn("SubmissionTime", DateTime)
+    last_update_time = NullColumn("LastUpdateTime", DateTime)
+    status = Column("Status", String(32), default="Unknown")
+    status_reason = Column("StatusReason", String(255), default="Unknown")
+    accounting_sent = Column("AccountingSent", EnumBackedBool(), default=False)
 
     __table_args__ = (
         Index("PilotJobReference", "PilotJobReference"),
@@ -43,9 +45,9 @@ class PilotAgents(PilotAgentsDBBase):
 class JobToPilotMapping(PilotAgentsDBBase):
     __tablename__ = "JobToPilotMapping"
 
-    PilotID = Column("PilotID", Integer, primary_key=True)
-    JobID = Column("JobID", Integer, primary_key=True)
-    StartTime = Column("StartTime", DateTime)
+    pilot_id = Column("PilotID", Integer, primary_key=True)
+    job_id = Column("JobID", Integer, primary_key=True)
+    start_time = Column("StartTime", DateTime)
 
     __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
 
@@ -53,6 +55,6 @@ class JobToPilotMapping(PilotAgentsDBBase):
 class PilotOutput(PilotAgentsDBBase):
     __tablename__ = "PilotOutput"
 
-    PilotID = Column("PilotID", Integer, primary_key=True)
-    StdOutput = Column("StdOutput", Text)
-    StdError = Column("StdError", Text)
+    pilot_id = Column("PilotID", Integer, primary_key=True)
+    std_output = Column("StdOutput", Text)
+    std_error = Column("StdError", Text)

diracx/db/sql/sandbox_metadata/db.py

@@ -2,13 +2,15 @@ from __future__ import annotations
 
 from typing import Any
 
-import sqlalchemy
+from sqlalchemy import Executable, delete, insert, literal, select, update
+from sqlalchemy.exc import IntegrityError, NoResultFound
 
+from diracx.core.exceptions import SandboxNotFoundError
 from diracx.core.models import SandboxInfo, SandboxType, UserInfo
 from diracx.db.sql.utils import BaseSQLDB, utcnow
 
 from .schema import Base as SandboxMetadataDBBase
-from .schema import sb_EntityMapping, sb_Owners, sb_SandBoxes
+from .schema import SandBoxes, SBEntityMapping, SBOwners
 
 
 class SandboxMetadataDB(BaseSQLDB):
@@ -17,16 +19,16 @@ class SandboxMetadataDB(BaseSQLDB):
     async def upsert_owner(self, user: UserInfo) -> int:
         """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,
-            sb_Owners.OwnerGroup == user.dirac_group,
-            sb_Owners.VO == user.vo,
+        stmt = select(SBOwners.OwnerID).where(
+            SBOwners.Owner == user.preferred_username,
+            SBOwners.OwnerGroup == user.dirac_group,
+            SBOwners.VO == user.vo,
         )
         result = await self.conn.execute(stmt)
         if owner_id := result.scalar_one_or_none():
             return owner_id
 
-        stmt = sqlalchemy.insert(sb_Owners).values(
+        stmt = insert(SBOwners).values(
             Owner=user.preferred_username,
             OwnerGroup=user.dirac_group,
             VO=user.vo,
@@ -53,7 +55,7 @@ class SandboxMetadataDB(BaseSQLDB):
         """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(
+        stmt = insert(SandBoxes).values(
             OwnerId=owner_id,
             SEName=se_name,
             SEPFN=pfn,
@@ -63,27 +65,31 @@ class SandboxMetadataDB(BaseSQLDB):
         )
         try:
             result = await self.conn.execute(stmt)
-        except sqlalchemy.exc.IntegrityError:
+        except IntegrityError:
             await self.update_sandbox_last_access_time(se_name, pfn)
         else:
             assert result.rowcount == 1
 
     async def update_sandbox_last_access_time(self, se_name: str, pfn: str) -> None:
         stmt = (
-            sqlalchemy.update(sb_SandBoxes)
-            .where(sb_SandBoxes.SEName == se_name, sb_SandBoxes.SEPFN == pfn)
+            update(SandBoxes)
+            .where(SandBoxes.SEName == se_name, SandBoxes.SEPFN == pfn)
             .values(LastAccessTime=utcnow())
         )
         result = await self.conn.execute(stmt)
         assert result.rowcount == 1
 
-    async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool:
+    async def sandbox_is_assigned(self, pfn: str, se_name: str) -> bool | None:
         """Checks if a sandbox exists and has been assigned."""
-        stmt: sqlalchemy.Executable = sqlalchemy.select(sb_SandBoxes.Assigned).where(
-            sb_SandBoxes.SEName == se_name, sb_SandBoxes.SEPFN == pfn
+        stmt: Executable = select(SandBoxes.Assigned).where(
+            SandBoxes.SEName == se_name, SandBoxes.SEPFN == pfn
         )
         result = await self.conn.execute(stmt)
-        is_assigned = result.scalar_one()
+        try:
+            is_assigned = result.scalar_one()
+        except NoResultFound as e:
+            raise SandboxNotFoundError(pfn, se_name) from e
+
         return is_assigned
 
     @staticmethod
@@ -97,11 +103,11 @@ class SandboxMetadataDB(BaseSQLDB):
         """Get the sandbox assign to job."""
         entity_id = self.jobid_to_entity_id(job_id)
         stmt = (
-            sqlalchemy.select(sb_SandBoxes.SEPFN)
-            .where(sb_SandBoxes.SBId == sb_EntityMapping.SBId)
+            select(SandBoxes.SEPFN)
+            .where(SandBoxes.SBId == SBEntityMapping.SBId)
             .where(
-                sb_EntityMapping.EntityId == entity_id,
-                sb_EntityMapping.Type == sb_type,
+                SBEntityMapping.EntityId == entity_id,
+                SBEntityMapping.Type == sb_type,
             )
         )
         result = await self.conn.execute(stmt)
@@ -118,24 +124,20 @@ class SandboxMetadataDB(BaseSQLDB):
         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)
-            select_sb_id = sqlalchemy.select(
-                sb_SandBoxes.SBId,
-                sqlalchemy.literal(entity_id).label("EntityId"),
-                sqlalchemy.literal(sb_type).label("Type"),
+            select_sb_id = select(
+                SandBoxes.SBId,
+                literal(entity_id).label("EntityId"),
+                literal(sb_type).label("Type"),
             ).where(
-                sb_SandBoxes.SEName == se_name,
-                sb_SandBoxes.SEPFN == pfn,
+                SandBoxes.SEName == se_name,
+                SandBoxes.SEPFN == pfn,
             )
-            stmt = sqlalchemy.insert(sb_EntityMapping).from_select(
+            stmt = insert(SBEntityMapping).from_select(
                 ["SBId", "EntityId", "Type"], select_sb_id
             )
             await self.conn.execute(stmt)
 
-            stmt = (
-                sqlalchemy.update(sb_SandBoxes)
-                .where(sb_SandBoxes.SEPFN == pfn)
-                .values(Assigned=True)
-            )
+            stmt = update(SandBoxes).where(SandBoxes.SEPFN == pfn).values(Assigned=True)
             result = await self.conn.execute(stmt)
             assert result.rowcount == 1
 
@@ -143,29 +145,29 @@ class SandboxMetadataDB(BaseSQLDB):
         """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)
+            sb_sel_stmt = select(SandBoxes.SBId)
             sb_sel_stmt = sb_sel_stmt.join(
-                sb_EntityMapping, sb_EntityMapping.SBId == sb_SandBoxes.SBId
+                SBEntityMapping, SBEntityMapping.SBId == SandBoxes.SBId
             )
-            sb_sel_stmt = sb_sel_stmt.where(sb_EntityMapping.EntityId == entity_id)
+            sb_sel_stmt = sb_sel_stmt.where(SBEntityMapping.EntityId == entity_id)
 
             result = await self.conn.execute(sb_sel_stmt)
             sb_ids = [row.SBId for row in result]
 
-            del_stmt = sqlalchemy.delete(sb_EntityMapping).where(
-                sb_EntityMapping.EntityId == entity_id
+            del_stmt = delete(SBEntityMapping).where(
+                SBEntityMapping.EntityId == entity_id
             )
             await self.conn.execute(del_stmt)
 
-            sb_entity_sel_stmt = sqlalchemy.select(sb_EntityMapping.SBId).where(
-                sb_EntityMapping.SBId.in_(sb_ids)
+            sb_entity_sel_stmt = select(SBEntityMapping.SBId).where(
+                SBEntityMapping.SBId.in_(sb_ids)
             )
             result = await self.conn.execute(sb_entity_sel_stmt)
             remaining_sb_ids = [row.SBId for row in result]
             if not remaining_sb_ids:
                 unassign_stmt = (
-                    sqlalchemy.update(sb_SandBoxes)
-                    .where(sb_SandBoxes.SBId.in_(sb_ids))
+                    update(SandBoxes)
+                    .where(SandBoxes.SBId.in_(sb_ids))
                     .values(Assigned=False)
                 )
                 await self.conn.execute(unassign_stmt)

diracx/db/sql/sandbox_metadata/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     BigInteger,
     Boolean,
@@ -14,7 +16,7 @@ from diracx.db.sql.utils import Column, DateNowColumn
 Base = declarative_base()
 
 
-class sb_Owners(Base):
+class SBOwners(Base):
     __tablename__ = "sb_Owners"
     OwnerID = Column(Integer, autoincrement=True)
     Owner = Column(String(32))
@@ -23,7 +25,7 @@ class sb_Owners(Base):
     __table_args__ = (PrimaryKeyConstraint("OwnerID"),)
 
 
-class sb_SandBoxes(Base):
+class SandBoxes(Base):
     __tablename__ = "sb_SandBoxes"
     SBId = Column(Integer, autoincrement=True)
     OwnerId = Column(Integer)
@@ -40,7 +42,7 @@ class sb_SandBoxes(Base):
     )
 
 
-class sb_EntityMapping(Base):
+class SBEntityMapping(Base):
     __tablename__ = "sb_EntityMapping"
     SBId = Column(Integer)
     EntityId = Column(String(128))

diracx/db/sql/task_queue/db.py

@@ -121,12 +121,12 @@ class TaskQueueDB(BaseSQLDB):
         # 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)
+        num_owners = len(data)
         # If there are no owners do now
-        if numOwners == 0:
+        if num_owners == 0:
             return
         # Split the share amongst the number of owners
-        entities_shares = {row[0]: job_share / numOwners for row in data}
+        entities_shares = {row[0]: job_share / num_owners for row in data}
 
         # TODO: implement the following
         # If corrector is enabled let it work it's magic

diracx/db/sql/task_queue/schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from sqlalchemy import (
     BigInteger,
     Boolean,