|
5 | 5 |
|
6 | 6 | from __future__ import annotations
|
7 | 7 |
|
8 |
| -from datetime import datetime |
| 8 | +from datetime import datetime, timezone |
9 | 9 | from enum import StrEnum
|
10 | 10 | from typing import Literal
|
11 | 11 |
|
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)] |
14 | 38 |
|
15 | 39 |
|
16 | 40 | class ScalarSearchOperator(StrEnum):
|
@@ -56,7 +80,7 @@ class InsertedJob(TypedDict):
|
56 | 80 | JobID: int
|
57 | 81 | Status: str
|
58 | 82 | MinorStatus: str
|
59 |
| - TimeStamp: datetime |
| 83 | + TimeStamp: DiracUTCDatetime |
60 | 84 |
|
61 | 85 |
|
62 | 86 | class JobSummaryParams(BaseModel):
|
@@ -101,7 +125,7 @@ class JobLoggingRecord(BaseModel):
|
101 | 125 | status: JobStatus | Literal["idem"]
|
102 | 126 | minor_status: str
|
103 | 127 | application_status: str
|
104 |
| - date: datetime |
| 128 | + date: DiracUTCDatetime |
105 | 129 | source: str
|
106 | 130 |
|
107 | 131 |
|
@@ -130,10 +154,10 @@ class SetJobStatusReturnSuccess(BaseModel):
|
130 | 154 | Status: JobStatus | None = None
|
131 | 155 | MinorStatus: str | None = None
|
132 | 156 | 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 |
137 | 161 |
|
138 | 162 | success: dict[int, SetJobStatusReturnSuccess]
|
139 | 163 | failed: dict[int, dict[str, str]]
|
|
0 commit comments