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

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

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

{diracx_db-0.0.1a46 → diracx_db-0.0.1a48}/src/diracx/db/os/job_parameters.py

@@ -9,13 +9,19 @@ class JobParametersDB(BaseOSDB):
     fields = {
         "JobID": {"type": "long"},
         "timestamp": {"type": "date"},
+        "PilotAgent": {"type": "keyword"},
+        "Pilot_Reference": {"type": "keyword"},
+        "JobGroup": {"type": "keyword"},
         "CPUNormalizationFactor": {"type": "long"},
         "NormCPUTime(s)": {"type": "long"},
-        "Memory(kB)": {"type": "long"},
+        "Memory(MB)": {"type": "long"},
+        "LocalAccount": {"type": "keyword"},
         "TotalCPUTime(s)": {"type": "long"},
-        "MemoryUsed(kb)": {"type": "long"},
-        "HostName": {"type": "keyword"},
+        "PayloadPID": {"type": "long"},
+        "HostName": {"type": "text"},
         "GridCE": {"type": "keyword"},
+        "CEQueue": {"type": "keyword"},
+        "BatchSystem": {"type": "keyword"},
         "ModelName": {"type": "keyword"},
         "Status": {"type": "keyword"},
         "JobType": {"type": "keyword"},

{diracx_db-0.0.1a46 → diracx_db-0.0.1a48}/src/diracx/db/sql/dummy/db.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from sqlalchemy import func, insert, select
+from sqlalchemy import insert
 from uuid_utils import UUID
 
-from diracx.db.sql.utils import BaseSQLDB, apply_search_filters
+from diracx.db.sql.utils import BaseSQLDB
 
 from .schema import Base as DummyDBBase
 from .schema import Cars, Owners
@@ -22,18 +22,7 @@ class DummyDB(BaseSQLDB):
     metadata = DummyDBBase.metadata
 
     async def summary(self, group_by, search) -> list[dict[str, str | int]]:
-        columns = [Cars.__table__.columns[x] for x in group_by]
-
-        stmt = select(*columns, func.count(Cars.license_plate).label("count"))
-        stmt = apply_search_filters(Cars.__table__.columns.__getitem__, stmt, search)
-        stmt = stmt.group_by(*columns)
-
-        # Execute the query
-        return [
-            dict(row._mapping)
-            async for row in (await self.conn.stream(stmt))
-            if row.count > 0  # type: ignore
-        ]
+        return await self._summary(Cars, group_by, search)
 
     async def insert_owner(self, name: str) -> int:
         stmt = insert(Owners).values(name=name)

{diracx_db-0.0.1a46 → diracx_db-0.0.1a48}/src/diracx/db/sql/job/db.py

@@ -5,15 +5,16 @@ __all__ = ["JobDB"]
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Iterable
 
-from sqlalchemy import bindparam, case, delete, func, 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
 
-from ..utils import BaseSQLDB, apply_search_filters, apply_sort_constraints
+from ..utils import BaseSQLDB, _get_columns
 from ..utils.functions import utcnow
 from .schema import (
     HeartBeatLoggingInfo,
@@ -25,17 +26,6 @@ from .schema import (
 )
 
 
-def _get_columns(table, parameters):
-    columns = [x for x in table.columns]
-    if parameters:
-        if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
-            raise InvalidQueryError(
-                f"Unrecognised parameters requested {unrecognised_parameters}"
-            )
-        columns = [c for c in columns if c.name in parameters]
-    return columns
-
-
 class JobDB(BaseSQLDB):
     metadata = JobDBBase.metadata
 
@@ -54,20 +44,11 @@ class JobDB(BaseSQLDB):
     # to find a way to make it dynamic
     jdl_2_db_parameters = ["JobName", "JobType", "JobGroup"]
 
-    async def summary(self, group_by, search) -> list[dict[str, str | int]]:
+    async def summary(
+        self, group_by: list[str], search: list[SearchSpec]
+    ) -> list[dict[str, str | int]]:
         """Get a summary of the jobs."""
-        columns = _get_columns(Jobs.__table__, group_by)
-
-        stmt = select(*columns, func.count(Jobs.job_id).label("count"))
-        stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
-        stmt = stmt.group_by(*columns)
-
-        # Execute the query
-        return [
-            dict(row._mapping)
-            async for row in (await self.conn.stream(stmt))
-            if row.count > 0  # type: ignore
-        ]
+        return await self._summary(table=Jobs, group_by=group_by, search=search)
 
     async def search(
         self,
@@ -80,34 +61,15 @@ class JobDB(BaseSQLDB):
         page: int | None = None,
     ) -> tuple[int, list[dict[Any, Any]]]:
         """Search for jobs in the database."""
-        # Find which columns to select
-        columns = _get_columns(Jobs.__table__, parameters)
-
-        stmt = select(*columns)
-
-        stmt = apply_search_filters(Jobs.__table__.columns.__getitem__, stmt, search)
-        stmt = apply_sort_constraints(Jobs.__table__.columns.__getitem__, stmt, sorts)
-
-        if distinct:
-            stmt = stmt.distinct()
-
-        # Calculate total count before applying pagination
-        total_count_subquery = stmt.alias()
-        total_count_stmt = select(func.count()).select_from(total_count_subquery)
-        total = (await self.conn.execute(total_count_stmt)).scalar_one()
-
-        # Apply pagination
-        if page is not None:
-            if page < 1:
-                raise InvalidQueryError("Page must be a positive integer")
-            if per_page < 1:
-                raise InvalidQueryError("Per page must be a positive integer")
-            stmt = stmt.offset((page - 1) * per_page).limit(per_page)
-
-        # Execute the query
-        return total, [
-            dict(row._mapping) async for row in (await self.conn.stream(stmt))
-        ]
+        return await self._search(
+            table=Jobs,
+            parameters=parameters,
+            search=search,
+            sorts=sorts,
+            distinct=distinct,
+            per_page=per_page,
+            page=page,
+        )
 
     async def create_job(self, compressed_original_jdl: str):
         """Used to insert a new job with original JDL. Returns inserted job id."""
@@ -167,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(
@@ -199,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
@@ -232,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.1a46 → 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.1a46 → 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.1a46 → diracx_db-0.0.1a48}/src/diracx/db/sql/utils/__init__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from .base import (
     BaseSQLDB,
     SQLDBUnavailableError,
+    _get_columns,
     apply_search_filters,
     apply_sort_constraints,
 )
@@ -10,6 +11,7 @@ from .functions import hash, substract_date, utcnow
 from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn
 
 __all__ = (
+    "_get_columns",
     "utcnow",
     "Column",
     "NullColumn",

{diracx_db-0.0.1a46 → diracx_db-0.0.1a48}/src/diracx/db/sql/utils/base.py

@@ -8,18 +8,23 @@ from abc import ABCMeta
 from collections.abc import AsyncIterator
 from contextvars import ContextVar
 from datetime import datetime
-from typing import Self, cast
+from typing import Any, Self, cast
 
 from pydantic import TypeAdapter
-from sqlalchemy import DateTime, MetaData, select
+from sqlalchemy import DateTime, MetaData, func, select
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.extensions import select_from_extension
-from diracx.core.models import SortDirection
+from diracx.core.models import (
+    SearchSpec,
+    SortDirection,
+    SortSpec,
+)
 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
 
@@ -227,6 +232,71 @@ class BaseSQLDB(metaclass=ABCMeta):
         except OperationalError as e:
             raise SQLDBUnavailableError("Cannot ping the DB") from e
 
+    async def _search(
+        self,
+        table: Any,
+        parameters: list[str] | None,
+        search: list[SearchSpec],
+        sorts: list[SortSpec],
+        *,
+        distinct: bool = False,
+        per_page: int = 100,
+        page: int | None = None,
+    ) -> tuple[int, list[dict[str, Any]]]:
+        """Search for elements in a table."""
+        # Find which columns to select
+        columns = _get_columns(table.__table__, parameters)
+
+        stmt = select(*columns)
+
+        stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
+        stmt = apply_sort_constraints(table.__table__.columns.__getitem__, stmt, sorts)
+
+        if distinct:
+            stmt = stmt.distinct()
+
+        # Calculate total count before applying pagination
+        total_count_subquery = stmt.alias()
+        total_count_stmt = select(func.count()).select_from(total_count_subquery)
+        total = (await self.conn.execute(total_count_stmt)).scalar_one()
+
+        # Apply pagination
+        if page is not None:
+            if page < 1:
+                raise InvalidQueryError("Page must be a positive integer")
+            if per_page < 1:
+                raise InvalidQueryError("Per page must be a positive integer")
+            stmt = stmt.offset((page - 1) * per_page).limit(per_page)
+
+        # Execute the query
+        return total, [
+            dict(row._mapping) async for row in (await self.conn.stream(stmt))
+        ]
+
+    async def _summary(
+        self, table: Any, group_by: list[str], search: list[SearchSpec]
+    ) -> list[dict[str, str | int]]:
+        """Get a summary of the elements of a table."""
+        columns = _get_columns(table.__table__, group_by)
+
+        pk_columns = list(table.__table__.primary_key.columns)
+        if not pk_columns:
+            raise ValueError(
+                "Model has no primary key and no count_column was provided."
+            )
+        count_col = pk_columns[0]
+
+        stmt = select(*columns, func.count(count_col).label("count"))
+        stmt = apply_search_filters(table.__table__.columns.__getitem__, stmt, search)
+        stmt = stmt.group_by(*columns)
+
+        # Execute the query
+        return [
+            dict(row._mapping)
+            async for row in (await self.conn.stream(stmt))
+            if row.count > 0  # type: ignore
+        ]
+
 
 def find_time_resolution(value):
     if isinstance(value, datetime):
@@ -258,6 +328,17 @@ def find_time_resolution(value):
     raise InvalidQueryError(f"Cannot parse {value=}")
 
 
+def _get_columns(table, parameters):
+    columns = [x for x in table.columns]
+    if parameters:
+        if unrecognised_parameters := set(parameters) - set(table.columns.keys()):
+            raise InvalidQueryError(
+                f"Unrecognised parameters requested {unrecognised_parameters}"
+            )
+        columns = [c for c in columns if c.name in parameters]
+    return columns
+
+
 def apply_search_filters(column_mapping, stmt, search):
     for query in search:
         try:
@@ -265,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:
@@ -300,6 +381,15 @@ def apply_search_filters(column_mapping, stmt, search):
             expr = column.like(query["value"])
         elif query["operator"] in "ilike":
             expr = column.ilike(query["value"])
+        elif query["operator"] == "not like":
+            expr = column.not_like(query["value"])
+        elif query["operator"] == "regex":
+            # We check the regex validity here
+            try:
+                re.compile(query["value"])
+            except re.error as e:
+                raise InvalidQueryError(f"Invalid regex {query['value']}") from e
+            expr = column.regexp_match(query["value"])
         else:
             raise InvalidQueryError(f"Unknown filter {query=}")
         stmt = stmt.where(expr)

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.1a46 → 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:
@@ -187,6 +240,55 @@ async def test_search_conditions(populated_job_db):
         assert total == 0
         assert not result
 
+        # Search for a specific scalar condition: Owner not like 'owner1%'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.NOT_LIKE, value="owner1%"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100 - 11
+        assert result
+        assert len(result) == 100 - 11
+        assert all(not r["Owner"].startswith("owner1") for r in result)
+
+        # Search for a specific scalar condition: OwnerGroup not like 'owner_group2'
+        condition = ScalarSearchSpec(
+            parameter="OwnerGroup",
+            operator=ScalarSearchOperator.NOT_LIKE,
+            value="owner_group2",
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100 - 50
+        assert result
+        assert len(result) == 100 - 50
+        assert all(not r["OwnerGroup"] == "owner_group2" for r in result)
+
+        # Search for a specific scalar condition: Owner regex '^owner\d+$'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="^owner\\d+$"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 100
+        assert result
+        assert len(result) == 100
+
+        # Search for a specific scalar condition: JobID regex 'owner[0-3]+'
+        # owner0, owner1, owner2, owner3 (4 jobs)
+        # owner11 -> owner39 (30 jobs)
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="owner[0-3]+"
+        )
+        total, result = await job_db.search([], [condition], [])
+        assert total == 34
+        assert result
+        assert len(result) == 34
+
+        # Search for a specific scalar condition: JobID regex 'owner[1-'
+        condition = ScalarSearchSpec(
+            parameter="Owner", operator=ScalarSearchOperator.REGEX, value="owner[1-"
+        )
+        with pytest.raises(InvalidQueryError):
+            await job_db.search([], [condition], [])
+
 
 async def test_search_sorts(populated_job_db):
     """Test that we can search for jobs in the database and sort the results."""

{diracx_db-0.0.1a46 → diracx_db-0.0.1a48}/tests/test_dummy_db.py

@@ -129,7 +129,7 @@ async def test_failed_transaction(dummy_db):
 
     # The connection is created when the context manager is entered
     # This is our transaction
-    with pytest.raises(KeyError):
+    with pytest.raises(InvalidQueryError):
         async with dummy_db as dummy_db:
             assert dummy_db.conn
 

diracx_db-0.0.1a46/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=}")