From acc2e3c527e139b50400b65924f96c1c9389e84b Mon Sep 17 00:00:00 2001
From: Ryunosuke O'Neil <r.oneil@cern.ch>
Date: Sun, 13 Apr 2025 16:30:41 +0200
Subject: [PATCH] feat (JobDB): validate datetime objects to be aware and UTC

simplify validator

[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 diracx-core/src/diracx/core/models.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

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]]