diracx-db 0.0.1a47__py3-none-any.whl → 0.0.1a49__py3-none-any.whl

diracx/db/sql/job/db.py

@@ -5,10 +5,11 @@ __all__ = ["JobDB"]
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Iterable
 
-from sqlalchemy import bindparam, case, delete, insert, select, update
+from sqlalchemy import bindparam, case, delete, literal, select, update
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BindParameter
+from sqlalchemy.sql import expression
 
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.models import JobCommand, SearchSpec, SortSpec
@@ -128,27 +129,14 @@ class JobDB(BaseSQLDB):
             ],
         )
 
-    @staticmethod
-    def _set_job_attributes_fix_value(column, value):
-        """Apply corrections to the values before inserting them into the database.
-
-        TODO: Move this logic into the sqlalchemy model.
-        """
-        if column == "VerifiedFlag":
-            value_str = str(value)
-            if value_str in ("True", "False"):
-                return value_str
-        if column == "AccountedFlag":
-            value_str = str(value)
-            if value_str in ("True", "False", "Failed"):
-                return value_str
-        else:
-            return value
-        raise NotImplementedError(f"Unrecognized value for column {column}: {value}")
-
     async def set_job_attributes(self, job_data):
         """Update the parameters of the given jobs."""
         # TODO: add myDate and force parameters.
