diracx-db 0.0.1a46__py3-none-any.whl → 0.0.6__py3-none-any.whl

diracx/db/__main__.py

@@ -39,6 +39,7 @@ async def init_sql():
                 if db._db_url.startswith("sqlite"):
                     await conn.exec_driver_sql("PRAGMA foreign_keys=ON")
                 await conn.run_sync(db.metadata.create_all)
+                await db.post_create(conn)
 
 
 async def init_os():

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/sql/auth/db.py

@@ -1,16 +1,21 @@
 from __future__ import annotations
 
+import logging
 import secrets
+from datetime import UTC, datetime
+from itertools import pairwise
 
-from sqlalchemy import insert, select, update
+from dateutil.rrule import MONTHLY, rrule
+from sqlalchemy import insert, select, text, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
+from sqlalchemy.ext.asyncio import AsyncConnection
 from uuid_utils import UUID, uuid7
 
 from diracx.core.exceptions import (
     AuthorizationError,
     TokenNotFoundError,
 )
-from diracx.db.sql.utils import BaseSQLDB, hash, substract_date
+from diracx.db.sql.utils import BaseSQLDB, hash, substract_date, uuid7_from_datetime
 
 from .schema import (
     AuthorizationFlows,
@@ -25,10 +30,72 @@ from .schema import Base as AuthDBBase
 USER_CODE_ALPHABET = "BCDFGHJKLMNPQRSTVWXZ"
 MAX_RETRY = 5
 
+logger = logging.getLogger(__name__)
+
 
 class AuthDB(BaseSQLDB):
     metadata = AuthDBBase.metadata
 
+    @classmethod
+    async def post_create(cls, conn: AsyncConnection) -> None:
+        """Create partitions if it is a MySQL DB and it does not have
+        it yet and the table does not have any data yet.
+        We do this as a post_create step as sqlalchemy does not support
+        partition so well.
+        """
+        if conn.dialect.name == "mysql":
+            check_partition_query = text(
+                "SELECT PARTITION_NAME FROM information_schema.partitions "
+                "WHERE TABLE_NAME = 'RefreshTokens' AND PARTITION_NAME is not NULL"
+            )
+            partition_names = (await conn.execute(check_partition_query)).all()
+
+            if not partition_names:
+                # Create a monthly partition from today until 2 years
+                # The partition are named p_<year>_<month>
+                start_date = datetime.now(tz=UTC).replace(
+                    day=1, hour=0, minute=0, second=0, microsecond=0
+                )
+                end_date = start_date.replace(year=start_date.year + 2)
+
+                dates = [
+                    dt for dt in rrule(MONTHLY, dtstart=start_date, until=end_date)
+                ]
+
+                partition_list = []
+                for name, limit in pairwise(dates):
+                    partition_list.append(
+                        f"PARTITION p_{name.year}_{name.month} "
+                        f"VALUES LESS THAN ('{str(uuid7_from_datetime(limit, randomize=False)).replace('-', '')}')"
+                    )
+                partition_list.append("PARTITION p_future VALUES LESS THAN (MAXVALUE)")
+
+                alter_query = text(
+                    f"ALTER TABLE RefreshTokens PARTITION BY RANGE COLUMNS (JTI) ({','.join(partition_list)})"
+                )
+
+                check_table_empty_query = text("SELECT * FROM RefreshTokens LIMIT 1")
+                refresh_table_content = (
+                    await conn.execute(check_table_empty_query)
+                ).all()
+                if refresh_table_content:
+                    logger.warning(
+                        "RefreshTokens table not empty. Run the following query yourself"
+                    )
+                    logger.warning(alter_query)
+                    return
+
+                await conn.execute(alter_query)
+
+                partition_names = (
+                    await conn.execute(
+                        check_partition_query, {"table_name": "RefreshTokens"}
+                    )
+                ).all()
+                assert partition_names, (
+                    f"There should be partitions now {partition_names}"
+                )
+
     async def device_flow_validate_user_code(
         self, user_code: str, max_validity: int
     ) -> str:
@@ -111,6 +178,10 @@ class AuthDB(BaseSQLDB):
                 await self.conn.execute(stmt)
 
             except IntegrityError:
+                logger.warning(
+                    "Device flow code collision detected, retrying (user_code=%s)",
+                    user_code,
+                )
                 continue
 
             return user_code, device_code

diracx/db/sql/auth/schema.py

@@ -10,7 +10,12 @@ from sqlalchemy import (
 )
 from sqlalchemy.orm import declarative_base
 
-from diracx.db.sql.utils import Column, DateNowColumn, EnumColumn, NullColumn
+from diracx.db.sql.utils import (
+    Column,
+    DateNowColumn,
+    EnumColumn,
+    NullColumn,
+)
 
 USER_CODE_LENGTH = 8
 
@@ -92,7 +97,6 @@ class RefreshTokens(Base):
     status = EnumColumn(
         "Status", RefreshTokenStatus, server_default=RefreshTokenStatus.CREATED.name
     )
-    creation_time = DateNowColumn("CreationTime", index=True)
     scope = Column("Scope", String(1024))
 
     # User attributes bound to the refresh token

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/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/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/__init__.py

@@ -1,15 +1,7 @@
 from __future__ import annotations
 
-from .base import (
-    BaseSQLDB,
-    SQLDBUnavailableError,
-    apply_search_filters,
-    apply_sort_constraints,
-)
-from .functions import hash, substract_date, utcnow
-from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn
-
-__all__ = (
+__all__ = [
+    "_get_columns",
     "utcnow",
     "Column",
     "NullColumn",
@@ -22,4 +14,18 @@ __all__ = (
     "substract_date",
     "hash",
     "SQLDBUnavailableError",
+    "uuid7_from_datetime",
+    "uuid7_to_datetime",
+]
+
+from .base import (
+    BaseSQLDB,
+    SQLDBUnavailableError,
+    _get_columns,
+    apply_search_filters,
+    apply_sort_constraints,
+    uuid7_from_datetime,
+    uuid7_to_datetime,
 )
+from .functions import hash, substract_date, utcnow
+from .types import Column, DateNowColumn, EnumBackedBool, EnumColumn, NullColumn

diracx/db/sql/utils/base.py

@@ -7,19 +7,26 @@ import re
 from abc import ABCMeta
 from collections.abc import AsyncIterator
 from contextvars import ContextVar
-from datetime import datetime
-from typing import Self, cast
+from datetime import datetime, timezone
+from typing import Any, Self, cast
+from uuid import UUID as StdUUID  # noqa: N811
 
 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 uuid_utils import UUID, uuid7
 
 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
 
@@ -148,6 +155,11 @@ class BaseSQLDB(metaclass=ABCMeta):
                     raise
         return db_urls
 
+    @classmethod
+    async def post_create(cls, conn: AsyncConnection) -> None:
+        """Execute actions after the schema has been created."""
+        return
+
     @classmethod
     def transaction(cls) -> Self:
         raise NotImplementedError("This should never be called")
@@ -199,6 +211,12 @@ class BaseSQLDB(metaclass=ABCMeta):
         try:
             self._conn.set(await self.engine.connect().__aenter__())
         except Exception as e:
+            logger.warning(
+                "Database connection failed for %s: %s",
+                self.__class__.__name__,
+                e,
+                exc_info=True,
+            )
             raise SQLDBUnavailableError(
                 f"Cannot connect to {self.__class__.__name__}"
             ) from e
@@ -227,6 +245,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 +341,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 +359,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 +394,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)
@@ -326,3 +429,33 @@ def apply_sort_constraints(column_mapping, stmt, sorts):
     if sort_columns:
         stmt = stmt.order_by(*sort_columns)
     return stmt
+
+
+def uuid7_to_datetime(uuid: UUID | StdUUID | str) -> datetime:
+    """Convert a UUIDv7 to a datetime."""
+    if isinstance(uuid, StdUUID):
+        # Convert stdlib UUID to uuid_utils.UUID
+        uuid = UUID(str(uuid))
+    elif not isinstance(uuid, UUID):
+        # Convert string or other types to uuid_utils.UUID
+        uuid = UUID(uuid)
+    if uuid.version != 7:
+        raise ValueError(f"UUID {uuid} is not a UUIDv7")
+    return datetime.fromtimestamp(uuid.timestamp / 1000.0, tz=timezone.utc)
+
+
+def uuid7_from_datetime(dt: datetime, *, randomize: bool = True) -> UUID:
+    """Generate a UUIDv7 corresponding to the given datetime.
+
+    If randomize is True, the standard uuid7 function is used resulting in the
+    lowest 62-bits being random. If randomize is False, the UUIDv7 will be the
+    lowest possible UUIDv7 for the given datetime.
+    """
+    timestamp = dt.timestamp()
+    if randomize:
+        uuid = uuid7(int(timestamp), int((timestamp % 1) * 1e9))
+    else:
+        time_high = int(timestamp * 1000) >> 16
+        time_low = int(timestamp * 1000) & 0xFFFF
+        uuid = UUID.from_fields((time_high, time_low, 0x7000, 0x80, 0, 0))
+    return uuid

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.1a46.dist-info → diracx_db-0.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-db
-Version: 0.0.1a46
+Version: 0.0.6
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -12,8 +12,11 @@ Requires-Python: >=3.11
 Requires-Dist: diracx-core
 Requires-Dist: opensearch-py[async]
 Requires-Dist: pydantic>=2.10
+Requires-Dist: python-dateutil
 Requires-Dist: sqlalchemy[aiomysql,aiosqlite]>=2
 Requires-Dist: uuid-utils
 Provides-Extra: testing
 Requires-Dist: diracx-testing; extra == 'testing'
 Requires-Dist: freezegun; extra == 'testing'
+Provides-Extra: types
+Requires-Dist: types-python-dateutil; extra == 'types'

{diracx_db-0.0.1a46.dist-info → diracx_db-0.0.6.dist-info}/RECORD

@@ -1,37 +1,37 @@
 diracx/db/__init__.py,sha256=2oeUeVwZq53bo_ZOflEYZsBn7tcR5Tzb2AIu0TAWELM,109
-diracx/db/__main__.py,sha256=3yaUP1ig-yaPSQM4wy6CtSXXHivQg-hIz2FeBt7joBc,1714
+diracx/db/__main__.py,sha256=6YlmpiU1cLLHjKLy1DfdEOQUyvSla-MbJsJ7aQwAOVs,1757
 diracx/db/exceptions.py,sha256=1nn-SZLG-nQwkxbvHjZqXhE5ouzWj1f3qhSda2B4ZEg,83
 diracx/db/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diracx/db/os/__init__.py,sha256=IZr6z6SefrRvuC8sTC4RmB3_wwOyEt1GzpDuwSMH8O4,112
-diracx/db/os/job_parameters.py,sha256=loAc-bo3u-RMAp_H1g8VRt8T-rCCsXp_d9aCvg5OS-A,1225
+diracx/db/os/job_parameters.py,sha256=3w_CeA2z-cY5pWwXkGu-Fod27FobbUXuwVKK-jN037U,1479
 diracx/db/os/utils.py,sha256=V4T-taos64SFNcorfIr7mq5l5y88K6TzyCj1YqWk8VI,11562
 diracx/db/sql/__init__.py,sha256=JYu0b0IVhoXy3lX2m2r2dmAjsRS7IbECBUMEDvX0Te4,391
 diracx/db/sql/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/auth/db.py,sha256=QJtBqMrhOf97UvMG0WpyjsgIRiu19v04FoDzXAyXtT0,8952
-diracx/db/sql/auth/schema.py,sha256=x2PEbmM_bNPdZUN5BMGMrdSmX8zkDeJ3P9XfhLBGBTs,3173
+diracx/db/sql/auth/db.py,sha256=F9s05K-9C6kL2nUZY7P8zD79fiuo2biREMhfI7oCjh4,11979
+diracx/db/sql/auth/schema.py,sha256=9fUV7taDPnoAcoiwRAmQraOmF2Ytoizjs2TFvN7zsVs,3132
 diracx/db/sql/dummy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diracx/db/sql/dummy/db.py,sha256=IW4FzG7ERKbhZvC32KL7Rodu2u-zKAf8BryO4VAdJew,1650
+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=TnEc0fckiuMJAZg2v1_Pbwfn7kDPDam6TXp9ySuiddk,11910
-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
 diracx/db/sql/task_queue/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=QkvpqBuIAgkAOywAssYzdxSzUQVZlSUumK7mPxotXfM,547
-diracx/db/sql/utils/base.py,sha256=HYQuX16mgg9LAMtAEmbTmJFIN0OSMe1Hcb57dtl7LCc,12367
+diracx/db/sql/utils/__init__.py,sha256=k1DI4Idlqv36pXn2BhQysb947Peio9DnYaePslkTpUQ,685
+diracx/db/sql/utils/base.py,sha256=DqW-JYgjqvqkwLFqou5uzg73lZ83C0jHCgkt9qR1NTg,17255
 diracx/db/sql/utils/functions.py,sha256=_E4tc9Gti6LuSh7QEyoqPJSvCuByVqvRenOXCzxsulE,4014
-diracx/db/sql/utils/types.py,sha256=yU-tXsu6hFGPsr9ba1n3ZjGPnHQI_06lbpkTeDCWJtg,1287
-diracx_db-0.0.1a46.dist-info/METADATA,sha256=hu6BtBuYz30tuMO8Z40gyZhSf53QzqTKO45ryw6tTL4,675
-diracx_db-0.0.1a46.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-diracx_db-0.0.1a46.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
-diracx_db-0.0.1a46.dist-info/RECORD,,
+diracx/db/sql/utils/types.py,sha256=KNZWJfpvHTjfIPg6Nn7zY-rS0q3ybnirHcTcLAYSYbE,5118
+diracx_db-0.0.6.dist-info/METADATA,sha256=Lu8x2pR3BfnKGgHYz4w5Z4CTCf7tPi9p9tlldVFJiLo,780
+diracx_db-0.0.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+diracx_db-0.0.6.dist-info/entry_points.txt,sha256=UPqhLvb9gui0kOyWeI_edtefcrHToZmQt1p76vIwujo,317
+diracx_db-0.0.6.dist-info/RECORD,,

{diracx_db-0.0.1a46.dist-info → diracx_db-0.0.6.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any