diracx-db 0.0.1a6__py3-none-any.whl → 0.0.1a8__py3-none-any.whl

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/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/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/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.1a6.dist-info → diracx_db-0.0.1a8.dist-info}/METADATA

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

{diracx_db-0.0.1a6.dist-info → diracx_db-0.0.1a8.dist-info}/RECORD

@@ -6,7 +6,7 @@ diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
 diracx/db/os/job_parameters.py,sha256=Knca19uT2G-5FI7MOFlaOAXeHn4ecPVLIH30TiwhaTw,858
 diracx/db/os/utils.py,sha256=MACBZSn9mzbaY_HAPqk-OBmLI2OYCph_1aI7oqlr-K0,9124
 diracx/db/sql/__init__.py,sha256=R6tk5lo1EHbt8joGDesesYHcc1swIq9T4AaSixhh7lA,252
-diracx/db/sql/utils.py,sha256=DAqqPCQIloWyK_Tpv7CXP3hrO1nq1Bw3CC7Nwlg5Ji4,7849
+diracx/db/sql/utils.py,sha256=BuXjIuXN-_v8YkCoMoMhw2tHVUqG6lTBx-e4VEYWE8o,7857
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/auth/db.py,sha256=7GAdwD7g2YR86yQkP9ecuIbFd1h_wxVKnF_GfXfZqLA,9915
 diracx/db/sql/auth/schema.py,sha256=wutCjZ_uz21J0HHZjwoOXq3cLdlNY2lCR390yIJ_T60,2891
@@ -14,14 +14,14 @@ diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 diracx/db/sql/dummy/db.py,sha256=5PIPv6aKY7CGIwmvnGKowjVr9ZQWpbjFSd2PIX7YOUw,1627
 diracx/db/sql/dummy/schema.py,sha256=uEkGDNVZbmJecytkHY1CO-M1MiKxe5w1_h0joJMPC9E,680
 diracx/db/sql/jobs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/jobs/db.py,sha256=cprktaw2XuvjEsSv6GD1vPWmJTGjmb7RAIaXigufM8w,28337
-diracx/db/sql/jobs/schema.py,sha256=6K8zn7kqvfPsCTfuCntSblZU-wrAwzziLW_b26s5KDE,10128
-diracx/db/sql/jobs/status_utility.py,sha256=VLFTq6npFF6ACeKofxg70zOgZMOmjD7ynYGYKVAYQAc,10843
+diracx/db/sql/jobs/db.py,sha256=-hzxAoL6Pukpaf3lpMTrzR4s0Md3jENHVxRrmUBKS60,29781
+diracx/db/sql/jobs/schema.py,sha256=5KCMciWOymJqinIwd-gdHZXrMH8p9cPz7fZo0lYHDXk,10143
+diracx/db/sql/jobs/status_utility.py,sha256=0kAt623nh1O5wgsgktctdCmHEynO1nU0vn-7zakNeOA,10525
 diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/sandbox_metadata/db.py,sha256=HcmXlnObD10Lwe_VMyloTRa6TFZgEKOJn7ud8AkisbE,3464
 diracx/db/sql/sandbox_metadata/schema.py,sha256=VekS-tJqXFqn6DE3bMgZqKyBVkgFCm0mfPx7s9rZ6Bg,1372
-diracx_db-0.0.1a6.dist-info/METADATA,sha256=eFneeg6C3nt-R4ahzp7Tss_SP_YVKv3_VL8R5nYiMb0,680
-diracx_db-0.0.1a6.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-diracx_db-0.0.1a6.dist-info/entry_points.txt,sha256=xEFGu_zgmPgQPlUeFtdahQfQIboJ1ugFOK8eMio9gtw,271
-diracx_db-0.0.1a6.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
-diracx_db-0.0.1a6.dist-info/RECORD,,
+diracx_db-0.0.1a8.dist-info/METADATA,sha256=XUbk4aS7IZCL8ie9Rcpg57GSmiWjOolNF07ExUEd-DI,680
+diracx_db-0.0.1a8.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+diracx_db-0.0.1a8.dist-info/entry_points.txt,sha256=xEFGu_zgmPgQPlUeFtdahQfQIboJ1ugFOK8eMio9gtw,271
+diracx_db-0.0.1a8.dist-info/top_level.txt,sha256=vJx10tdRlBX3rF2Psgk5jlwVGZNcL3m_7iQWwgPXt-U,7
+diracx_db-0.0.1a8.dist-info/RECORD,,