fix(dbapi): Address timestamp with time zone localization issue

tests/integration/test_dbapi_integration.py

@@ -364,7 +364,7 @@ def test_datetime_with_numeric_offset_time_zone_query_param(trino_connection):
 def test_datetime_with_named_time_zone_query_param(trino_connection):
  cur = trino_connection.cursor()
 
- params = datetime(2020, 1, 1, 16, 43, 22, 320000, tzinfo=pytz.timezone('America/Los_Angeles'))
+ params = pytz.timezone('America/Los_Angeles').localize(datetime(2020, 1, 1, 16, 43, 22, 320000))
 
  cur.execute("SELECT ?", params=(params,))
  rows = cur.fetchall()
@@ -393,13 +393,44 @@ def test_null_datetime_with_time_zone(trino_connection):
  assert_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3)
 
 
-def test_datetime_with_time_zone_numeric_offset(trino_connection):
+@pytest.mark.parametrize(
+ 'operation,result,precision',
+ [
+ (
+ "SELECT TIMESTAMP '2001-08-22 03:04:05.321 -08:00'",
+ datetime(2001, 8, 22, 3, 4, 5, 321000, tzinfo=timezone(-timedelta(hours=8))),
+ 3,
+
+ ),
+ (
+ "SELECT TIMESTAMP '2001-08-22 03:04:05.321 America/Los_Angeles'",
+ pytz.timezone("America/Los_Angeles").localize(datetime(2001, 8, 22, 3, 4, 5, 321000)),
+ 3,
+ ),
+ ],
+)
+def test_datetime_with_time_zone(trino_connection, operation, result, precision):
+ cur = trino_connection.cursor()
+
+ cur.execute(operation)
+ rows = cur.fetchall()
+
+ assert rows[0][0] == result
+
+ assert_cursor_description(
+ cur,
+ trino_type=f"timestamp({precision}) with time zone",
+ precision=precision,
+ )
+
+
+def test_datetime_with_at_time_zone(trino_connection):
  cur = trino_connection.cursor()
 
- cur.execute("SELECT TIMESTAMP '2001-08-22 03:04:05.321 -08:00'")
+ cur.execute("SELECT TIMESTAMP '2001-08-22 03:04:05.321 UTC' AT TIME ZONE 'America/Los_Angeles'")
  rows = cur.fetchall()
 
- assert rows[0][0] == datetime.strptime("2001-08-22 03:04:05.321 -08:00", "%Y-%m-%d %H:%M:%S.%f %z")
+ assert rows[0][0] == pytz.timezone('America/Los_Angeles').localize(datetime(2001, 8, 21, 20, 4, 5, 321000))
  assert_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3)
 
 
@@ -416,23 +447,24 @@ def test_datetimes_with_time_zone_in_dst_gap_query_param(trino_connection):
 def test_doubled_datetimes(trino_connection):
  # Trino doesn't distinguish between doubled datetimes that lie within a DST transition. See also
  # See also https://github.com/trinodb/trino/issues/5781
+ dt = datetime(2002, 10, 27, 1, 30, 0)
  cur = trino_connection.cursor()
 
- params = pytz.timezone('US/Eastern').localize(datetime(2002, 10, 27, 1, 30, 0), is_dst=True)
+ params = pytz.timezone('US/Eastern').localize(dt, is_dst=True)
 
  cur.execute("SELECT ?", params=(params,))
  rows = cur.fetchall()
 
- assert rows[0][0] == datetime(2002, 10, 27, 1, 30, 0, tzinfo=pytz.timezone('US/Eastern'))
+ assert rows[0][0] == pytz.timezone('US/Eastern').localize(dt)
 
  cur = trino_connection.cursor()
 
- params = pytz.timezone('US/Eastern').localize(datetime(2002, 10, 27, 1, 30, 0), is_dst=False)
+ params = pytz.timezone('US/Eastern').localize(dt, is_dst=False)
 
  cur.execute("SELECT ?", params=(params,))
  rows = cur.fetchall()
 
- assert rows[0][0] == datetime(2002, 10, 27, 1, 30, 0, tzinfo=pytz.timezone('US/Eastern'))
+ assert rows[0][0] == pytz.timezone('US/Eastern').localize(dt)
 
 
 def test_date_query_param(trino_connection):

trino/client.py

@@ -1108,8 +1108,10 @@ def map(self, value) -> Optional[datetime]:
  datetime_with_fraction, timezone_part = value.rsplit(' ', 1)
  whole_python_temporal_value = datetime_with_fraction[:self.datetime_default_size]
  remaining_fractional_seconds = datetime_with_fraction[self.datetime_default_size + 1:]
+ tz = _create_tzinfo(timezone_part)
+ dt = datetime.fromisoformat(whole_python_temporal_value)
  return TimestampWithTimeZone(
- datetime.fromisoformat(whole_python_temporal_value).replace(tzinfo=_create_tzinfo(timezone_part)),
+ dt.replace(tzinfo=tz) if isinstance(tz, timezone) else tz.localize(dt),
  _fraction_to_decimal(remaining_fractional_seconds),
  ).round_to(self.precision).to_python_type()