]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Add datetime binary adapter
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Wed, 12 May 2021 23:13:20 +0000 (01:13 +0200)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 13 May 2021 01:43:19 +0000 (03:43 +0200)
Only for naive datetimes, timezone not handled yet.

psycopg3/psycopg3/types/__init__.py
psycopg3/psycopg3/types/date.py
tests/fix_faker.py
tests/types/test_date.py

index abf69b902ddeb9de78f54cad0660dc6a04547ede..058fa46e2b76c774545cc9cd1e8758d983c9c17e 100644 (file)
@@ -87,7 +87,9 @@ from .date import (
     TimeTzDumper as TimeTzDumper,
     TimeTzBinaryDumper as TimeTzBinaryDumper,
     DateTimeTzDumper as DateTimeTzDumper,
+    DateTimeTzBinaryDumper as DateTimeTzBinaryDumper,
     DateTimeDumper as DateTimeDumper,
+    DateTimeBinaryDumper as DateTimeBinaryDumper,
     TimeDeltaDumper as TimeDeltaDumper,
     DateLoader as DateLoader,
     DateBinaryLoader as DateBinaryLoader,
@@ -96,6 +98,7 @@ from .date import (
     TimeTzLoader as TimeTzLoader,
     TimeTzBinaryLoader as TimeTzBinaryLoader,
     TimestampLoader as TimestampLoader,
+    TimestampBinaryLoader as TimestampBinaryLoader,
     TimestamptzLoader as TimestamptzLoader,
     IntervalLoader as IntervalLoader,
 )
@@ -214,6 +217,7 @@ def register_default_globals(ctx: AdaptContext) -> None:
     TimeDumper.register("datetime.time", ctx)
     TimeBinaryDumper.register("datetime.time", ctx)
     DateTimeTzDumper.register("datetime.datetime", ctx)
+    DateTimeTzBinaryDumper.register("datetime.datetime", ctx)
     TimeDeltaDumper.register("datetime.timedelta", ctx)
     DateLoader.register("date", ctx)
     DateBinaryLoader.register("date", ctx)
@@ -222,6 +226,7 @@ def register_default_globals(ctx: AdaptContext) -> None:
     TimeTzLoader.register("timetz", ctx)
     TimeTzBinaryLoader.register("timetz", ctx)
     TimestampLoader.register("timestamp", ctx)
+    TimestampBinaryLoader.register("timestamp", ctx)
     TimestamptzLoader.register("timestamptz", ctx)
     IntervalLoader.register("interval", ctx)
 
index 762d955fca19f17df261fac3f3efe40fadace5fa..a10cfa40daa881a5d20ec8f4c74c5482741287c4 100644 (file)
@@ -29,9 +29,10 @@ _unpack_timetz = cast(
 )
 _pack_timetz = cast(Callable[[int, int], bytes], struct.Struct("!qi").pack)
 
-_pg_date_epoch = date(2000, 1, 1).toordinal()
-_py_date_min = date.min.toordinal()
-_py_date_max = date.max.toordinal()
+_pg_date_epoch_days = date(2000, 1, 1).toordinal()
+_pg_datetime_epoch = datetime(2000, 1, 1)
+_pg_datetimetz_epoch = datetime(2000, 1, 1, tzinfo=timezone.utc)
+_py_date_min_days = date.min.toordinal()
 
 
 class DateDumper(Dumper):
@@ -51,7 +52,7 @@ class DateBinaryDumper(Dumper):
     _oid = builtins["date"].oid
 
     def dump(self, obj: date) -> bytes:
-        days = obj.toordinal() - _pg_date_epoch
+        days = obj.toordinal() - _pg_date_epoch_days
         return _pack_int4(days)
 
 
@@ -123,18 +124,11 @@ class TimeTzBinaryDumper(TimeBinaryDumper):
         return _pack_timetz(ms, -int(off.total_seconds()))
 
 
-class DateTimeTzDumper(Dumper):
-
-    format = Format.TEXT
+class _BaseDateTimeDumper(Dumper):
 
     # Can change to timestamp type if the object dumped is naive
     _oid = builtins["timestamptz"].oid
 
-    def dump(self, obj: datetime) -> bytes:
-        # NOTE: whatever the PostgreSQL DateStyle input format (DMY, MDY, YMD)
-        # the YYYY-MM-DD is always understood correctly.
-        return str(obj).encode("utf8")
-
     def get_key(
         self, obj: datetime, format: Pg3Format
     ) -> Union[type, Tuple[type]]:
@@ -145,6 +139,19 @@ class DateTimeTzDumper(Dumper):
         else:
             return (self.cls,)
 
+    def upgrade(self, obj: datetime, format: Pg3Format) -> "Dumper":
+        raise NotImplementedError
+
+
+class DateTimeTzDumper(_BaseDateTimeDumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj: datetime) -> bytes:
+        # NOTE: whatever the PostgreSQL DateStyle input format (DMY, MDY, YMD)
+        # the YYYY-MM-DD is always understood correctly.
+        return str(obj).encode("utf8")
+
     def upgrade(self, obj: datetime, format: Pg3Format) -> "Dumper":
         if obj.tzinfo:
             return self
@@ -156,6 +163,42 @@ class DateTimeDumper(DateTimeTzDumper):
     _oid = builtins["timestamp"].oid
 
 
+class DateTimeTzBinaryDumper(_BaseDateTimeDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj: datetime) -> bytes:
+        raise NotImplementedError
+
+    def upgrade(self, obj: datetime, format: Pg3Format) -> "Dumper":
+        if obj.tzinfo:
+            return self
+        else:
+            return DateTimeBinaryDumper(self.cls)
+
+
+class DateTimeBinaryDumper(DateTimeTzBinaryDumper):
+    _oid = builtins["timestamp"].oid
+
+    # Somewhere, between year 2270 and 2275, float rounding in total_seconds
+    # cause us errors: switch to an algorithm without rounding before then.
+    _delta_prec_loss = (
+        datetime(2250, 1, 1) - _pg_datetime_epoch
+    ).total_seconds()
+
+    def dump(self, obj: datetime) -> bytes:
+        delta = obj - _pg_datetime_epoch
+        secs = delta.total_seconds()
+        if secs < self._delta_prec_loss:
+            micros = int(1_000_000 * secs)
+        else:
+            micros = (
+                1_000_000 * (86_400 * delta.days + delta.seconds)
+                + delta.microseconds
+            )
+        return _pack_int8(micros)
+
+
 class TimeDeltaDumper(Dumper):
 
     format = Format.TEXT
@@ -250,11 +293,11 @@ class DateBinaryLoader(Loader):
     format = Format.BINARY
 
     def load(self, data: Buffer) -> date:
-        days = _unpack_int4(data)[0] + _pg_date_epoch
-        if _py_date_min <= days <= _py_date_max:
+        days = _unpack_int4(data)[0] + _pg_date_epoch_days
+        try:
             return date.fromordinal(days)
-        else:
-            if days < _py_date_min:
+        except ValueError:
+            if days < _py_date_min_days:
                 raise DataError("date too small (before year 1)")
             else:
                 raise DataError("date too large (after year 10K)")
@@ -431,6 +474,21 @@ class TimestampLoader(DateLoader):
                 return 0
 
 
+class TimestampBinaryLoader(Loader):
+
+    format = Format.BINARY
+
+    def load(self, data: Buffer) -> datetime:
+        micros = _unpack_int8(data)[0]
+        try:
+            return _pg_datetime_epoch + timedelta(microseconds=micros)
+        except OverflowError:
+            if micros <= 0:
+                raise DataError("timestamp too small (before year 1)")
+            else:
+                raise DataError("timestamp too large (after year 10K)")
+
+
 class TimestamptzLoader(TimestampLoader):
 
     format = Format.TEXT
index fcc8103abda341c2029f31f09ee5500962179c07..db159fdb734bf25bfc7c39406eb84310fd9815f7 100644 (file)
@@ -235,6 +235,11 @@ class Faker:
         day = randrange(dt.date.max.toordinal())
         return dt.date.fromordinal(day + 1)
 
+    def make_datetime(self, spec):
+        delta = dt.datetime.max - dt.datetime.min
+        micros = randrange((delta.days + 1) * 24 * 60 * 60 * 1_000_000)
+        return dt.datetime.min + dt.timedelta(microseconds=micros)
+
     def make_Decimal(self, spec):
         if random() >= 0.99:
             if self.conn.info.server_version >= 140000:
index cdb19affcf800213e4d0397c9c4861351ee57a21..27b48d265465825665364a6f45b10fab9592f578 100644 (file)
@@ -106,27 +106,23 @@ def test_load_date_overflow_binary(conn, val):
         ("min", "0001-01-01 00:00"),
         ("1000,1,1,0,0", "1000-01-01 00:00"),
         ("2000,1,1,0,0", "2000-01-01 00:00"),
-        ("2000,12,31,23,59,59,999999", "2000-12-31 23:59:59.999999"),
-        ("3000,1,1,0,0", "3000-01-01 00:00"),
+        ("2000,1,2,3,4,5,6", "2000-01-02 03:04:05.000006"),
+        ("2000,1,2,3,4,5,678", "2000-01-02 03:04:05.000678"),
+        ("2000,1,2,3,0,0,456789", "2000-01-02 03:00:00.456789"),
+        ("2000,1,1,0,0,0,1", "2000-01-01 00:00:00.000001"),
+        ("2200,1,1,0,0,0,1", "2200-01-01 00:00:00.000001"),
+        ("2300,1,1,0,0,0,1", "2300-01-01 00:00:00.000001"),
+        ("7000,1,1,0,0,0,1", "7000-01-01 00:00:00.000001"),
         ("max", "9999-12-31 23:59:59.999999"),
     ],
 )