+
+        if not job_data:
+            # nothing to do!
+            raise ValueError("job_data is empty")
+
         for job_id in job_data.keys():
             if "Status" in job_data[job_id]:
                 job_data[job_id].update(
@@ -160,7 +148,11 @@ class JobDB(BaseSQLDB):
                 *[
                     (
                         Jobs.__table__.c.JobID == job_id,
-                        self._set_job_attributes_fix_value(column, attrs[column]),
+                        # Since the setting of the new column value is obscured by the CASE statement,
+                        # ensure that SQLAlchemy renders the new column value with the correct type
+                        literal(attrs[column], type_=Jobs.__table__.c[column].type)
+                        if not isinstance(attrs[column], expression.FunctionElement)
+                        else attrs[column],
                     )
                     for job_id, attrs in job_data.items()
                     if column in attrs
@@ -193,7 +185,7 @@ class JobDB(BaseSQLDB):
     async def set_job_commands(self, commands: list[tuple[int, str, str]]) -> None:
         """Store a command to be passed to the job together with the next heart beat."""
         await self.conn.execute(
-            insert(JobCommands),
+            JobCommands.__table__.insert(),
             [
                 {
                     "JobID": job_id,

diracx/db/sql/job/schema.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import sqlalchemy.types as types
 from sqlalchemy import (
-    DateTime,
     ForeignKey,
     Index,
     Integer,
@@ -11,6 +10,8 @@ from sqlalchemy import (
 )
 from sqlalchemy.orm import declarative_base
 
+from diracx.db.sql.utils.types import SmarterDateTime
+
 from ..utils import Column, EnumBackedBool, NullColumn
 
 JobDBBase = declarative_base()
@@ -19,11 +20,8 @@ JobDBBase = declarative_base()
 class AccountedFlagEnum(types.TypeDecorator):
     """Maps a ``AccountedFlagEnum()`` column to True/False in Python."""
 
-    impl = types.Enum
-    cache_ok: bool = True
-
-    def __init__(self) -> None:
-        super().__init__("True", "False", "Failed")
+    impl = types.Enum("True", "False", "Failed", name="accounted_flag_enum")
+    cache_ok = True
 
     def process_bind_param(self, value, dialect) -> str:
         if value is True:
@@ -63,12 +61,30 @@ class Jobs(JobDBBase):
     owner = Column("Owner", String(64), default="Unknown")
     owner_group = Column("OwnerGroup", String(128), default="Unknown")
     vo = Column("VO", String(32))
-    submission_time = NullColumn("SubmissionTime", DateTime)
-    reschedule_time = NullColumn("RescheduleTime", DateTime)
-    last_update_time = NullColumn("LastUpdateTime", DateTime)
-    start_exec_time = NullColumn("StartExecTime", DateTime)
-    heart_beat_time = NullColumn("HeartBeatTime", DateTime)
-    end_exec_time = NullColumn("EndExecTime", DateTime)
+    submission_time = NullColumn(
+        "SubmissionTime",
+        SmarterDateTime(),
+    )
+    reschedule_time = NullColumn(
+        "RescheduleTime",
+        SmarterDateTime(),
+    )
+    last_update_time = NullColumn(
+        "LastUpdateTime",
+        SmarterDateTime(),
+    )
+    start_exec_time = NullColumn(
+        "StartExecTime",
+        SmarterDateTime(),
+    )
+    heart_beat_time = NullColumn(
+        "HeartBeatTime",
+        SmarterDateTime(),
+    )
+    end_exec_time = NullColumn(
+        "EndExecTime",
+        SmarterDateTime(),
+    )
     status = Column("Status", String(32), default="Received")
     minor_status = Column("MinorStatus", String(128), default="Unknown")
     application_status = Column("ApplicationStatus", String(255), default="Unknown")
@@ -143,7 +159,11 @@ class HeartBeatLoggingInfo(JobDBBase):
     )
     name = Column("Name", String(100), primary_key=True)
     value = Column("Value", Text)
-    heart_beat_time = Column("HeartBeatTime", DateTime, primary_key=True)
+    heart_beat_time = Column(
+        "HeartBeatTime",
+        SmarterDateTime(),
+        primary_key=True,
+    )
 
 
 class JobCommands(JobDBBase):
@@ -154,5 +174,12 @@ class JobCommands(JobDBBase):
     command = Column("Command", String(100))
     arguments = Column("Arguments", String(100))
     status = Column("Status", String(64), default="Received")
-    reception_time = Column("ReceptionTime", DateTime, primary_key=True)
-    execution_time = NullColumn("ExecutionTime", DateTime)
+    reception_time = Column(
+        "ReceptionTime",
+        SmarterDateTime(),
+        primary_key=True,
+    )
+    execution_time = NullColumn(
+        "ExecutionTime",
+        SmarterDateTime(),
+    )

diracx/db/sql/pilot_agents/schema.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from sqlalchemy import (
-    DateTime,
     Double,
     Index,
     Integer,
@@ -10,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()
@@ -29,8 +30,8 @@ class PilotAgents(PilotAgentsDBBase):
     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)
+    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)
@@ -47,7 +48,7 @@ class JobToPilotMapping(PilotAgentsDBBase):
 
     pilot_id = Column("PilotID", Integer, primary_key=True)
     job_id = Column("JobID", Integer, primary_key=True)
-    start_time = Column("StartTime", DateTime)
+    start_time = Column("StartTime", SmarterDateTime)
 
     __table_args__ = (Index("JobID", "JobID"), Index("PilotID", "PilotID"))
 

diracx/db/sql/utils/base.py

@@ -24,6 +24,7 @@ from diracx.core.models import (
 )
 from diracx.core.settings import SqlalchemyDsn
 from diracx.db.exceptions import DBUnavailableError
+from diracx.db.sql.utils.types import SmarterDateTime
 
 from .functions import date_trunc
 
@@ -345,7 +346,7 @@ def apply_search_filters(column_mapping, stmt, search):
         except KeyError as e:
             raise InvalidQueryError(f"Unknown column {query['parameter']}") from e
 
-        if isinstance(column.type, DateTime):
+        if isinstance(column.type, (DateTime, SmarterDateTime)):
             if "value" in query and isinstance(query["value"], str):
                 resolution, value = find_time_resolution(query["value"])
                 if resolution:

diracx/db/sql/utils/types.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+from datetime import datetime
 from functools import partial
+from zoneinfo import ZoneInfo
 
 import sqlalchemy.types as types
 from sqlalchemy import Column as RawColumn
@@ -20,11 +22,8 @@ def EnumColumn(name, enum_type, **kwargs):  # noqa: N802
 class EnumBackedBool(types.TypeDecorator):
     """Maps a ``EnumBackedBool()`` column to True/False in Python."""
 
-    impl = types.Enum
-    cache_ok: bool = True
-
-    def __init__(self) -> None:
-        super().__init__("True", "False")
+    impl = types.Enum("True", "False", name="enum_backed_bool")
+    cache_ok = True
 
     def process_bind_param(self, value, dialect) -> str:
         if value is True:
@@ -41,3 +40,98 @@ class EnumBackedBool(types.TypeDecorator):
             return False
         else:
             raise NotImplementedError(f"Unknown {value=}")
+
+
+class SmarterDateTime(types.TypeDecorator):
+    """A DateTime type that also accepts ISO8601 strings.
+
+    Takes into account converting timezone aware datetime objects into
+    naive form and back when needed.
+
+    """
+
+    impl = DateTime()
+    cache_ok = True
+
+    def __init__(
+        self,
+        stored_tz: ZoneInfo | None = ZoneInfo("UTC"),
+        returned_tz: ZoneInfo = ZoneInfo("UTC"),
+        stored_naive_sqlite=True,
+        stored_naive_mysql=True,
+        stored_naive_postgres=False,  # Forces timezone-awareness
+    ):
+        self._stored_naive_dialect = {
+            "sqlite": stored_naive_sqlite,
+            "mysql": stored_naive_mysql,
+            "postgres": stored_naive_postgres,
+        }
+        self._stored_tz: ZoneInfo | None = stored_tz  # None = Local timezone
+        self._returned_tz: ZoneInfo = returned_tz
+
+    def _stored_naive(self, dialect):
+        if dialect.name not in self._stored_naive_dialect:
+            raise NotImplementedError(dialect.name)
+        return self._stored_naive_dialect.get(dialect.name)
+
+    def process_bind_param(self, value, dialect):
+        if value is None:
+            return None
+
+        if isinstance(value, str):
+            try:
+                value: datetime = datetime.fromisoformat(value)
+            except ValueError as err:
+                raise ValueError(f"Unable to parse datetime string: {value}") from err
+
+        if not isinstance(value, datetime):
+            raise ValueError(f"Expected datetime or ISO8601 string, but got {value!r}")
+
+        if not value.tzinfo:
+            raise ValueError(
+                f"Provided timestamp {value=} has no tzinfo -"
+                " this is problematic and may cause inconsistencies in stored timestamps.\n"
+                " Please always work with tz-aware datetimes / attach tzinfo to your datetime objects:"
+                " e.g. datetime.now(tz=timezone.utc) or use datetime_obj.astimezone() with no arguments if you need to "
+                "attach the local timezone to a local naive timestamp."
+            )
+
+        # Check that we need to convert the timezone to match self._stored_tz timezone:
+        if self._stored_naive(dialect):
+            # if self._stored_tz is None, we use our local/system timezone.
+            stored_tz = self._stored_tz
+
+            # astimezone converts to the stored timezone (local timezone if None)
+            # replace strips the TZ info --> naive datetime object
+            value = value.astimezone(tz=stored_tz).replace(tzinfo=None)
+
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        if not isinstance(value, datetime):
+            raise NotImplementedError(f"{value=} not a datetime object")
+
+        if self._stored_naive(dialect):
+            # Here we add back the tzinfo to the naive timestamp
+            # from the DB to make it aware again.
+            if value.tzinfo is None:
+                # we are definitely given a naive timestamp, so handle it.
+                # add back the timezone info if stored_tz is set
+                if self._stored_tz:
+                    value = value.replace(tzinfo=self._stored_tz)
+                else:
+                    # if stored as a local time, add back the system timezone info...
+                    value = value.astimezone()
+            else:
+                raise ValueError(
+                    f"stored_naive is True for {dialect.name=}, but the database engine returned "
+                    "a tz-aware datetime. You need to check the SQLAlchemy model is consistent with the DB schema."
+                )
+
+        # finally, convert the datetime according to the "returned_tz"
+        value = value.astimezone(self._returned_tz)
+
+        # phew...
+        return value

{diracx_db-0.0.1a47.dist-info → diracx_db-0.0.1a49.dist-info}/METADATA

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

{diracx_db-0.0.1a47.dist-info → diracx_db-0.0.1a49.dist-info}/RECORD

@@ -13,14 +13,14 @@ diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 diracx/db/sql/dummy/db.py,sha256=MKSUSJI1BlRgK08tjCfkCkOz02asvJAeBw60pAdiGV8,1212
 diracx/db/sql/dummy/schema.py,sha256=9zI53pKlzc6qBezsyjkatOQrNZdGCjwgjQ8Iz_pyAXs,789
 diracx/db/sql/job/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/job/db.py,sha256=TsHbMVUO-87228hVbodGQclTgY2b7fI0XBsbNbCVgc4,10298
-diracx/db/sql/job/schema.py,sha256=eFgZshe6NEzOM2qI0HI9Y3abrqDMoQIwa9L0vZugHcU,5431
+diracx/db/sql/job/db.py,sha256=bX-4OMyW4h9tqeTE3OvonxTXlL6j_Qvv9uEtK5SthN8,10120
+diracx/db/sql/job/schema.py,sha256=fJdmiLp6psdAjo_CoBfSAGSYk2NJkSBwvik9tznESD0,5740
 diracx/db/sql/job_logging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/job_logging/db.py,sha256=hyklARuEj3R1sSJ8UaObRprmsRx7RjbKAcbfgT9BwRg,5496
 diracx/db/sql/job_logging/schema.py,sha256=k6uBw-RHAcJ5GEleNpiWoXEJBhCiNG-y4xAgBKHZjjM,2524
 diracx/db/sql/pilot_agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/pilot_agents/db.py,sha256=6CQ0QGV4NhsGKVCygEtE4kmIjT89xJwrIMuYZTslWFE,1231
-diracx/db/sql/pilot_agents/schema.py,sha256=KeWnFSpYOTrT3-_rOCFjbjNnPNXKnUZiJVsu4vv5U2U,2149
+diracx/db/sql/pilot_agents/schema.py,sha256=BTFLuiwcxAvAtTvTP9C7DbGtXoM-IHVDG9k7HMx62AA,2211
 diracx/db/sql/sandbox_metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/sql/sandbox_metadata/db.py,sha256=FtyPx6GAGJAH-lmuw8PQj6_KGHG6t3AC3-E9uWf-JNs,10236
 diracx/db/sql/sandbox_metadata/schema.py,sha256=V5gV2PHwzTbBz_th9ribLfE7Lqk8YGemDmvqq4jWQJ4,1530
@@ -28,10 +28,10 @@ diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 diracx/db/sql/task_queue/db.py,sha256=2qul1D2tX2uCI92N591WK5xWHakG0pNibzDwKQ7W-I8,6246
 diracx/db/sql/task_queue/schema.py,sha256=5efAgvNYRkLlaJ2NzRInRfmVa3tyIzQu2l0oRPy4Kzw,3258
 diracx/db/sql/utils/__init__.py,sha256=XYbv-AJAPl7bb8dETpjc07olmtXQ0h1MFUbLqjAphQE,585
-diracx/db/sql/utils/base.py,sha256=snZFJmUJV-wweZLpio29MxuPFghfugpVMDC1iE_jM7w,15568
+diracx/db/sql/utils/base.py,sha256=Dn7a-ICuFNTT3w0QtbS63uF1S4Wnn5NK3cu6Pzn4a6A,15641
 diracx/db/sql/utils/functions.py,sha256=_E4tc9Gti6LuSh7QEyoqPJSvCuByVqvRenOXCzxsulE,4014
-diracx/db/sql/utils/types.py,sha256=yU-tXsu6hFGPsr9ba1n3ZjGPnHQI_06lbpkTeDCWJtg,1287
-diracx_db-0.0.1a47.dist-info/METADATA,sha256=z9GxhxY4-mwWkwLQ1Ue72mU28cXPAB2W4tr_UU2J6yA,675
-diracx_db-0.0.1a47.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-diracx_db-0.0.1a47.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
-diracx_db-0.0.1a47.dist-info/RECORD,,
+diracx/db/sql/utils/types.py,sha256=KNZWJfpvHTjfIPg6Nn7zY-rS0q3ybnirHcTcLAYSYbE,5118
+diracx_db-0.0.1a49.dist-info/METADATA,sha256=hPW3mb1Ain8hSsRRDWlL3kGeRNKpYc1bHtJ9hlSALIE,675
+diracx_db-0.0.1a49.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+diracx_db-0.0.1a49.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
+diracx_db-0.0.1a49.dist-info/RECORD,,