Timestamptz autoconversion

piccolo/apps/migrations/auto/serialisation.py

@@ -12,11 +12,17 @@
 from dataclasses import dataclass, field
 from enum import Enum
 
-from piccolo.columns import Column
+from piccolo.columns import Column, Timestamptz
 from piccolo.columns.defaults.base import Default
+from piccolo.columns.defaults.timestamptz import (
+    TimestamptzCustom,
+    TimestamptzNow,
+    TimestamptzOffset,
+)
 from piccolo.columns.reference import LazyTableReference
 from piccolo.table import Table
 from piccolo.utils.repr import repr_class_instance
+from piccolo.utils.zoneinfo import ZoneInfo
 
 from .serialisation_legacy import deserialise_legacy_params
 
@@ -546,6 +552,30 @@ def serialise_params(params: t.Dict[str, t.Any]) -> SerialisedParams:
                     expect_conflict_with_global_name=UniqueGlobalNames.DEFAULT,
                 )
             )
+            # ZoneInfo for Timestamptz* instances
+            in_group = (
+                Timestamptz,
+                TimestamptzNow,
+                TimestamptzCustom,
+                TimestamptzOffset,
+            )
+            if isinstance(value, in_group):
+                extra_imports.append(
+                    Import(
+                        module=ZoneInfo.__module__,
+                        target=None,
+                    )
+                )
+            continue
+
+        # ZoneInfo instances
+        if isinstance(value, ZoneInfo):
+            extra_imports.append(
+                Import(
+                    module=value.__class__.__module__,
+                    target=None,
+                )
+            )
             continue
 
         # Dates and times
@@ -633,6 +663,7 @@ def serialise_params(params: t.Dict[str, t.Any]) -> SerialisedParams:
                     extra_imports.append(
                         Import(module=module_name, target=type_.__name__)
                     )
+            continue
 
         # Functions
         if inspect.isfunction(value):

piccolo/columns/base.py

@@ -763,6 +763,9 @@ def as_alias(self, name: str) -> Column:
         column._alias = name
         return column
 
+    def _get_alias(self) -> str:
+        return self._alias or self._meta.get_default_alias()
+
     def join_on(self, column: Column) -> ForeignKey:
         """
         Joins are typically performed via foreign key columns. For example,
@@ -945,8 +948,8 @@ def ddl(self) -> str:
 
         return query
 
-    def copy(self) -> Column:
-        column: Column = copy.copy(self)
+    def copy(self: Self) -> Self:
+        column = copy.copy(self)
         column._meta = self._meta.copy()
         return column
 
@@ -971,3 +974,6 @@ def __repr__(self):
             f"{table_class_name}.{self._meta.name} - "
             f"{self.__class__.__name__}"
         )
+
+
+Self = t.TypeVar("Self", bound=Column)

piccolo/columns/column_types.py

@@ -63,6 +63,7 @@ class Band(Table):
 from piccolo.querystring import QueryString, Unquoted
 from piccolo.utils.encoding import dump_json
 from piccolo.utils.warnings import colored_warning
+from piccolo.utils.zoneinfo import ZoneInfo
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.columns.base import ColumnMeta
@@ -955,30 +956,33 @@ def __set__(self, obj, value: t.Union[datetime, None]):
 class Timestamptz(Column):
     """
     Used for storing timezone aware datetimes. Uses the ``datetime`` type for
-    values. The values are converted to UTC in the database, and are also
-    returned as UTC.
+    values. The values are converted to UTC when saved into the database and
+    are converted back into the timezone of the column on select queries.
 
     **Example**
 
     .. code-block:: python
 
         import datetime
+        from zoneinfo import ZoneInfo
 
-        class Concert(Table):
-            starts = Timestamptz()
+        class TallinnConcerts(Table):
+            event_start = Timestamptz(at_time_zone=ZoneInfo("Europe/Tallinn"))
 
         # Create