-def test_dump_datetime(conn, val, expr):
-    cur = conn.cursor()
-    cur.execute("set timezone to '+02:00'")
-    cur.execute(f"select '{expr}'::timestamp = %s", (as_dt(val),))
-    assert cur.fetchone()[0] is True
-
-
-@pytest.mark.xfail  # TODO: binary dump
-@pytest.mark.parametrize(
-    "val, expr",
-    [("2000,1,1,0,0", "'2000-01-01 00:00'::timestamp")],
-)
-def test_dump_datetime_binary(conn, val, expr):
+@pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
+def test_dump_datetime(conn, val, expr, fmt_in):
     cur = conn.cursor()
     cur.execute("set timezone to '+02:00'")
-    cur.execute(f"select {expr} = %b", (as_dt(val),))
+    cur.execute(f"select %{fmt_in}", (as_dt(val),))
+    print(cur.fetchone()[0])
+    cur.execute(f"select '{expr}'::timestamp = %{fmt_in}", (as_dt(val),))
     assert cur.fetchone()[0] is True
 
 
@@ -135,26 +131,26 @@ def test_dump_datetime_datestyle(conn, datestyle_in):
     cur = conn.cursor()
     cur.execute(f"set datestyle = ISO, {datestyle_in}")
     cur.execute(
-        "select 'epoch'::timestamp + '1d 3h 4m 5s'::interval = %s",
+        "select 'epoch'::timestamp + '1d 3h 4m 5s'::interval = %t",
         (dt.datetime(1970, 1, 2, 3, 4, 5),),
     )
     assert cur.fetchone()[0] is True
 
 
