diff --git a/docs/src/piccolo/functions/datetime.rst b/docs/src/piccolo/functions/datetime.rst index ee33e8c4a..0c7323b2f 100644 --- a/docs/src/piccolo/functions/datetime.rst +++ b/docs/src/piccolo/functions/datetime.rst @@ -6,12 +6,16 @@ Datetime functions Postgres / Cockroach -------------------- +AtTimeZone +~~~~~~~~~~ + +.. autoclass:: AtTimeZone + Extract ~~~~~~~ .. autoclass:: Extract - SQLite ------ diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py index 1effd6922..5b94b2f01 100644 --- a/piccolo/apps/playground/commands/run.py +++ b/piccolo/apps/playground/commands/run.py @@ -25,6 +25,7 @@ Numeric, Serial, Text, + Time, Timestamp, Timestamptz, Varchar, @@ -143,6 +144,8 @@ class RecordingStudio(Table): id: Serial name = Varchar(length=100) facilities = JSON(null=True) + opens_at = Time() + closes_at = Time() @classmethod def get_readable(cls) -> Readable: @@ -279,6 +282,8 @@ def populate(): {"name": "Bob Williams"}, ], }, + RecordingStudio.opens_at: datetime.time(9, 0, 0), + RecordingStudio.closes_at: datetime.time(17, 0, 0), } ) recording_studio_1.save().run_sync() @@ -294,6 +299,8 @@ def populate(): {"name": "Frank Smith"}, ], }, + RecordingStudio.opens_at: datetime.time(10, 0, 0), + RecordingStudio.closes_at: datetime.time(20, 0, 0), }, ) recording_studio_2.save().run_sync() diff --git a/piccolo/query/functions/datetime.py b/piccolo/query/functions/datetime.py index 146fc5c42..edc935591 100644 --- a/piccolo/query/functions/datetime.py +++ b/piccolo/query/functions/datetime.py @@ -1,4 +1,8 @@ +from __future__ import annotations + +import datetime from typing import Literal, Optional, Union, get_args +from zoneinfo import ZoneInfo from piccolo.columns.base import Column from piccolo.columns.column_types import ( @@ -248,6 +252,81 @@ def Second( ) +class AtTimeZone(QueryString): + def __init__( + self, + identifier: Union[Time, Timestamp, Timestamptz, QueryString], + timezone: ZoneInfo | str | datetime.timedelta, + alias: Optional[str] = None, + ): + """ + .. note:: This is for Postgres / Cockroach only. + + Convert the column to the given timezone. See the + `Postgres docs `_ + for more information. + + For example:: + + class Signing(Table): + starts = Timestamptz() + + >>> await Signing.select( + ... AtTimeZone(Signing.starts, 'EST', alias='starts_est'), + ... Signing.starts, + ... ) + [{ + 'starts_est': datetime.datetime( + 2026, 12, 20, 5, 0 + ), + 'starts': datetime.datetime( + 2026, 12, 20, 10, 0, tzinfo=datetime.timezone.utc + ) + }] + + :param timezone: + Valid arguments are ``'EST'``, ``ZoneInfo('EST')`` and + ``timedelta(hours=5)``. + + """ # noqa: E501 + # Preserve the original alias from the column. + + from piccolo.columns import Column + + if isinstance(identifier, Column): + alias = ( + alias + or identifier._alias + or identifier._meta.get_default_alias() + ) + elif isinstance(identifier, QueryString): + alias = alias or identifier._alias + + ####################################################################### + + if isinstance(timezone, str): + # Validate it's a correct timezone + timezone = ZoneInfo(timezone) + + if isinstance(timezone, datetime.timedelta): + total_seconds = timezone.total_seconds() + prefix = "+" if total_seconds >= 0 else "-" + timezone = ( + f"{prefix}{datetime.timedelta(seconds=abs(total_seconds))}" + ) + + if isinstance(timezone, ZoneInfo): + # Validate it's a correct timezone + timezone = timezone.key + + super().__init__( + "{} AT TIME ZONE {}", + identifier, + timezone, + alias=alias, + ) + + __all__ = ( "Extract", "Strftime", @@ -257,4 +336,5 @@ def Second( "Hour", "Minute", "Second", + "AtTimeZone", ) diff --git a/tests/query/functions/test_datetime.py b/tests/query/functions/test_datetime.py index 382f688ab..5780fb690 100644 --- a/tests/query/functions/test_datetime.py +++ b/tests/query/functions/test_datetime.py @@ -1,7 +1,9 @@ import datetime +import zoneinfo -from piccolo.columns import Timestamp +from piccolo.columns import Timestamp, Timestamptz from piccolo.query.functions.datetime import ( + AtTimeZone, Day, Extract, Hour, @@ -12,7 +14,7 @@ Year, ) from piccolo.table import Table -from piccolo.testing.test_case import TableTest +from piccolo.testing.test_case import AsyncTableTest, TableTest from tests.base import engines_only, sqlite_only @@ -111,3 +113,33 @@ def test_second(self): ).run_sync(), [{"starts_second": self.concert.starts.second}], ) + + +class Signing(Table): + starts = Timestamptz() + + +@engines_only("cockroach", "postgres") +class TestAtTimeZone(AsyncTableTest): + + tables = [Signing] + + async def test_at_time_zone(self): + await Signing( + starts=datetime.datetime( + 2026, 5, 11, 7, 0, 0, tzinfo=datetime.timezone.utc + ) + ).save() + + for timezone in ( + "EST", + zoneinfo.ZoneInfo("EST"), + datetime.timedelta(hours=5), + ): + self.assertEqual( + await Signing.select(AtTimeZone(Signing.starts, timezone)), + [{"starts": datetime.datetime(2026, 5, 11, 2, 0, 0)}], + ) + + with self.assertRaises(zoneinfo.ZoneInfoNotFoundError): + AtTimeZone(Signing.starts, "ABC")