diracx-db 0.0.1a21__py3-none-any.whl → 0.0.6__py3-none-any.whl

diracx/db/sql/job_logging/db.py

@@ -1,39 +1,15 @@
 from __future__ import annotations
 
-import time
+from collections import defaultdict
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel
-from sqlalchemy import delete, func, insert, select
+from typing import Iterable
 
-if TYPE_CHECKING:
-    pass
+from sqlalchemy import delete, func, select
 
-from collections import defaultdict
-
-from diracx.core.exceptions import JobNotFound
-from diracx.core.models import (
-    JobStatus,
-    JobStatusReturn,
-)
+from diracx.core.models import JobLoggingRecord, JobStatusReturn
 
 from ..utils import BaseSQLDB
-from .schema import (
-    JobLoggingDBBase,
-    LoggingInfo,
-)
-
-MAGIC_EPOC_NUMBER = 1270000000
-
-
-class JobLoggingRecord(BaseModel):
-    job_id: int
-    status: JobStatus
-    minor_status: str
-    application_status: str
-    date: datetime
-    source: str
+from .schema import JobLoggingDBBase, LoggingInfo
 
 
 class JobLoggingDB(BaseSQLDB):
@@ -41,115 +17,79 @@ class JobLoggingDB(BaseSQLDB):
 
     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 bulk_insert_record(
+    async def insert_records(
         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)
+                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))}
+        seqnums = {
+            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),
-            [
+        values = []
+        for record in records:
+            if record.job_id not in seqnums:
+                seqnums[record.job_id] = 1
+
+            values.append(
                 {
                     "JobID": record.job_id,
-                    "SeqNum": seqnum.get(record.job_id, 1),
+                    "SeqNum": seqnums[record.job_id],
                     "Status": record.status,
                     "MinorStatus": record.minor_status,
                     "ApplicationStatus": record.application_status[:255],
                     "StatusTime": record.date,
-                    "StatusTimeOrder": get_epoc(record.date),
-                    "Source": record.source[:32],
+                    "StatusTimeOrder": record.date,
+                    "StatusSource": record.source[:32],
                 }
-                for record in records
-            ],
+            )
+            seqnums[record.job_id] = seqnums[record.job_id] + 1
+
+        await self.conn.execute(
+            LoggingInfo.__table__.insert(),
+            values,
         )
 
-    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 +101,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 +119,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,42 +134,19 @@ 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):
-        """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
-
-    async def get_wms_time_stamps_bulk(self, job_ids):
+    async def get_wms_time_stamps(
+        self, job_ids: Iterable[int]
+    ) -> dict[int, dict[str, datetime]]:
         """Get TimeStamps for job MajorState transitions for multiple jobs at once
         return a {JobID: {State:timestamp}} dictionary.
         """
-        result = defaultdict(dict)
+        result: defaultdict[int, dict[str, datetime]] = 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
+            LoggingInfo.job_id, LoggingInfo.status, LoggingInfo.status_time_order
+        ).where(LoggingInfo.job_id.in_(job_ids))
+        for job_id, event, etime in await self.conn.execute(stmt):
+            result[job_id][event] = etime
+        return dict(result)

diracx/db/sql/job_logging/schema.py

@@ -1,9 +1,8 @@
-from sqlalchemy import (
-    Integer,
-    Numeric,
-    PrimaryKeyConstraint,
-    String,
-)
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from sqlalchemy import Integer, Numeric, PrimaryKeyConstraint, String, TypeDecorator
 from sqlalchemy.orm import declarative_base
 
 from ..utils import Column, DateNowColumn
@@ -11,15 +10,55 @@ from ..utils import Column, DateNowColumn
 JobLoggingDBBase = declarative_base()
 
 
+class MagicEpochDateTime(TypeDecorator):
+    """A SQLAlchemy type that stores a datetime as a numeric value representing the
+    seconds elapsed since MAGIC_EPOC_NUMBER. The underlying column is defined as
+    Numeric(12,3) which provides a fixed-precision representation.
+    """
+
+    impl = Numeric(12, 3)
+    cache_ok = True
+
+    MAGIC_EPOC_NUMBER = 1270000000
+
+    def process_bind_param(self, value, dialect):
+        """Convert a Python datetime to a numeric value: (timestamp - MAGIC_EPOC_NUMBER).
+        The result is rounded to three decimal places.
+        """
+        if value is None:
+            return None
+        if isinstance(value, datetime):
+            # Convert datetime to seconds since the Unix epoch, subtract our magic epoch,
+            # and round to three decimal places.
+            epoch_seconds = (
+                value.replace(tzinfo=UTC).timestamp() - self.MAGIC_EPOC_NUMBER
+            )
+            return round(epoch_seconds, 3)
+        raise ValueError(
+            "Expected a datetime object for MagicEpochDateTime bind parameter."
+        )
+
+    def process_result_value(self, value, dialect):
+        """Convert the numeric database value back into a Python datetime by reversing the
+        stored difference (adding MAGIC_EPOC_NUMBER).
+        """
+        if value is None:
+            return None
+        # Carefully convert from Decimal to datetime to avoid losing precision
+        value += self.MAGIC_EPOC_NUMBER
+        value_int = int(value)
+        result = datetime.fromtimestamp(value_int, tz=UTC)
+        return result.replace(microsecond=int((value - value_int) * 1_000_000))
+
+
 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")
+    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")
+    status_time_order = Column("StatusTimeOrder", MagicEpochDateTime, default=0)
+    source = Column("StatusSource", String(32), default="Unknown")
     __table_args__ = (PrimaryKeyConstraint("JobID", "SeqNum"),)

diracx/db/sql/pilot_agents/db.py

@@ -20,7 +20,6 @@ class PilotAgentsDB(BaseSQLDB):
         grid_type: str = "DIRAC",
         pilot_stamps: dict | None = None,
     ) -> None:
-
         if pilot_stamps is None:
             pilot_stamps = {}
 

diracx/db/sql/pilot_agents/schema.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 from sqlalchemy import (
-    DateTime,
     Double,
     Index,
     Integer,
@@ -8,6 +9,8 @@ from sqlalchemy import (
 )
 from sqlalchemy.orm import declarative_base
 
+from diracx.db.sql.utils.types import SmarterDateTime
+
 from ..utils import Column, EnumBackedBool, NullColumn
 
 PilotAgentsDBBase = declarative_base()
@@ -16,22 +19,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", SmarterDateTime)
+    last_update_time = NullColumn("LastUpdateTime", SmarterDateTime)
+    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 +46,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", SmarterDateTime)
 
     __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
 
@@ -53,6 +56,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)