Merge pull request piccolo-orm#12 from piccolo-orm/interval_column_type

added Interval column type

Author: dantownsend

docs/src/piccolo/schema/column_types.rst

@@ -117,6 +117,12 @@ Date
 
 .. autoclass:: Date
 
+========
+Interval
+========
+
+.. autoclass:: Interval
+
 ====
 Time
 ====

piccolo/apps/playground/commands/run.py

@@ -12,6 +12,7 @@
     Varchar,
     ForeignKey,
     Integer,
+    Interval,
     Numeric,
     Timestamp,
     UUID,
@@ -40,6 +41,7 @@ class Concert(Table):
     band_2 = ForeignKey(Band)
     venue = ForeignKey(Venue)
     starts = Timestamp()
+    duration = Interval()
 
 
 class Ticket(Table):
@@ -91,6 +93,7 @@ def populate():
         band_2=rustaceans.id,
         venue=venue.id,
         starts=datetime.datetime.now() + datetime.timedelta(days=7),
+        duration=datetime.timedelta(hours=2),
     )
     concert.save().run_sync()
 

piccolo/columns/__init__.py

@@ -4,6 +4,7 @@
     Float,
     ForeignKey,
     Integer,
+    Interval,
     Numeric,
     PrimaryKey,
     Real,

piccolo/columns/column_types.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import copy
-from datetime import datetime, date, time
+from datetime import datetime, date, time, timedelta
 import decimal
 import typing as t
 import uuid
@@ -9,6 +9,7 @@
 from piccolo.columns.operators.string import ConcatPostgres, ConcatSQLite
 from piccolo.columns.defaults.date import DateArg, DateNow, DateCustom
 from piccolo.columns.defaults.time import TimeArg, TimeNow, TimeCustom
