🐛 Fix timedelta arithmetic with UTC datetime fields

sqlmodel/sql/sqltypes.py

@@ -1,8 +1,9 @@
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from typing import Any, cast
 
 from sqlalchemy import types
 from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.sql.operators import OperatorType
 
 
 class UTCDateTime(types.TypeDecorator[datetime]):
@@ -17,6 +18,13 @@ def __init__(self) -> None:
     def __repr__(self) -> str:
         return "UTCDateTime()"
 
+    def coerce_compared_value(
+        self, op: OperatorType | None, value: Any
+    ) -> types.TypeEngine[Any]:
+        if isinstance(value, timedelta):
+            return types.Interval()
+        return self
+
     def process_bind_param(
         self, value: datetime | None, dialect: Dialect
     ) -> datetime | None:

tests/test_datetime.py

@@ -1,5 +1,7 @@
+from collections.abc import Callable
 from datetime import date, datetime, time, timedelta, timezone, tzinfo
 from io import StringIO
+from operator import add, sub
 from typing import Annotated, Any
 
 import pytest
@@ -264,6 +266,34 @@ def test_postgresql_datetime_processing() -> None:
     assert result(None) is None
 
 
+@pytest.mark.parametrize(
+    "operation, reverse",
+    [(add, False), (sub, False), (add, True)],
+)
+def test_postgresql_datetime_interval_arithmetic(
+    operation: Callable[[Any, Any], Any], reverse: bool
+) -> None:
+    class Event(SQLModel, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        occurred_at: datetime
+
+    occurred_at = SQLModel.metadata.tables["event"].c.occurred_at
+    duration = timedelta(days=30)
+    expression = (
+        operation(duration, occurred_at)
+        if reverse
+        else operation(occurred_at, duration)
+    )
+    # Adding or subtracting a duration still produces a UTC datetime.
+    assert isinstance(expression.type, UTCDateTime)
+    # Exercise the actual parameter processor used for database execution.
+    parameter = expression.left if reverse else expression.right
+    dialect = postgresql.dialect()
+    bind = parameter.type.dialect_impl(dialect).bind_processor(dialect)
+    bound = bind(parameter.value) if bind is not None else parameter.value
+    assert bound == duration
+
+
 def test_existing_sqlite_utc_data() -> None:
     class Event(SQLModel, table=True):
         id: int | None = Field(default=None, primary_key=True)