diff --git a/diracx-core/src/diracx/core/models.py b/diracx-core/src/diracx/core/models.py index 415e36295..73f3b4544 100644 --- a/diracx-core/src/diracx/core/models.py +++ b/diracx-core/src/diracx/core/models.py @@ -5,12 +5,20 @@ from __future__ import annotations -from datetime import datetime +from datetime import UTC, datetime from enum import StrEnum from typing import Literal -from pydantic import BaseModel, Field -from typing_extensions import TypedDict +from pydantic import AfterValidator, AwareDatetime, BaseModel, Field +from typing_extensions import Annotated, TypedDict + + +def good_utc_dt(v): + """A validator that ensures that the aware datetime is in UTC timezone.""" + return v.astimezone(UTC) + + +DiracUTCDatetime = Annotated[AwareDatetime, AfterValidator(good_utc_dt)] class ScalarSearchOperator(StrEnum): @@ -56,7 +64,7 @@ class InsertedJob(TypedDict): JobID: int Status: str MinorStatus: str - TimeStamp: datetime + TimeStamp: DiracUTCDatetime class JobSummaryParams(BaseModel): @@ -101,7 +109,7 @@ class JobLoggingRecord(BaseModel): status: JobStatus | Literal["idem"] minor_status: str application_status: str - date: datetime + date: DiracUTCDatetime source: str @@ -130,10 +138,10 @@ class SetJobStatusReturnSuccess(BaseModel): Status: JobStatus | None = None MinorStatus: str | None = None ApplicationStatus: str | None = None - HeartBeatTime: datetime | None = None - StartExecTime: datetime | None = None - EndExecTime: datetime | None = None - LastUpdateTime: datetime | None = None + HeartBeatTime: DiracUTCDatetime | None = None + StartExecTime: DiracUTCDatetime | None = None + EndExecTime: DiracUTCDatetime | None = None + LastUpdateTime: DiracUTCDatetime | None = None success: dict[int, SetJobStatusReturnSuccess] failed: dict[int, dict[str, str]]