+from piccolo.columns.defaults.interval import IntervalArg, IntervalCustom
 from piccolo.columns.defaults.timestamp import (
     TimestampArg,
     TimestampNow,
@@ -612,6 +613,57 @@ def __init__(self, default: TimeArg = TimeNow(), **kwargs) -> None:
         super().__init__(**kwargs)
 
 
+class Interval(Column):  # lgtm [py/missing-equals]
+    """
+    Used for storing timedeltas. Uses the ``timedelta`` type for values.
+
+    **Example**
+
+    .. code-block:: python
+
+        from datetime import timedelta
+
+        class Concert(Table):
+            duration = Interval()
+
+        # Create
+        >>> Concert(
+        >>>    duration=timedelta(hours=2)
+        >>> ).save().run_sync()
+
+        # Query
+        >>> Concert.select(Concert.duration).run_sync()
+        {'duration': datetime.timedelta(seconds=7200)}
+
+    """
+
+    value_type = timedelta
+
+    def __init__(
+        self, default: IntervalArg = IntervalCustom(), **kwargs
+    ) -> None:
+        self._validate_default(default, IntervalArg.__args__)  # type: ignore
+
+        if isinstance(default, timedelta):
+            default = IntervalCustom.from_timedelta(default)
+
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+    @property
+    def column_type(self):
+        engine_type = self._meta.table._meta.db.engine_type
+        if engine_type == "postgres":
+            return "INTERVAL"
+        elif engine_type == "sqlite":
+            # We can't use 'INTERVAL' because the type affinity in SQLite would
+            # make it an integer - but we need a numeric field.
+            # https://sqlite.org/datatype3.html#determination_of_column_affinity
+            return "SECONDS"
+        raise Exception("Unrecognized engine type")
+
+
 ###############################################################################
 
 

piccolo/columns/defaults/interval.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+import datetime
+import typing as t
+
+from .base import Default
+
+
+class IntervalCustom(Default):  # lgtm [py/missing-equals]
+    def __init__(
+        self,
+        weeks: int = 0,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: int = 0,
+        milliseconds: int = 0,
+        microseconds: int = 0,
+    ):
+        self.weeks = weeks
+        self.days = days
+        self.hours = hours
+        self.minutes = minutes
+        self.seconds = seconds
+        self.milliseconds = milliseconds
+        self.microseconds = microseconds
+
+    @property
+    def timedelta(self):
+        return datetime.timedelta(
+            weeks=self.weeks,
+            days=self.days,
+            hours=self.hours,
+            minutes=self.minutes,
+            seconds=self.seconds,
+            milliseconds=self.milliseconds,
+            microseconds=self.microseconds,
+        )
+
+    @property
+    def postgres(self):
+        value = self.get_postgres_interval_string(
+            attributes=[
+                "weeks",
+                "days",
+                "hours",
+                "minutes",
+                "seconds",
+                "milliseconds",
+                "microseconds",
+            ]
+        )
+        return f"'{value}'"
+
+    @property
+    def sqlite(self):
+        return self.timedelta.total_seconds()
+
+    def python(self):
+        return self.timedelta
+
+    @classmethod
+    def from_timedelta(cls, instance: datetime.timedelta):
+        return cls(
+            days=instance.days,
+            seconds=instance.seconds,
+            microseconds=instance.microseconds,
+        )
+
+
+###############################################################################
+
+IntervalArg = t.Union[
+    IntervalCustom, None, datetime.timedelta,
+]
+
+
+__all__ = [
+    "IntervalArg",
+    "IntervalCustom",
+]

piccolo/engine/sqlite.py

@@ -51,6 +51,13 @@ def convert_date_in(value):
     return value.isoformat()
 
 
+def convert_timedelta_in(value):
+    """
+    Converts the timedelta value being passed into sqlite.
+    """
+    return value.total_seconds()
+
+
 def convert_numeric_out(value: bytes) -> Decimal:
     """
     Convert float values into Decimals.
@@ -78,21 +85,30 @@ def convert_date_out(value: bytes) -> datetime.date:
 
 def convert_time_out(value: bytes) -> datetime.time:
     """
-    If the value is a uuid, convert it to a UUID instance.
+    If the value is a time, convert it to a UUID instance.
     """
     return datetime.time.fromisoformat(value.decode("utf8"))
 
 
+def convert_seconds_out(value: bytes) -> datetime.timedelta:
+    """
+    If the value is from a seconds column, convert it to a timedelta instance.
+    """
+    return datetime.timedelta(seconds=float(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_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.timedelta, convert_timedelta_in)
 
 ###############################################################################
 
@@ -317,6 +333,8 @@ def create_db(self, migrate=False):
         else:
             raise Exception(f"Database at {self.path} already exists")
         if migrate:
+            # Commented out for now, as migrations for SQLite aren't as
+            # well supported as Postgres.
             # from piccolo.commands.migration.forwards import (
             #     ForwardsMigrationManager,
             # )

piccolo/query/base.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import datetime
 import itertools
 from time import time
 import typing as t
@@ -26,7 +27,9 @@ def default(obj):
     """
     Used for handling edge cases which orjson can't serialise out of the box.
     """
-    if isinstance(obj, UUID):
+    # This is the asyncpg UUID, not the builtin Python UUID, which orjon can
+    # serialise out of the box.
+    if isinstance(obj, (UUID, datetime.timedelta)):
         return str(obj)
     raise TypeError
 

requirements.txt

@@ -3,6 +3,6 @@ asgiref==3.2.10
 asyncpg==0.21.0
 black==19.10b0
 colorama==0.4.3
-orjson==3.3.1
+orjson==3.4.0
 Jinja2==2.11.2
 targ==0.1.*

tests/columns/test_interval.py

@@ -0,0 +1,102 @@
+import datetime
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Interval
+from piccolo.columns.defaults.interval import IntervalCustom
+
+
+class MyTable(Table):
+    interval = Interval()
+
+
+class MyTableDefault(Table):
+    interval = Interval(default=IntervalCustom(days=1))
+
+
+class TestInterval(TestCase):
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_interval(self):
+        # Test a range of different timedeltas
+        intervals = [
+            datetime.timedelta(weeks=1),
+            datetime.timedelta(days=1),
+            datetime.timedelta(hours=1),
+            datetime.timedelta(minutes=1),
+            datetime.timedelta(seconds=1),
+            datetime.timedelta(milliseconds=1),
+            datetime.timedelta(microseconds=1),
+        ]
+
+        for interval in intervals:
+            # Make sure that saving it works OK.
+            row = MyTable(interval=interval)
+            row.save().run_sync()
+
+            # Make sure that fetching it back works OK.
+            result = (
+                MyTable.objects()
+                .where(MyTable.id == row.id)
+                .first()
+                .run_sync()
+            )
+            self.assertEqual(result.interval, interval)
+
+    def test_interval_where_clause(self):
+        """
+        Make sure a range of where clauses resolve correctly.
+        """
+        interval = datetime.timedelta(hours=2)
+        row = MyTable(interval=interval)
+        row.save().run_sync()
+
+        result = (
+            MyTable.objects()
+            .where(MyTable.interval < datetime.timedelta(hours=3))
+            .first()
+            .run_sync()
+        )
+        self.assertTrue(result is not None)
+
+        result = (
+            MyTable.objects()
+            .where(MyTable.interval < datetime.timedelta(days=1))
+            .first()
+            .run_sync()
+        )
+        self.assertTrue(result is not None)
+
+        result = (
+            MyTable.objects()
+            .where(MyTable.interval > datetime.timedelta(minutes=1))
+            .first()
+            .run_sync()
+        )
+        self.assertTrue(result is not None)
+
+        result = (
+            MyTable.exists()
+            .where(MyTable.interval == datetime.timedelta(hours=2))
+            .run_sync()
+        )
+        self.assertTrue(result)
+
+
+class TestIntervalDefault(TestCase):
+    def setUp(self):
+        MyTableDefault.create_table().run_sync()
+
+    def tearDown(self):
+        MyTableDefault.alter().drop_table().run_sync()
+
+    def test_interval(self):
+        row = MyTableDefault()
+        row.save().run_sync()
+
+        result = MyTableDefault.objects().first().run_sync()
+        self.assertTrue(result.interval.days == 1)