-        >>> await Concert(
-        ...    starts=datetime.datetime(
-        ...        year=2050, month=1, day=1, tzinfo=datetime.timezone.tz
+        >>> await TallinnConcerts(
+        ...    event_start=datetime.datetime(
+        ...        year=2050, month=1, day=1, hour=20
         ...    )
         ... ).save()
 
         # Query
-        >>> await Concert.select(Concert.starts)
+        >>> await TallinnConcerts.select(TallinnConcerts.event_start)
         {
-            'starts': datetime.datetime(
-                2050, 1, 1, 0, 0, tzinfo=datetime.timezone.utc
+            'event_start': datetime.datetime(
+                2050, 1, 1, 20, 0, tzinfo=zoneinfo.ZoneInfo(
+                    key='Europe/Tallinn'
+                )
             )
         }
 
@@ -993,22 +997,59 @@ class Concert(Table):
     timedelta_delegate = TimedeltaDelegate()
 
     def __init__(
-        self, default: TimestamptzArg = TimestamptzNow(), **kwargs
+        self,
+        default: TimestamptzArg = TimestamptzNow(),
+        at_time_zone: ZoneInfo = ZoneInfo("UTC"),
+        **kwargs,
     ) -> None:
         self._validate_default(
             default, TimestamptzArg.__args__  # type: ignore
         )
 
         if isinstance(default, datetime):
-            default = TimestamptzCustom.from_datetime(default)
+            default = TimestamptzCustom.from_datetime(default, at_time_zone)
 
         if default == datetime.now:
-            default = TimestamptzNow()
+            default = TimestamptzNow(tz=at_time_zone)
 
+        self._at_time_zone = at_time_zone
         self.default = default
-        kwargs.update({"default": default})
+        kwargs.update({"default": default, "at_time_zone": at_time_zone})
         super().__init__(**kwargs)
 
+    ###########################################################################
+
+    def at_time_zone(self, time_zone: t.Union[ZoneInfo, str]) -> Timestamptz:
+        """
+        By default, the database returns the value in UTC. This lets us get
+        the value converted to the specified timezone.
+        """
+        time_zone = (
+            ZoneInfo(time_zone) if isinstance(time_zone, str) else time_zone
+        )
+        instance = self.copy()
+        instance._at_time_zone = time_zone
+        return instance
+
+    ###########################################################################
+
+    def get_select_string(
+        self, engine_type: str, with_alias: bool = True
+    ) -> str:
+        select_string = self._meta.get_full_name(with_alias=False)
+
+        if self._at_time_zone != ZoneInfo("UTC"):
+            # SQLite doesn't support `AT TIME ZONE`, so we have to do it in
+            # Python instead (see ``Select.response_handler``).
+            if self._meta.engine_type in ("postgres", "cockroach"):
+                select_string += f" AT TIME ZONE '{self._at_time_zone.key}'"
+
+        if with_alias:
+            alias = self._get_alias()
+            select_string += f' AS "{alias}"'
+
+        return select_string
+
     ###########################################################################
     # For update queries
 
@@ -2305,7 +2346,7 @@ def arrow(self, key: str) -> JSONB:
         Allows part of the JSON structure to be returned - for example,
         for {"a": 1}, and a key value of "a", then 1 will be returned.
         """
-        instance = t.cast(JSONB, self.copy())
+        instance = self.copy()
         instance.json_operator = f"-> '{key}'"
         return instance
 
@@ -2318,7 +2359,7 @@ def get_select_string(
             select_string += f" {self.json_operator}"
 
         if with_alias:
-            alias = self._alias or self._meta.get_default_alias()
+            alias = self._get_alias()
             select_string += f' AS "{alias}"'
 
         return select_string
@@ -2623,7 +2664,7 @@ def get_select_string(self, engine_type: str, with_alias=True) -> str:
             select_string += f"[{self.index}]"
 
         if with_alias:
-            alias = self._alias or self._meta.get_default_alias()
+            alias = self._get_alias()
             select_string += f' AS "{alias}"'
 
         return select_string

piccolo/columns/defaults/timestamptz.py

@@ -1,13 +1,28 @@
 from __future__ import annotations
 
-import datetime
+import datetime as pydatetime
 import typing as t
 from enum import Enum
 
+from piccolo.utils.zoneinfo import ZoneInfo
+
 from .timestamp import TimestampCustom, TimestampNow, TimestampOffset
 
 
 class TimestamptzOffset(TimestampOffset):
+    def __init__(
+        self,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: int = 0,
+        tz: ZoneInfo = ZoneInfo("UTC"),
+    ):
+        self.tz = tz
+        super().__init__(
+            days=days, hours=hours, minutes=minutes, seconds=seconds
+        )
+
     @property
     def cockroach(self):
         interval_string = self.get_postgres_interval_string(
@@ -16,9 +31,7 @@ def cockroach(self):
         return f"CURRENT_TIMESTAMP + INTERVAL '{interval_string}'"
 
     def python(self):
-        return datetime.datetime.now(
-            tz=datetime.timezone.utc
-        ) + datetime.timedelta(
+        return pydatetime.datetime.now(tz=self.tz) + pydatetime.timedelta(
             days=self.days,
             hours=self.hours,
             minutes=self.minutes,
@@ -27,35 +40,60 @@ def python(self):
 
 
 class TimestamptzNow(TimestampNow):
+    def __init__(self, tz: ZoneInfo = ZoneInfo("UTC")):
+        self.tz = tz
+
     @property
     def cockroach(self):
         return "current_timestamp"
 
     def python(self):
-        return datetime.datetime.now(tz=datetime.timezone.utc)
+        return pydatetime.datetime.now(tz=self.tz)
 
 
 class TimestamptzCustom(TimestampCustom):
+    def __init__(
+        self,
+        year: int = 2000,
+        month: int = 1,
+        day: int = 1,
+        hour: int = 0,
+        second: int = 0,
+        microsecond: int = 0,
+        tz: ZoneInfo = ZoneInfo("UTC"),
+    ):
+        self.tz = tz
+        super().__init__(
+            year=year,
+            month=month,
+            day=day,
+            hour=hour,
+            second=second,
+            microsecond=microsecond,
+        )
+
     @property
     def cockroach(self):
         return "'{}'".format(self.datetime.isoformat().replace("T", " "))
 
     @property
     def datetime(self):
-        return datetime.datetime(
+        return pydatetime.datetime(
             year=self.year,
             month=self.month,
             day=self.day,
             hour=self.hour,
             second=self.second,
             microsecond=self.microsecond,
-            tzinfo=datetime.timezone.utc,
+            tzinfo=self.tz,
         )
 
     @classmethod
-    def from_datetime(cls, instance: datetime.datetime):  # type: ignore
+    def from_datetime(
+        cls, instance: pydatetime.datetime, tz: ZoneInfo = ZoneInfo("UTC")
+    ):  # type: ignore
         if instance.tzinfo is not None:
-            instance = instance.astimezone(datetime.timezone.utc)
+            instance = instance.astimezone(tz)
         return cls(
             year=instance.year,
             month=instance.month,
@@ -72,7 +110,7 @@ def from_datetime(cls, instance: datetime.datetime):  # type: ignore
     TimestamptzOffset,
     Enum,
     None,
-    datetime.datetime,
+    pydatetime.datetime,
 ]