From 1898190fa105be689b16c382233e9ac6631eaf7a Mon Sep 17 00:00:00 2001 From: Przemek Denkiewicz Date: Fri, 19 Jul 2024 14:09:36 +0200 Subject: [PATCH 1/2] Map INTERVAL types to Python types --- tests/integration/test_types_integration.py | 26 +++++++++++++--- trino/mapper.py | 33 +++++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py index dc8fe683..fa4eec81 100644 --- a/tests/integration/test_types_integration.py +++ b/tests/integration/test_types_integration.py @@ -6,6 +6,7 @@ from zoneinfo import ZoneInfo import pytest +from dateutil.relativedelta import relativedelta import trino from tests.integration.conftest import trino_version @@ -733,13 +734,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() diff --git a/trino/mapper.py b/trino/mapper.py index b56037d7..39939d48 100644 --- a/trino/mapper.py +++ b/trino/mapper.py @@ -8,6 +8,8 @@ from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar from zoneinfo import ZoneInfo +from dateutil.relativedelta import relativedelta + import trino.exceptions from trino.types import ( POWERS_OF_TEN, @@ -167,6 +169,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 @@ -271,6 +300,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': From 7472f37e30148bc78105db2297f3a8f163f9e942 Mon Sep 17 00:00:00 2001 From: Damian Owsianny Date: Wed, 9 Oct 2024 22:08:26 +0200 Subject: [PATCH 2/2] Raise TrinoDataError if the INTERVAL type value exceeds the maximum or minimum limit. Add test cases for edge scenarios --- tests/integration/test_types_integration.py | 91 ++++++++++++++++----- trino/mapper.py | 9 +- 2 files changed, 78 insertions(+), 22 deletions(-) diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py index fa4eec81..73be522b 100644 --- a/tests/integration/test_types_integration.py +++ b/tests/integration/test_types_integration.py @@ -735,30 +735,79 @@ def create_timezone(timezone_str: str) -> tzinfo: 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() + ( + 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)) + # max supported INTERVAL in Trino + .add_field( + sql="INTERVAL '178956970-7' YEAR TO MONTH", + python=relativedelta(years=178956970, months=7)) + # min supported INTERVAL in Trino + .add_field( + sql="INTERVAL '-178956970-8' YEAR TO MONTH", + python=relativedelta(years=-178956970, months=-8)) + ).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 '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() + ( + SqlTest(trino_connection) + .add_field( + sql="CAST(null AS INTERVAL DAY TO SECOND)", + python=None) + .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)) + # max supported timedelta in Python + .add_field( + sql="INTERVAL '999999999 23:59:59.999' DAY TO SECOND", + python=timedelta(days=999999999, hours=23, minutes=59, seconds=59, milliseconds=999)) + # min supported timedelta in Python + .add_field( + sql="INTERVAL '-999999999' DAY", + python=timedelta(days=-999999999)) + ).execute() + + SqlExpectFailureTest(trino_connection).execute("INTERVAL '1000000000' DAY") + SqlExpectFailureTest(trino_connection).execute("INTERVAL '-999999999 00:00:00.001' DAY TO SECOND") def test_array(trino_connection): diff --git a/trino/mapper.py b/trino/mapper.py index 39939d48..53ed6193 100644 --- a/trino/mapper.py +++ b/trino/mapper.py @@ -193,7 +193,14 @@ def map(self, value: Any) -> Optional[timedelta]: 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) + try: + return timedelta(days=days, hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds) + except OverflowError as e: + error_str = ( + f"Could not convert '{value}' into the associated python type, as the value " + "exceeds the maximum or minimum limit." + ) + raise trino.exceptions.TrinoDataError(error_str) from e class ArrayValueMapper(ValueMapper[List[Optional[Any]]]):