Skip to content

Commit 74e9829

Browse files
committed
feat: validate datetime objects to be aware and UTC
simplify validator
1 parent 76ebeff commit 74e9829

File tree

1 file changed

+17
-9
lines changed

1 file changed

+17
-9
lines changed

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

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,20 @@
55

66
from __future__ import annotations
77

8-
from datetime import datetime
8+
from datetime import datetime, UTC
99
from enum import StrEnum
1010
from typing import Literal
1111

12-
from pydantic import BaseModel, Field
13-
from typing_extensions import TypedDict
12+
from pydantic import AwareDatetime, BaseModel, AfterValidator, Field
13+
from typing_extensions import Annotated, TypedDict
14+
15+
16+
def good_utc_dt(v):
17+
"""A validator that ensures that the aware datetime is in UTC timezone."""
18+
return v.astimezone(UTC)
19+
20+
21+
DiracUTCDatetime = Annotated[AwareDatetime, AfterValidator(good_utc_dt)]
1422

1523

1624
class ScalarSearchOperator(StrEnum):
@@ -56,7 +64,7 @@ class InsertedJob(TypedDict):
5664
JobID: int
5765
Status: str
5866
MinorStatus: str
59-
TimeStamp: datetime
67+
TimeStamp: DiracUTCDatetime
6068

6169

6270
class JobSummaryParams(BaseModel):
@@ -101,7 +109,7 @@ class JobLoggingRecord(BaseModel):
101109
status: JobStatus | Literal["idem"]
102110
minor_status: str
103111
application_status: str
104-
date: datetime
112+
date: DiracUTCDatetime
105113
source: str
106114

107115

@@ -130,10 +138,10 @@ class SetJobStatusReturnSuccess(BaseModel):
130138
Status: JobStatus | None = None
131139
MinorStatus: str | None = None
132140
ApplicationStatus: str | None = None
133-
HeartBeatTime: datetime | None = None
134-
StartExecTime: datetime | None = None
135-
EndExecTime: datetime | None = None
136-
LastUpdateTime: datetime | None = None
141+
HeartBeatTime: DiracUTCDatetime | None = None
142+
StartExecTime: DiracUTCDatetime | None = None
143+
EndExecTime: DiracUTCDatetime | None = None
144+
LastUpdateTime: DiracUTCDatetime | None = None
137145

138146
success: dict[int, SetJobStatusReturnSuccess]
139147
failed: dict[int, dict[str, str]]

0 commit comments

Comments
 (0)