Map INTERVAL types to Python types

Author: hovaesco

tests/integration/test_types_integration.py

@@ -4,6 +4,8 @@
 from datetime import date, datetime, time, timedelta, timezone, tzinfo
 from decimal import Decimal
 
+from dateutil.relativedelta import relativedelta
+
 try:
     from zoneinfo import ZoneInfo
 except ModuleNotFoundError:
@@ -737,13 +739,30 @@ def create_timezone(timezone_str: str) -> tzinfo:
         return ZoneInfo(timezone_str)
 
 
-def test_interval(trino_connection):
+def test_interval_year_to_month(trino_connection):
     SqlTest(trino_connection) \
         .add_field(sql="CAST(null AS INTERVAL YEAR TO MONTH)", python=None) \
+        .add_field(sql="INTERVAL '10' YEAR", python=relativedelta(years=10)) \
+        .add_field(sql="INTERVAL '-5' YEAR", python=relativedelta(years=-5)) \
+        .add_field(sql="INTERVAL '3' MONTH", python=relativedelta(months=3)) \
+        .add_field(sql="INTERVAL '-18' MONTH", python=relativedelta(years=-1, months=-6)) \
+        .add_field(sql="INTERVAL '30' MONTH", python=relativedelta(years=2, months=6)) \
+        .add_field(sql="INTERVAL '124-30' YEAR TO MONTH", python=relativedelta(years=126, months=6)) \
+        .execute()
+
+
+def test_interval_day_to_second(trino_connection):
+    SqlTest(trino_connection) \
         .add_field(sql="CAST(null AS INTERVAL DAY TO SECOND)", python=None) \
-        .add_field(sql="INTERVAL '3' MONTH", python='0-3') \
-        .add_field(sql="INTERVAL '2' DAY", python='2 00:00:00.000') \
-        .add_field(sql="INTERVAL '-2' DAY", python='-2 00:00:00.000') \
+        .add_field(sql="INTERVAL '2' DAY", python=timedelta(days=2)) \
+        .add_field(sql="INTERVAL '-2' DAY", python=timedelta(days=-2)) \
+        .add_field(sql="INTERVAL '-2' SECOND", python=timedelta(seconds=-2)) \
+        .add_field(sql="INTERVAL '1 11:11:11.116555' DAY TO SECOND",
+                   python=timedelta(days=1, seconds=40271, microseconds=116000)) \
+        .add_field(sql="INTERVAL '-5 23:59:57.000' DAY TO SECOND", python=timedelta(days=-6, seconds=3)) \
+        .add_field(sql="INTERVAL '12 10:45' DAY TO MINUTE", python=timedelta(days=12, seconds=38700)) \
+        .add_field(sql="INTERVAL '45:32.123' MINUTE TO SECOND", python=timedelta(seconds=2732, microseconds=123000)) \
+        .add_field(sql="INTERVAL '32.123' SECOND", python=timedelta(seconds=32, microseconds=123000)) \
         .execute()
 
 

trino/mapper.py

@@ -8,6 +8,8 @@
 from decimal import Decimal
 from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar
 
+from dateutil.relativedelta import relativedelta
+
 if sys.version_info >= (3, 9):
     from zoneinfo import ZoneInfo
 else:
@@ -172,6 +174,33 @@ def _fraction_to_decimal(fractional_str: str) -> Decimal:
     return Decimal(fractional_str or 0) / POWERS_OF_TEN[len(fractional_str)]
 
 
+class IntervalYearToMonthMapper(ValueMapper[relativedelta]):
+    def map(self, value: Any) -> Optional[relativedelta]:
+        if value is None:
+            return None
+        is_negative = value[0] == "-"
+        years, months = (value[1:] if is_negative else value).split('-')
+        years, months = int(years), int(months)
+        if is_negative:
+            years, months = -years, -months
+        return relativedelta(years=years, months=months)
+
+
+class IntervalDayToSecondMapper(ValueMapper[timedelta]):
+    def map(self, value: Any) -> Optional[timedelta]:
+        if value is None:
+            return None
+        is_negative = value[0] == "-"
+        days, time = (value[1:] if is_negative else value).split(' ')
+        hours, minutes, seconds_milliseconds = time.split(':')
+        seconds, milliseconds = seconds_milliseconds.split('.')
+        days, hours, minutes, seconds, milliseconds = (int(days), int(hours), int(minutes), int(seconds),
+                                                       int(milliseconds))
+        if is_negative:
+            days, hours, minutes, seconds, milliseconds = -days, -hours, -minutes, -seconds, -milliseconds
+        return timedelta(days=days, hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds)
+
+
 class ArrayValueMapper(ValueMapper[List[Optional[Any]]]):
     def __init__(self, mapper: ValueMapper[Any]):
         self.mapper = mapper
@@ -276,6 +305,10 @@ def _create_value_mapper(self, column: Dict[str, Any]) -> ValueMapper[Any]:
             return TimestampValueMapper(self._get_precision(column))
         if col_type == 'timestamp with time zone':
             return TimestampWithTimeZoneValueMapper(self._get_precision(column))
+        if col_type == 'interval year to month':
+            return IntervalYearToMonthMapper()
+        if col_type == 'interval day to second':
+            return IntervalDayToSecondMapper()
 
         # structural types
         if col_type == 'array':