diracx-db 0.0.1a47__tar.gz → 0.0.1a48__tar.gz

{diracx_db-0.0.1a47 → diracx_db-0.0.1a48}/PKG-INFO

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

{diracx_db-0.0.1a47 → diracx_db-0.0.1a48}/src/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-0.0.1a47 → diracx_db-0.0.1a48}/src/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-0.0.1a47 → diracx_db-0.0.1a48}/src/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-0.0.1a47 → diracx_db-0.0.1a48}/src/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-0.0.1a48/src/diracx/db/sql/utils/types.py

@@ -0,0 +1,137 @@
+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
+from sqlalchemy import DateTime, Enum
+
+from .functions import utcnow
+
+Column: partial[RawColumn] = partial(RawColumn, nullable=False)
+NullColumn: partial[RawColumn] = partial(RawColumn, nullable=True)
+DateNowColumn = partial(Column, type_=DateTime(timezone=True), server_default=utcnow())
+
+
+def EnumColumn(name, enum_type, **kwargs):  # noqa: N802
+    return Column(name, Enum(enum_type, native_enum=False, length=16), **kwargs)
+
+
+class EnumBackedBool(types.TypeDecorator):
+    """Maps a ``EnumBackedBool()`` column to True/False in Python."""
+
+    impl = types.Enum("True", "False", name="enum_backed_bool")
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect) -> str:
+        if value is True:
+            return "True"
+        elif value is False:
+            return "False"
+        else:
+            raise NotImplementedError(value, dialect)
+
+    def process_result_value(self, value, dialect) -> bool:
+        if value == "True":
+            return True
+        elif value == "False":
+            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 → diracx_db-0.0.1a48}/tests/jobs/test_job_db.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
 import pytest
+import sqlalchemy
 from sqlalchemy.exc import IntegrityError
 
 from diracx.core.exceptions import InvalidQueryError
@@ -47,6 +51,55 @@ async def populated_job_db(job_db):
     yield job_db
 
 
+async def test_bad_naive_datetime_used(populated_job_db):
+    async with populated_job_db as db:
+        compressed_jdl = "CompressedJDL0001111BadJob"
+        job_id = await db.create_job(compressed_jdl)
+        jobs_to_insert = {}
+        jobs_to_insert[job_id] = {
+            "JobID": job_id,
+            "Status": "New",
+            "Owner": "owner0101010101",
+            "OwnerGroup": "owner_group1",
+            "VO": "lhcb",
+            "HeartBeatTime": datetime.now().replace(tzinfo=None),  # noqa
+        }
+        with pytest.raises(sqlalchemy.exc.StatementError):
+            await db.insert_job_attributes(jobs_to_insert)  # should complain
+
+
+async def test_timezone_converted_back_to_utc(populated_job_db):
+    current_utc_dt = datetime.now(tz=ZoneInfo("UTC"))
+
+    async with populated_job_db as db:
+        compressed_jdl = "CompressedJDL0001111BadJob"
+        job_id = await db.create_job(compressed_jdl)
+        jobs_to_insert = {}
+        jobs_to_insert[job_id] = {
+            "JobID": job_id,
+            "Status": "New",
+            "Owner": "owner0101010101",
+            "OwnerGroup": "owner_group1",
+            "VO": "lhcb",
+            "HeartBeatTime": current_utc_dt.astimezone(ZoneInfo("Asia/Tokyo")),
+        }
+        await db.insert_job_attributes(jobs_to_insert)
+
+        total, result = await db.search(
+            ["JobID", "HeartBeatTime"],
+            [
+                ScalarSearchSpec(
+                    parameter="JobID",
+                    operator=ScalarSearchOperator.EQUAL,
+                    value=int(job_id),
+                )
+            ],
+            [],
+        )
+        assert total == 1
+        assert result[0]["HeartBeatTime"] == current_utc_dt, result
+
+
 async def test_search_parameters(populated_job_db):
     """Test that we can search specific parameters for jobs in the database."""
     async with populated_job_db as job_db:

diracx_db-0.0.1a47/src/diracx/db/sql/utils/types.py

@@ -1,43 +0,0 @@
-from __future__ import annotations
-
-from functools import partial
-
-import sqlalchemy.types as types
-from sqlalchemy import Column as RawColumn
-from sqlalchemy import DateTime, Enum
-
-from .functions import utcnow
-
-Column: partial[RawColumn] = partial(RawColumn, nullable=False)
-NullColumn: partial[RawColumn] = partial(RawColumn, nullable=True)
-DateNowColumn = partial(Column, type_=DateTime(timezone=True), server_default=utcnow())
-
-
-def EnumColumn(name, enum_type, **kwargs):  # noqa: N802
-    return Column(name, Enum(enum_type, native_enum=False, length=16), **kwargs)
-
-
-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")
-
-    def process_bind_param(self, value, dialect) -> str:
-        if value is True:
-            return "True"
-        elif value is False:
-            return "False"
-        else:
-            raise NotImplementedError(value, dialect)
-
-    def process_result_value(self, value, dialect) -> bool:
-        if value == "True":
-            return True
-        elif value == "False":
-            return False
-        else:
-            raise NotImplementedError(f"Unknown {value=}")