Merge pull request piccolo-orm#52 from piccolo-orm/timestamptz_column

Timestamptz column

Author: dantownsend

docs/src/piccolo/schema/column_types.rst

@@ -146,6 +146,12 @@ Timestamp
 
 .. autoclass:: Timestamp
 
+===========
+Timestamptz
+===========
+
+.. autoclass:: Timestamptz
+
 -------------------------------------------------------------------------------
 
 ****

piccolo/columns/__init__.py

@@ -16,6 +16,7 @@
     Serial,
     Text,
     Timestamp,
+    Timestamptz,
     UUID,
     Varchar,
 )

piccolo/columns/column_types.py

@@ -20,6 +20,11 @@
     TimestampCustom,
     TimestampNow,
 )
+from piccolo.columns.defaults.timestamptz import (
+    TimestamptzArg,
+    TimestamptzCustom,
+    TimestamptzNow,
+)
 from piccolo.columns.defaults.uuid import UUID4, UUIDArg
 from piccolo.columns.operators.string import ConcatPostgres, ConcatSQLite
 from piccolo.columns.reference import LazyTableReference
@@ -545,6 +550,11 @@ def __init__(
         self._validate_default(default, TimestampArg.__args__)  # type: ignore
 
         if isinstance(default, datetime):
+            if default.tzinfo is not None:
+                raise ValueError(
+                    "Timestamp only stores timezone naive datetime objects - "
+                    "use Timestamptz instead."
+                )
             default = TimestampCustom.from_datetime(default)
 
         if default == datetime.now:
@@ -555,6 +565,58 @@ def __init__(
         super().__init__(**kwargs)
 
 
+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.
+
+    **Example**
+
+    .. code-block:: python
+
+        import datetime
+
+        class Concert(Table):
+            starts = Timestamptz()
+
+        # Create
+        >>> Concert(
+        >>>    starts=datetime.datetime(
+        >>>        year=2050, month=1, day=1, tzinfo=datetime.timezone.tz
+        >>>    )
+        >>> ).save().run_sync()
+
+        # Query
+        >>> Concert.select(Concert.starts).run_sync()
+        {
+            'starts': datetime.datetime(
+                2050, 1, 1, 0, 0, tzinfo=datetime.timezone.utc
+            )
+        }
+
+    """
+
+    value_type = datetime
+
+    def __init__(
+        self, default: TimestamptzArg = TimestamptzNow(), **kwargs
+    ) -> None:
+        self._validate_default(
+            default, TimestamptzArg.__args__  # type: ignore
+        )
+
+        if isinstance(default, datetime):
+            default = TimestamptzCustom.from_datetime(default)
+
+        if default == datetime.now:
+            default = TimestamptzNow()
+
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+
 class Date(Column):
     """
     Used for storing dates. Uses the ``date`` type for values.

piccolo/columns/defaults/timestamp.py

@@ -53,12 +53,12 @@ def python(self):
 class TimestampCustom(Default):
     def __init__(
         self,
-        year: int,
-        month: int,
-        day: int,
-        hour: int,
-        second: int,
-        microsecond: int,
+        year: int = 2000,
+        month: int = 1,
+        day: int = 1,
+        hour: int = 0,
+        second: int = 0,
+        microsecond: int = 0,
     ):
         self.year = year
         self.month = month

piccolo/columns/defaults/timestamptz.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+import datetime
+import typing as t
+
+from .timestamp import TimestampOffset, TimestampNow, TimestampCustom
+
+
+class TimestamptzOffset(TimestampOffset):
+    def python(self):
+        return datetime.datetime.now(
+            tz=datetime.timezone.utc
+        ) + datetime.timedelta(
+            days=self.days,
+            hours=self.hours,
+            minutes=self.minutes,
+            seconds=self.seconds,
+        )
+
+
+class TimestamptzNow(TimestampNow):
+    def python(self):
+        return datetime.datetime.now(tz=datetime.timezone.utc)
+
+
+class TimestamptzCustom(TimestampCustom):
+    @property
+    def datetime(self):
+        return datetime.datetime(
+            year=self.year,
+            month=self.month,
+            day=self.day,
+            hour=self.hour,
+            second=self.second,
+            microsecond=self.microsecond,
+            tzinfo=datetime.timezone.utc,
+        )
+
+    @classmethod
+    def from_datetime(cls, instance: datetime.datetime):  # type: ignore
+        if instance.tzinfo is not None:
+            instance = instance.astimezone(datetime.timezone.utc)
+        return cls(
+            year=instance.year,
+            month=instance.month,
+            day=instance.month,
+            hour=instance.hour,
+            second=instance.second,
+            microsecond=instance.microsecond,
+        )
+
+
+TimestamptzArg = t.Union[
+    TimestamptzCustom,
+    TimestamptzNow,
+    TimestamptzOffset,
+    None,
+    datetime.datetime,
+]
+
+
+__all__ = [
+    "TimestamptzArg",
+    "TimestamptzCustom",
+    "TimestamptzNow",
+    "TimestamptzOffset",
+]

piccolo/engine/sqlite.py

@@ -22,21 +22,24 @@
 # consistent with the Postgres engine.
 
 
+# In
+
+
 def convert_numeric_in(value):
     """
     Convert any Decimal values into floats.
     """
     return float(value)
 
 
-def convert_uuid_in(value):
+def convert_uuid_in(value) -> str:
     """
     Converts the UUID value being passed into sqlite.
     """
     return str(value)
 
 
-def convert_time_in(value):
+def convert_time_in(value) -> str:
     """
     Converts the time value being passed into sqlite.
     """
@@ -50,13 +53,27 @@ def convert_date_in(value):
     return value.isoformat()
 
 
+def convert_datetime_in(value: datetime.datetime) -> str:
+    """
+    Converts the datetime into a string. If it's timezone aware, we want to
+    convert it to UTC first. This is to replicate Postgres, which stores
+    timezone aware datetimes in UTC.
+    """
+    if value.tzinfo is not None:
+        value = value.astimezone(datetime.timezone.utc)
+    return str(value)
+
+
 def convert_timedelta_in(value):
     """
     Converts the timedelta value being passed into sqlite.
     """
     return value.total_seconds()
 
 
+# Out
+
+
 def convert_numeric_out(value: bytes) -> Decimal:
     """
     Convert float values into Decimals.
@@ -104,18 +121,28 @@ def convert_boolean_out(value: bytes) -> bool:
     return _value == "1"
 
 
+def convert_timestamptz_out(value: bytes) -> datetime.datetime:
+    """
+    If the value is from a timstamptz column, convert it to a datetime value,
+    with a timezone of UTC.
+    """
+    return datetime.datetime.fromisoformat(value.decode("utf8"))
+
+
 sqlite3.register_converter("Numeric", convert_numeric_out)
 sqlite3.register_converter("Integer", convert_int_out)
 sqlite3.register_converter("UUID", convert_uuid_out)
 sqlite3.register_converter("Date", convert_date_out)
 sqlite3.register_converter("Time", convert_time_out)
 sqlite3.register_converter("Seconds", convert_seconds_out)
 sqlite3.register_converter("Boolean", convert_boolean_out)
+sqlite3.register_converter("Timestamptz", convert_timestamptz_out)
 
 sqlite3.register_adapter(Decimal, convert_numeric_in)
 sqlite3.register_adapter(uuid.UUID, convert_uuid_in)
 sqlite3.register_adapter(datetime.time, convert_time_in)
 sqlite3.register_adapter(datetime.date, convert_date_in)
+sqlite3.register_adapter(datetime.datetime, convert_datetime_in)
 sqlite3.register_adapter(datetime.timedelta, convert_timedelta_in)
 
 ###############################################################################

tests/columns/test_timestamp.py

@@ -1,16 +1,21 @@
 import datetime
 from unittest import TestCase
 
-from piccolo.table import Table
 from piccolo.columns.column_types import Timestamp
 from piccolo.columns.defaults.timestamp import TimestampNow
+from piccolo.table import Table
 
 
 class MyTable(Table):
     created_on = Timestamp()
 
 
 class MyTableDefault(Table):
+    """
+    A table containing all of the possible `default` arguments for
+    `Timestamp`.
+    """
+
     created_on = Timestamp(default=TimestampNow())
 
 
@@ -22,13 +27,23 @@ def tearDown(self):
         MyTable.alter().drop_table().run_sync()
 
     def test_timestamp(self):
+        """
+        Make sure a datetime can be stored and retrieved.
+        """
         created_on = datetime.datetime.now()
         row = MyTable(created_on=created_on)
         row.save().run_sync()
 
         result = MyTable.objects().first().run_sync()
         self.assertEqual(result.created_on, created_on)
 
+    def test_timezone_aware(self):
+        """
+        Raise an error if a timezone aware datetime is given as a default.
+        """
+        with self.assertRaises(ValueError):
+            Timestamp(default=datetime.datetime.now(tz=datetime.timezone.utc))
+
 
 class TestTimestampDefault(TestCase):
     def setUp(self):
@@ -38,6 +53,9 @@ def tearDown(self):
         MyTableDefault.alter().drop_table().run_sync()
 
     def test_timestamp(self):
+        """
+        Make sure the default values get created correctly.
+        """
         created_on = datetime.datetime.now()
         row = MyTableDefault()
         row.save().run_sync()
@@ -46,3 +64,4 @@ def test_timestamp(self):
         self.assertTrue(
             result.created_on - created_on < datetime.timedelta(seconds=1)
         )
+        self.assertTrue(result.created_on.tzinfo is None)