API: normalize naive datetimes to UTC in validators

Author: xsedla1o

dp3/api/internal/models.py

@@ -1,11 +1,10 @@
-from datetime import datetime
 from typing import Annotated, Any, Literal, Optional, Union
 
 from pydantic import BaseModel, Field, TypeAdapter, create_model, model_validator
 
 from dp3.api.internal.config import MODEL_SPEC
 from dp3.api.internal.helpers import api_to_dp3_datapoint
-from dp3.common.types import T2Datetime
+from dp3.common.types import AwareDatetime, T2Datetime
 
 
 class DataPoint(BaseModel):
@@ -27,7 +26,7 @@ class DataPoint(BaseModel):
     id: Any
     attr: str
     v: Any
-    t1: Optional[datetime] = None
+    t1: Optional[AwareDatetime] = None
     t2: Optional[T2Datetime] = Field(None, validate_default=True)
     c: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0
     src: Optional[str] = None

dp3/api/routers/entity.py

@@ -20,7 +20,7 @@
 from dp3.api.internal.response_models import ErrorResponse, RequestValidationError, SuccessResponse
 from dp3.common.attrspec import AttrType
 from dp3.common.task import DataPointTask, task_context
-from dp3.common.types import UTC
+from dp3.common.types import UTC, AwareDatetime
 from dp3.database.database import DatabaseError
 
 
@@ -43,7 +43,7 @@ async def parse_eid(etype: str, eid: str):
 
 
 def get_eid_master_record_handler(
-    e: EntityId, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: EntityId, date_from: Optional[AwareDatetime] = None, date_to: Optional[AwareDatetime] = None
 ):
     """Handler for getting master record of EID"""
     # TODO: This is probably not the most efficient way. Maybe gather only
