Skip to content

Commit 92d7943

Browse files
committed
Bring epoch calculation and date validation in line with DIRAC
Stricter models. precommit Don't use TimeUtilities revert JobLoggingDB changes
1 parent 76ebeff commit 92d7943

File tree

1 file changed

+33
-9
lines changed

1 file changed

+33
-9
lines changed

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

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

66
from __future__ import annotations
77

8-
from datetime import datetime
8+
from datetime import datetime, timezone
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, BeforeValidator, Field
13+
from typing_extensions import Annotated, TypedDict
14+
15+
16+
def good_utc_dt(v):
17+
# We need to specify that timezone is UTC because otherwise timestamp
18+
# assumes local time while we mean UTC.
19+
if not v:
20+
# Make the UTC datetime
21+
v = datetime.now(tz=timezone.utc)
22+
23+
if isinstance(v, str):
24+
# The date is provided as a string.
25+
v = datetime.fromisoformat(v)
26+
27+
if isinstance(v, datetime):
28+
if not v.tzinfo:
29+
raise ValueError(f"datetime object provided but tzinfo was None. got {v=}")
30+
return v.astimezone(timezone.utc)
31+
32+
raise ValueError(
33+
f"Incorrect date for the logging record: got {v=}. Should be a UTC datetime object."
34+
)
35+
36+
37+
DiracUTCDatetime = Annotated[AwareDatetime, BeforeValidator(good_utc_dt)]
1438

1539

1640
class ScalarSearchOperator(StrEnum):
@@ -56,7 +80,7 @@ class InsertedJob(TypedDict):
5680
JobID: int
5781
Status: str
5882
MinorStatus: str
59-
TimeStamp: datetime
83+
TimeStamp: DiracUTCDatetime
6084

6185

6286
class JobSummaryParams(BaseModel):
@@ -101,7 +125,7 @@ class JobLoggingRecord(BaseModel):
101125
status: JobStatus | Literal["idem"]
102126
minor_status: str
103127
application_status: str
104-
date: datetime
128+
date: DiracUTCDatetime
105129
source: str
106130

107131

@@ -130,10 +154,10 @@ class SetJobStatusReturnSuccess(BaseModel):
130154
Status: JobStatus | None = None
131155
MinorStatus: str | None = None
132156
ApplicationStatus: str | None = None
133-
HeartBeatTime: datetime | None = None
134-
StartExecTime: datetime | None = None
135-
EndExecTime: datetime | None = None
136-
LastUpdateTime: datetime | None = None
157+
HeartBeatTime: DiracUTCDatetime | None = None
158+
StartExecTime: DiracUTCDatetime | None = None
159+
EndExecTime: DiracUTCDatetime | None = None
160+
LastUpdateTime: DiracUTCDatetime | None = None
137161

138162
success: dict[int, SetJobStatusReturnSuccess]
139163
failed: dict[int, dict[str, str]]

0 commit comments

Comments
 (0)