From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 23 Aug 2020 20:41:15 +0000 (+0100)
Subject: Raise a polite exception for dates >= Y10K
X-Git-Tag: 3.0.dev0~438
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f7030f3bb3d22b98b34130e4292e961c1534920c;p=thirdparty%2Fpsycopg.git

Raise a polite exception for dates >= Y10K
---

diff --git a/psycopg3/psycopg3/types/date.py b/psycopg3/psycopg3/types/date.py
index 6a7d0b538..51de1b9ee 100644
--- a/psycopg3/psycopg3/types/date.py
+++ b/psycopg3/psycopg3/types/date.py
@@ -4,6 +4,7 @@ Adapters for date/time types.
 
 # Copyright (C) 2020 The Psycopg Team
 
+import re
 import codecs
 from datetime import date
 
@@ -37,14 +38,13 @@ class DateLoader(Loader):
     def __init__(self, oid: int, context: AdaptContext):
         super().__init__(oid, context)
 
-        if self.connection:
-            ds = self.connection.pgconn.parameter_status(b"DateStyle")
-            if not ds or ds.startswith(b"ISO"):
-                pass    # Default: YMD
-            elif ds.startswith(b"German"):
-                self.load = self.load_dmy  # type: ignore
-            elif ds.startswith(b"SQL") or ds.startswith(b"Postgres"):
-                self.load = self.load_mdy  # type: ignore
+        ds = self._get_datestyle()
+        if ds == b"ISO":
+            pass  # Default: YMD
+        elif ds == b"German":
+            self.load = self.load_dmy  # type: ignore
+        elif ds == b"SQL" or ds == b"Postgres":
+            self.load = self.load_mdy  # type: ignore
 
     def load_ymd(self, data: bytes) -> date:
         try:
@@ -72,12 +72,37 @@ class DateLoader(Loader):
 
         return self._raise_error(data, exc)
 
+    def _get_datestyle(self) -> bytes:
+        """Return the PostgreSQL output datestyle of the connection."""
+        if self.connection:
+            ds = self.connection.pgconn.parameter_status(b"DateStyle")
+            if ds:
+                return ds.split(b",", 1)[0]
+
+        return b"ISO"
+
     def _raise_error(self, data: bytes, exc: ValueError) -> date:
         # Most likely we received a BC date, which Python doesn't support
         # Otherwise the unexpected value is displayed in the exception.
         if data.endswith(b"BC"):
             raise InterfaceError(
-                "BC date not supported by Python:"
-                f" {data.decode('utf8', 'replace')}")
+                "Python doesn't support BC date:"
+                f" got {data.decode('utf8', 'replace')}"
+            )
+
+        # Find the year from the date. This is not the fast path so we don't
+        # need crazy speed.
+        ds = self._get_datestyle()
+        if ds == b"ISO":
+            year = int(data.split(b"-", 1)[0])
         else:
-            raise exc
+            year = int(re.split(rb"[-/\.]", data)[-1])
+
+        if year > 9999:
+            raise InterfaceError(
+                "Python date doesn't support years after 9999:"
+                f" got {data.decode('utf8', 'replace')}"
+            )
+
+        # We genuinely received something we cannot parse
+        raise exc
diff --git a/tests/types/test_date.py b/tests/types/test_date.py
index dcfdbadb2..4fd76ac01 100644
--- a/tests/types/test_date.py
+++ b/tests/types/test_date.py
@@ -9,7 +9,6 @@ from psycopg3.adapt import Format
     "val, expr",
     [
         (dt.date.min, "'0001-01-01'::date"),
-        (dt.date(1, 1, 1), "'0001-01-01'::date"),
         (dt.date(1000, 1, 1), "'1000-01-01'::date"),
         (dt.date(2000, 1, 1), "'2000-01-01'::date"),
         (dt.date(2000, 12, 31), "'2000-12-31'::date"),
@@ -44,11 +43,12 @@ def test_dump_date_datestyle(conn, datestyle_in):
 @pytest.mark.parametrize(
     "val, expr",
     [
-        (dt.date(1, 1, 1), "'0001-01-01'::date"),
+        (dt.date.min, "'0001-01-01'::date"),
         (dt.date(1000, 1, 1), "'1000-01-01'::date"),
         (dt.date(2000, 1, 1), "'2000-01-01'::date"),
         (dt.date(2000, 12, 31), "'2000-12-31'::date"),
         (dt.date(3000, 1, 1), "'3000-01-01'::date"),
+        (dt.date.max, "'9999-12-31'::date"),
     ],
 )
 def test_load_date(conn, val, expr):
@@ -79,6 +79,15 @@ def test_load_date_datestyle(conn, datestyle_out):
 def test_load_date_bc(conn, datestyle_out):
     cur = conn.cursor()
     cur.execute(f"set datestyle = {datestyle_out}, YMD")
-    cur.execute("select '0001-01-01'::date - 1")
+    cur.execute("select %s - 1", (dt.date.min,))
+    with pytest.raises(psycopg3.InterfaceError):
+        cur.fetchone()[0]
+
+
+@pytest.mark.parametrize("datestyle_out", ["ISO", "Postgres", "SQL", "German"])
+def test_load_date_too_large(conn, datestyle_out):
+    cur = conn.cursor()
+    cur.execute(f"set datestyle = {datestyle_out}, YMD")
+    cur.execute("select %s + 1", (dt.date.max,))
     with pytest.raises(psycopg3.InterfaceError):
         cur.fetchone()[0]