feat: validate datetime objects to be aware and UTC

simplify validator

Author: ryuwd

diracx-core/src/diracx/core/models.py

@@ -5,12 +5,20 @@
 
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import datetime, UTC
 from enum import StrEnum
 from typing import Literal
 
-from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
+from pydantic import AwareDatetime, BaseModel, AfterValidator, 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]]