feat: Fir 18774 new date types (#229)

Author: stepansergeevitch

setup.cfg

@@ -29,6 +29,7 @@ install_requires =
     cryptography>=3.4.0
     httpx[http2]==0.23.0
     pydantic[dotenv]>=1.8.2,<1.10
+    python-dateutil>=2.8.2
     readerwriterlock==1.0.9
     sqlparse>=0.4.2
     trio<0.22.0

src/firebolt/async_db/_types.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from collections import namedtuple
 from datetime import date, datetime, timezone
 from decimal import Decimal
@@ -18,13 +19,26 @@
 try:
     from ciso8601 import parse_datetime  # type: ignore
 except ImportError:
-    # Unfortunately, there seems to be no support for optional bits in strptime
-    def parse_datetime(date_string: str) -> datetime:  # type: ignore
-        format = "%Y-%m-%d %H:%M:%S.%f"
-        # fromisoformat doesn't support milliseconds
-        if "." in date_string:
-            return datetime.strptime(date_string, format)
-        return datetime.fromisoformat(date_string)
+    unsupported_milliseconds_re = re.compile(r"(?<=\.)\d{1,5}(?!\d)")
+
+    def _fix_milliseconds(datetime_string: str) -> str:
+        # Fill milliseconds with 0 to have exactly 6 digits
+        # Python parser only supports 3 or 6 digit milliseconds untill 3.11
+        def align_ms(match: re.Match) -> str:
+            ms = match.group()
+            return ms + "0" * (6 - len(ms))
+
+        return re.sub(unsupported_milliseconds_re, align_ms, datetime_string)
+
+    def _fix_timezone(datetime_string: str) -> str:
+        # timezone, provided as +/-dd is not supported by datetime.
+        # We need to append :00 to it
+        if datetime_string[-3] in "+-":
+            return datetime_string + ":00"
+        return datetime_string
+
+    def parse_datetime(datetime_string: str) -> datetime:
+        return datetime.fromisoformat(_fix_timezone(_fix_milliseconds(datetime_string)))
 
 
 from firebolt.utils.exception import (
@@ -179,9 +193,12 @@ class _InternalType(Enum):
     # DATE
     Date = "Date"
     Date32 = "Date32"
+    PGDate = "PGDate"
 
     # DATETIME, TIMESTAMP
     DateTime = "DateTime"
+    TimestampNtz = "TimestampNtz"
+    TimestampTz = "TimestampTz"
 
     # Nullable(Nothing)
     Nothing = "Nothing"
@@ -203,7 +220,10 @@ def python_type(self) -> type:
             _InternalType.String: str,
             _InternalType.Date: date,
             _InternalType.Date32: date,
+            _InternalType.PGDate: date,
             _InternalType.DateTime: datetime,
+            _InternalType.TimestampNtz: datetime,
+            _InternalType.TimestampTz: datetime,
             # For simplicity, this could happen only during 'select null' query
             _InternalType.Nothing: str,
         }

tests/integration/dbapi/async/test_queries_async.py

@@ -97,13 +97,15 @@ async def test_connect_engine_name(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Connecting with engine name is handled properly."""
     await test_select(
         connection_engine_name,
         all_types_query,
         all_types_query_description,
         all_types_query_response,
+        timezone_name,
     )
 
 
@@ -112,13 +114,15 @@ async def test_connect_no_engine(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Connecting with engine name is handled properly."""
     await test_select(
         connection_no_engine,
         all_types_query,
         all_types_query_description,
         all_types_query_response,
+        timezone_name,
     )
 
 
@@ -127,9 +131,16 @@ async def test_select(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Select handles all data types properly."""
     with connection.cursor() as c:
+        assert (
+            await c.execute(f"SET advanced_mode=1") == -1
+        ), "Invalid set statment row count"
+        assert (
+            await c.execute(f"SET time_zone={timezone_name}") == -1
+        ), "Invalid set statment row count"
         assert await c.execute(all_types_query) == 1, "Invalid row count returned"
         assert c.rowcount == 1, "Invalid rowcount value"
         data = await c.fetchall()

tests/integration/dbapi/conftest.py

@@ -1,4 +1,4 @@
-from datetime import date, datetime
+from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
 from logging import getLogger
 from typing import List
@@ -67,8 +67,11 @@ def all_types_query() -> str:
         "'text' as \"string\", "
         "CAST('2021-03-28' AS DATE) as \"date\", "
         "CAST('1860-03-04' AS DATE_EXT) as \"date32\","
+        "pgdate '0001-01-01' as pgdate, "
         "CAST('2019-07-31 01:01:01' AS DATETIME) as \"datetime\", "
         "CAST('2019-07-31 01:01:01.1234' AS TIMESTAMP_EXT(4)) as \"datetime64\", "
+        "CAST('1111-01-05 17:04:42.123456' as timestampntz) as timestampntz, "
+        "'1111-01-05 17:04:42.123456'::timestamptz as timestamptz,"
         'true as "bool",'
         "[1,2,3,4] as \"array\", cast('1231232.123459999990457054844258706536' as "
         'decimal(38,30)) as "decimal", '
@@ -92,8 +95,11 @@ def all_types_query_description() -> List[Column]:
         Column("string", str, None, None, None, None, None),
         Column("date", date, None, None, None, None, None),
         Column("date32", date, None, None, None, None, None),
+        Column("pgdate", date, None, None, None, None, None),
         Column("datetime", datetime, None, None, None, None, None),
         Column("datetime64", DATETIME64(4), None, None, None, None, None),
+        Column("timestampntz", datetime, None, None, None, None, None),
+        Column("timestamptz", datetime, None, None, None, None, None),
         Column("bool", int, None, None, None, None, None),
         Column("array", ARRAY(int), None, None, None, None, None),
         Column("decimal", DECIMAL(38, 30), None, None, None, None, None),
@@ -102,7 +108,7 @@ def all_types_query_description() -> List[Column]:
 
 
 @fixture
-def all_types_query_response() -> List[ColType]:
+def all_types_query_response(timezone_offset_seconds: int) -> List[ColType]:
     return [
         [
             1,
@@ -118,8 +124,20 @@ def all_types_query_response() -> List[ColType]:
             "text",
             date(2021, 3, 28),
             date(1860, 3, 4),
+            date(1, 1, 1),
             datetime(2019, 7, 31, 1, 1, 1),
             datetime(2019, 7, 31, 1, 1, 1, 123400),
+            datetime(1111, 1, 5, 17, 4, 42, 123456),
+            datetime(
+                1111,
+                1,
+                5,
+                17,
+                4,
+                42,
+                123456,
+                tzinfo=timezone(timedelta(seconds=timezone_offset_seconds)),
+            ),
             1,
             [1, 2, 3, 4],
             Decimal("1231232.123459999990457054844258706536"),
@@ -128,6 +146,16 @@ def all_types_query_response() -> List[ColType]:
     ]
 
 
+@fixture
+def timezone_name() -> str:
+    return "Asia/Calcutta"
+
+
+@fixture
+def timezone_offset_seconds() -> int:
+    return 5 * 3600 + 53 * 60 + 28  # 05:53:28
+
+
 @fixture
 def create_drop_description() -> List[Column]:
     return [

tests/integration/dbapi/sync/test_queries.py

@@ -56,13 +56,15 @@ def test_connect_engine_name(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Connecting with engine name is handled properly."""
     test_select(
         connection_engine_name,
         all_types_query,
         all_types_query_description,
         all_types_query_response,
+        timezone_name,
     )
 
 
@@ -71,13 +73,15 @@ def test_connect_no_engine(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Connecting with engine name is handled properly."""
     test_select(
         connection_no_engine,
         all_types_query,
         all_types_query_description,
         all_types_query_response,
+        timezone_name,
     )
 
 
@@ -86,9 +90,15 @@ def test_select(
     all_types_query: str,
     all_types_query_description: List[Column],
     all_types_query_response: List[ColType],
+    timezone_name: str,
 ) -> None:
     """Select handles all data types properly."""
     with connection.cursor() as c:
+        assert c.execute(f"SET advanced_mode=1") == -1, "Invalid set statment row count"
+        assert (
+            c.execute(f"SET time_zone={timezone_name}") == -1
+        ), "Invalid set statment row count"
+
         assert c.execute(all_types_query) == 1, "Invalid row count returned"
         assert c.rowcount == 1, "Invalid rowcount value"
         data = c.fetchall()

tests/unit/async_db/conftest.py

@@ -50,8 +50,11 @@ def types_map() -> Dict[str, type]:
         "String": str,
         "Date": date,
         "Date32": date,
+        "PGDate": date,
         "DateTime": datetime,
         "DateTime64(7)": DATETIME64(7),
+        "TimestampNtz": datetime,
+        "TimestampTz": datetime,
         "Nullable(Nothing)": str,
         "Decimal(123, 4)": DECIMAL(123, 4),
         "Decimal(38,0)": DECIMAL(38, 0),

tests/unit/async_db/test_typing_parse.py

@@ -1,8 +1,8 @@
-from datetime import date, datetime
+from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
-from typing import Dict
+from typing import Dict, Optional
 
-from pytest import raises
+from pytest import mark, raises
 
 from firebolt.async_db import (
     ARRAY,
@@ -72,22 +72,73 @@ def test_parse_value_str() -> None:
     assert parse_value(None, str) is None, "Error parsing str: provided None"
 
 
-def test_parse_value_datetime() -> None:
-    """parse_value parses all date and datetime values correctly."""
-    # Date
-    assert parse_value("2021-12-31", date) == date(
-        2021, 12, 31
-    ), "Error parsing date: str provided"
-    assert parse_value("1860-12-31", date) == date(
-        1860, 12, 31
-    ), "Error parsing extended date: str provided"
-
-    assert parse_value(None, date) is None, "Error parsing date: None provided"
+@mark.parametrize(
+    "value,expected,case",
+    [
+        ("2021-12-31", date(2021, 12, 31), "str provided"),
+        ("0001-01-01", date(1, 1, 1), "range low provided"),
+        ("9999-12-31", date(9999, 12, 31), "range high provided"),
+        (None, None, "None provided"),
+        ("2021-12-31 23:59:59", date(2021, 12, 31), "datetime provided"),
+    ],
+)
+def test_parse_value_date(value: Optional[str], expected: Optional[date], case: str):
+    """parse_value parses all date values correctly."""
+    assert parse_value(value, date) == expected, f"Error parsing date: {case}"
+
+
+@mark.parametrize(
+    "value,expected,case",
+    [
+        (
+            "2021-12-31 23:59:59.1234",
+            datetime(2021, 12, 31, 23, 59, 59, 123400),
+            "str provided",
+        ),
+        (
+            "0001-01-01 00:00:00.000000",
+            datetime(1, 1, 1, 0, 0, 0, 0),
+            "range low provided",
+        ),
+        (
+            "9999-12-31 23:59:59.999999",
+            datetime(9999, 12, 31, 23, 59, 59, 999999),
+            "range high provided",
+        ),
+        (
+            "2021-12-31 23:59:59.1234-03",
+            datetime(
+                2021, 12, 31, 23, 59, 59, 123400, tzinfo=timezone(timedelta(hours=-3))
+            ),
+            "timezone provided",
+        ),
+        (
+            "2021-12-31 23:59:59.1234+05:30:12",
+            datetime(
+                2021,
+                12,
+                31,
+                23,
+                59,
+                59,
+                123400,
+                tzinfo=timezone(timedelta(hours=5, minutes=30, seconds=12)),
+            ),
+            "timezone with seconds provided",
+        ),
+        (None, None, "None provided"),
+        ("2021-12-31", datetime(2021, 12, 31), "date provided"),
+    ],
+)
+def test_parse_value_datetime(
+    value: Optional[str], expected: Optional[date], case: str
+):
+    """parse_value parses all date values correctly."""
+    assert parse_value(value, datetime) == expected, f"Error parsing datetime: {case}"
 
-    assert parse_value("2021-12-31 23:59:59", date) == date(
-        2021, 12, 31
-    ), "Error parsing date: datetime string provided"
 
+def test_parse_value_datetime_errors() -> None:
+    """parse_value parses all date and datetime values correctly."""
     with raises(ValueError):
         parse_value("abd", date)