diracx-db 0.0.1a7__tar.gz → 0.0.1a8__tar.gz

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a7
+Version: 0.0.1a8
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/src/diracx/db/sql/jobs/db.py

@@ -3,11 +3,14 @@ from __future__ import annotations
 import logging
 import time
 from datetime import datetime, timezone
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from sqlalchemy import delete, func, insert, select, update
+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 (
     JobMinorStatus,
@@ -474,6 +477,44 @@ class JobDB(BaseSQLDB):
         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 = 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
+
 
 MAGIC_EPOC_NUMBER = 1270000000
 
@@ -521,12 +562,12 @@ class JobLoggingDB(BaseSQLDB):
             ApplicationStatus=application_status[:255],
             StatusTime=date,
             StatusTimeOrder=epoc,
-            StatusSource=source[:32],
+            Source=source[:32],
         )
         await self.conn.execute(stmt)
 
     async def get_records(self, job_id: int) -> list[JobStatusReturn]:
-        """Returns a Status,MinorStatus,ApplicationStatus,StatusTime,StatusSource tuple
+        """Returns a Status,MinorStatus,ApplicationStatus,StatusTime,Source tuple
         for each record found for job specified by its jobID in historical order
         """
 
@@ -536,7 +577,7 @@ class JobLoggingDB(BaseSQLDB):
                 LoggingInfo.MinorStatus,
                 LoggingInfo.ApplicationStatus,
                 LoggingInfo.StatusTime,
-                LoggingInfo.StatusSource,
+                LoggingInfo.Source,
             )
             .where(LoggingInfo.JobID == int(job_id))
             .order_by(LoggingInfo.StatusTimeOrder, LoggingInfo.StatusTime)
@@ -588,7 +629,7 @@ class JobLoggingDB(BaseSQLDB):
                         MinorStatus=minor_status,
                         ApplicationStatus=application_status,
                         StatusTime=status_time,
-                        StatusSource=status_source,
+                        Source=status_source,
                     )
                 )
 

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/src/diracx/db/sql/jobs/schema.py

@@ -204,7 +204,7 @@ class LoggingInfo(JobLoggingDBBase):
     StatusTime = DateNowColumn()
     # TODO: Check that this corresponds to the DOUBLE(12,3) type in MySQL
     StatusTimeOrder = Column(Numeric(precision=12, scale=3), default=0)
-    StatusSource = Column(String(32), default="Unknown")
+    Source = Column(String(32), default="Unknown", name="StatusSource")
     __table_args__ = (PrimaryKeyConstraint("JobID", "SeqNum"),)
 
 

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/src/diracx/db/sql/jobs/status_utility.py

@@ -41,7 +41,7 @@ async def set_job_status(
     # transform JobStateUpdate objects into dicts
     statusDict = {}
     for key, value in status.items():
-        statusDict[key] = value.dict(by_alias=True)
+        statusDict[key] = {k: v for k, v in value.dict().items() if v is not None}
 
     res = await job_db.search(
         parameters=["Status", "StartExecTime", "EndExecTime"],
@@ -110,7 +110,7 @@ async def set_job_status(
         #         return result
 
     for updTime in updateTimes:
-        if statusDict[updTime]["StatusSource"].startswith("Job"):
+        if statusDict[updTime]["Source"].startswith("Job"):
             job_data["HeartBeatTime"] = updTime
 
     if not startTime and newStartTime:
@@ -122,22 +122,16 @@ async def set_job_status(
     if job_data:
         await job_db.setJobAttributes(job_id, job_data)
 
-    # Update the JobLoggingDB records
-    # TODO: Because I really didn't liked the fact that the input field is called "Source"
-    # and the output field is called "StatusSource"
-    # I changed the name of the input field to "StatusSource"
-    # Meaning this change must be added to the transformation layer for DIRAC.
-
     for updTime in updateTimes:
         sDict = statusDict[updTime]
-        if not sDict["Status"]:
+        if not sDict.get("Status"):
             sDict["Status"] = "idem"
-        if not sDict["MinorStatus"]:
+        if not sDict.get("MinorStatus"):
             sDict["MinorStatus"] = "idem"
-        if not sDict["ApplicationStatus"]:
+        if not sDict.get("ApplicationStatus"):
             sDict["ApplicationStatus"] = "idem"
-        if not sDict["StatusSource"]:
-            sDict["StatusSource"] = "Unknown"
+        if not sDict.get("Source"):
+            sDict["Source"] = "Unknown"
 
         await job_logging_db.insert_record(
             job_id,
@@ -145,7 +139,7 @@ async def set_job_status(
             sDict["MinorStatus"],
             sDict["ApplicationStatus"],
             updTime,
-            sDict["StatusSource"],
+            sDict["Source"],
         )
 
     return SetJobStatusReturn(**job_data)
@@ -186,7 +180,7 @@ async def delete_jobs(
                         datetime.now(timezone.utc): JobStatusUpdate(
                             Status=JobStatus.DELETED,
                             MinorStatus="Checking accounting",
-                            StatusSource="job_manager",
+                            Source="job_manager",
                         )
                     },
                     job_db,
@@ -222,7 +216,7 @@ async def kill_jobs(
                         datetime.now(timezone.utc): JobStatusUpdate(
                             Status=JobStatus.KILLED,
                             MinorStatus="Marked for termination",
-                            StatusSource="job_manager",
+                            Source="job_manager",
                         )
                     },
                     job_db,
@@ -242,7 +236,7 @@ async def kill_jobs(
     #                 datetime.now(timezone.utc): JobStatusUpdate(
     #                     Status=JobStatus.KILLED,
     #                     MinorStatus="Marked for termination",
-    #                     StatusSource="job_manager",
+    #                     Source="job_manager",
     #                 )
     #             },
     #             job_db,

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/src/diracx/db/sql/utils.py

@@ -173,7 +173,7 @@ class BaseSQLDB(metaclass=ABCMeta):
             raise RuntimeError(f"{self.__class__} was used before entering")
         return cast(AsyncConnection, self._conn.get())
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         """
         Create a connection.
         This is called by the Dependency mechanism (see ``db_transaction``),

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/src/diracx_db.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diracx-db
-Version: 0.0.1a7
+Version: 0.0.1a8
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research

{diracx-db-0.0.1a7 → diracx-db-0.0.1a8}/tests/jobs/test_jobLoggingDB.py

@@ -32,9 +32,10 @@ async def test_insert_record(job_logging_db: JobLoggingDB):
 
         # Assert
         res = await job_logging_db.get_records(1)
+
         assert len(res) == 1
-        assert res[0].status == JobStatus.RECEIVED.value
-        assert res[0].minor_status == "minor_status"
-        assert res[0].application_status == "application_status"
-        assert res[0].status_time == date
-        assert res[0].status_source == "pytest"
+        assert res[0].Status == JobStatus.RECEIVED.value
+        assert res[0].MinorStatus == "minor_status"
+        assert res[0].ApplicationStatus == "application_status"
+        assert res[0].StatusTime == date
+        assert res[0].Source == "pytest"