Raise TrinoDataError if the INTERVAL type value exceeds the maximum o…

…r minimum limit. Add test cases for edge scenarios
trinodb · Oct 10, 2024 · d5f5e54 · d5f5e54
1 parent 78da73c
commit d5f5e54
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 24 deletions.
diff --git a/tests/integration/test_types_integration.py b/tests/integration/test_types_integration.py
@@ -740,30 +740,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
@@ -8,13 +8,13 @@
 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:
     from backports.zoneinfo import ZoneInfo
 
+from dateutil.relativedelta import relativedelta
+
 import trino.exceptions
 from trino.types import (
     POWERS_OF_TEN,
@@ -198,7 +198,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]]]):