-@pytest.mark.parametrize(
-    "val, expr",
-    [
-        ("min", "0001-01-01"),
-        ("1000,1,1", "1000-01-01"),
-        ("2000,1,1", "2000-01-01"),
-        ("2000,1,2,3,4,5,6", "2000-01-02 03:04:05.000006"),
-        ("2000,1,2,3,4,5,678", "2000-01-02 03:04:05.000678"),
-        ("2000,1,2,3,0,0,456789", "2000-01-02 03:00:00.456789"),
-        ("2000,12,31", "2000-12-31"),
-        ("3000,1,1", "3000-01-01"),
-        ("max", "9999-12-31 23:59:59.999999"),
-    ],
-)
+load_datetime_samples = [
+    ("min", "0001-01-01"),
+    ("1000,1,1", "1000-01-01"),
+    ("2000,1,1", "2000-01-01"),
+    ("2000,1,2,3,4,5,6", "2000-01-02 03:04:05.000006"),
+    ("2000,1,2,3,4,5,678", "2000-01-02 03:04:05.000678"),
+    ("2000,1,2,3,0,0,456789", "2000-01-02 03:00:00.456789"),
+    ("2000,12,31", "2000-12-31"),
+    ("3000,1,1", "3000-01-01"),
+    ("max", "9999-12-31 23:59:59.999999"),
+]
+
+
+@pytest.mark.parametrize("val, expr", load_datetime_samples)
 @pytest.mark.parametrize("datestyle_out", ["ISO", "Postgres", "SQL", "German"])
 @pytest.mark.parametrize("datestyle_in", ["DMY", "MDY", "YMD"])
 def test_load_datetime(conn, val, expr, datestyle_in, datestyle_out):
@@ -165,6 +161,14 @@ def test_load_datetime(conn, val, expr, datestyle_in, datestyle_out):
     assert cur.fetchone()[0] == as_dt(val)
 
 
+@pytest.mark.parametrize("val, expr", load_datetime_samples)
+def test_load_datetime_binary(conn, val, expr):
+    cur = conn.cursor(binary=True)
+    cur.execute("set timezone to '+02:00'")
+    cur.execute(f"select '{expr}'::timestamp")
+    assert cur.fetchone()[0] == as_dt(val)
+
+
 @pytest.mark.parametrize("val", ["min", "max"])
 @pytest.mark.parametrize("datestyle_out", ["ISO", "Postgres", "SQL", "German"])
 def test_load_datetime_overflow(conn, val, datestyle_out):
@@ -178,6 +182,17 @@ def test_load_datetime_overflow(conn, val, datestyle_out):
         cur.fetchone()[0]
 
 
+@pytest.mark.parametrize("val", ["min", "max"])
+def test_load_datetime_overflow_binary(conn, val):
+    cur = conn.cursor(binary=True)
+    cur.execute(
+        "select %t::timestamp + %s * '1s'::interval",
+        (as_dt(val), -1 if val == "min" else 1),
+    )
+    with pytest.raises(DataError):
+        cur.fetchone()[0]
+
+
 #
 # datetime+tz tests
 #