@@ -67,7 +67,7 @@ def get_eid_master_record_handler(
 
 
 def get_eid_snapshots_handler(
-    e: EntityId, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: EntityId, date_from: Optional[AwareDatetime] = None, date_to: Optional[AwareDatetime] = None
 ) -> list[dict[str, Any]]:
     """Handler for getting snapshots of EID"""
     snapshots = list(DB.snapshots.get_by_eid(e.type, e.id, t1=date_from, t2=date_to))
@@ -275,7 +275,7 @@ async def count_entity_type_eids(
 
 @router.get("/{etype}/{eid}")
 async def get_eid_data(
-    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[AwareDatetime] = None, date_to: Optional[AwareDatetime] = None
 ) -> EntityEidData:
     """Get data of `etype`'s `eid`.
 
@@ -295,15 +295,15 @@ async def get_eid_data(
 
 @router.get("/{etype}/{eid}/master")
 async def get_eid_master_record(
-    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[AwareDatetime] = None, date_to: Optional[AwareDatetime] = None
 ) -> EntityEidMasterRecord:
     """Get master record of `etype`'s `eid`."""
     return get_eid_master_record_handler(e, date_from, date_to)
 
 
 @router.get("/{etype}/{eid}/snapshots")
 async def get_eid_snapshots(
-    e: ParsedEid, date_from: Optional[datetime] = None, date_to: Optional[datetime] = None
+    e: ParsedEid, date_from: Optional[AwareDatetime] = None, date_to: Optional[AwareDatetime] = None
 ) -> EntityEidSnapshots:
     """Get snapshots of `etype`'s `eid`."""
     return get_eid_snapshots_handler(e, date_from, date_to)
@@ -313,8 +313,8 @@ async def get_eid_snapshots(
 async def get_eid_attr_value(
     e: ParsedEid,
     attr: str,
-    date_from: Optional[datetime] = None,
-    date_to: Optional[datetime] = None,
+    date_from: Optional[AwareDatetime] = None,
+    date_to: Optional[AwareDatetime] = None,
 ) -> EntityEidAttrValueOrHistory:
     """Get attribute value
 
@@ -394,7 +394,7 @@ async def get_distinct_attribute_values(etype: str, attr: str) -> dict[JsonVal,
 
 
 @router.post("/{etype}/{eid}/ttl")
-async def extend_eid_ttls(e: ParsedEid, body: dict[str, datetime]) -> SuccessResponse:
+async def extend_eid_ttls(e: ParsedEid, body: dict[str, AwareDatetime]) -> SuccessResponse:
     """Extend TTLs of the specified entity"""
     # Construct task
     with task_context(MODEL_SPEC):

dp3/common/datapoint.py

@@ -1,11 +1,10 @@
-from datetime import datetime
 from ipaddress import IPv4Address, IPv6Address
 from typing import Annotated, Any, Optional, Union
 
 from pydantic import BaseModel, BeforeValidator, Field, PlainSerializer
 
 from dp3.common.mac_address import MACAddress
-from dp3.common.types import T2Datetime
+from dp3.common.types import AwareDatetime, T2Datetime
 
 
 def to_json_friendly(v):
@@ -60,7 +59,7 @@ class DataPointObservationsBase(DataPointBase):
     Contains single raw data value received on API for observations attribute.
     """
 
-    t1: datetime
+    t1: AwareDatetime
     t2: T2Datetime = Field(None, validate_default=True)
     c: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0
 
@@ -71,7 +70,7 @@ class DataPointTimeseriesBase(DataPointBase):
     Contains single raw data value received on API for observations attribute.
     """
 
-    t1: datetime
+    t1: AwareDatetime
     t2: T2Datetime = Field(None, validate_default=True)
 
 

dp3/common/types.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timedelta, timezone
 from ipaddress import IPv4Address, IPv6Address
 from json import JSONEncoder
-from typing import Annotated, Any, Union
+from typing import Annotated, Any, Optional, Union
 
 from event_count_logger import DummyEventGroup, EventGroup
 from pydantic import AfterValidator, BeforeValidator
@@ -24,6 +24,18 @@ def parse_timedelta_or_passthrough(v):
 ParsedTimedelta = Annotated[timedelta, BeforeValidator(parse_timedelta_or_passthrough)]
 
 
+def ensure_timezone_aware(v: Optional[datetime]):
+    """Ensure datetime is timezone-aware by defaulting to UTC."""
+    if v is None:
+        return v
+    if v.tzinfo is None:
+        return v.replace(tzinfo=UTC)
+    return v
+
+
+AwareDatetime = Annotated[datetime, AfterValidator(ensure_timezone_aware)]
+
+
 def t2_implicity_t1(v, info: FieldValidationInfo):
     """If t2 is not specified, it is set to t1."""
     v = v or info.data.get("t1")
@@ -37,7 +49,11 @@ def t2_after_t1(v, info: FieldValidationInfo):
     return v
 
 
-T2Datetime = Annotated[datetime, BeforeValidator(t2_implicity_t1), AfterValidator(t2_after_t1)]
+T2Datetime = Annotated[
+    AwareDatetime,
+    BeforeValidator(t2_implicity_t1),
+    AfterValidator(t2_after_t1),
+]
 
 EventGroupType = Union[EventGroup, DummyEventGroup]
 

tests/test_common/test_types.py

@@ -0,0 +1,34 @@
+import unittest
+from datetime import datetime, timedelta, timezone
+
+from pydantic import BaseModel, Field
+
+from dp3.common.types import AwareDatetime, T2Datetime
+
+
+class _AwareModel(BaseModel):
+    dt: AwareDatetime
+
+
+class _T2Model(BaseModel):
+    t1: AwareDatetime
+    t2: T2Datetime = Field(None, validate_default=True)
+
+
+class TestAwareDatetime(unittest.TestCase):
+    def test_naive_datetime_defaults_to_utc(self):
+        model = _AwareModel(dt="2024-01-01T10:00:00")
+        self.assertEqual(model.dt.tzinfo, timezone.utc)
+
+    def test_existing_timezone_is_preserved(self):
+        cest_timezone = timezone(timedelta(hours=2), "CEST")
+        aware = datetime(2024, 1, 1, 10, 0, tzinfo=cest_timezone)
+        model = _AwareModel(dt=aware)
+        self.assertEqual(model.dt.tzinfo, cest_timezone)
+
+    def test_t2_datetime_inherits_timezone_when_missing(self):
+        model = _T2Model(t1="2024-01-01T00:00:00")
+        self.assertIsNotNone(model.t2)
+        self.assertEqual(model.t1.tzinfo, timezone.utc)
+        self.assertEqual(model.t2.tzinfo, timezone.utc)
+        self.assertEqual(model.t2, model.t1)