Enhance Cursor.description

tests/integration/test_dbapi_integration.py

@@ -119,6 +119,7 @@ def test_none_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="unknown")
 
 
 def test_string_query_param(trino_connection):
@@ -128,6 +129,7 @@ def test_string_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == "six'"
+ assert_cursor_description(cur, trino_type="varchar(4)", size=4)
 
 
 def test_execute_many(trino_connection):
@@ -241,10 +243,11 @@ def test_legacy_primitive_types_with_connection_and_cursor(
 def test_decimal_query_param(trino_connection):
  cur = trino_connection.cursor()
 
- cur.execute("SELECT ?", params=(Decimal('0.142857'),))
+ cur.execute("SELECT ?", params=(Decimal('1112.142857'),))
  rows = cur.fetchall()
 
- assert rows[0][0] == Decimal('0.142857')
+ assert rows[0][0] == Decimal('1112.142857')
+ assert_cursor_description(cur, trino_type="decimal(10, 6)", precision=10, scale=6)
 
 
 def test_null_decimal(trino_connection):
@@ -254,6 +257,7 @@ def test_null_decimal(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="decimal(38, 0)", precision=38, scale=0)
 
 
 def test_biggest_decimal(trino_connection):
@@ -264,6 +268,7 @@ def test_biggest_decimal(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
+ assert_cursor_description(cur, trino_type="decimal(38, 0)", precision=38, scale=0)
 
 
 def test_smallest_decimal(trino_connection):
@@ -274,6 +279,7 @@ def test_smallest_decimal(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
+ assert_cursor_description(cur, trino_type="decimal(38, 0)", precision=38, scale=0)
 
 
 def test_highest_precision_decimal(trino_connection):
@@ -284,6 +290,7 @@ def test_highest_precision_decimal(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
+ assert_cursor_description(cur, trino_type="decimal(38, 38)", precision=38, scale=38)
 
 
 def test_datetime_query_param(trino_connection):
@@ -295,7 +302,7 @@ def test_datetime_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
- assert cur.description[0][1] == "timestamp(6)"
+ assert_cursor_description(cur, trino_type="timestamp(6)", precision=6)
 
 
 def test_datetime_with_utc_time_zone_query_param(trino_connection):
@@ -307,7 +314,7 @@ def test_datetime_with_utc_time_zone_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
- assert cur.description[0][1] == "timestamp(6) with time zone"
+ assert_cursor_description(cur, trino_type="timestamp(6) with time zone", precision=6)
 
 
 def test_datetime_with_numeric_offset_time_zone_query_param(trino_connection):
@@ -321,7 +328,7 @@ def test_datetime_with_numeric_offset_time_zone_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
- assert cur.description[0][1] == "timestamp(6) with time zone"
+ assert_cursor_description(cur, trino_type="timestamp(6) with time zone", precision=6)
 
 
 def test_datetime_with_named_time_zone_query_param(trino_connection):
@@ -333,7 +340,7 @@ def test_datetime_with_named_time_zone_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
- assert cur.description[0][1] == "timestamp(6) with time zone"
+ assert_cursor_description(cur, trino_type="timestamp(6) with time zone", precision=6)
 
 
 def test_datetime_with_trailing_zeros(trino_connection):
@@ -343,6 +350,7 @@ def test_datetime_with_trailing_zeros(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == datetime.strptime("2001-08-22 03:04:05.321000", "%Y-%m-%d %H:%M:%S.%f")
+ assert_cursor_description(cur, trino_type="timestamp(6)", precision=6)
 
 
 def test_null_datetime_with_time_zone(trino_connection):
@@ -352,6 +360,7 @@ def test_null_datetime_with_time_zone(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3)
 
 
 def test_datetime_with_time_zone_numeric_offset(trino_connection):
@@ -361,6 +370,7 @@ def test_datetime_with_time_zone_numeric_offset(trino_connection):
  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_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3)
 
 
 def test_datetimes_with_time_zone_in_dst_gap_query_param(trino_connection):
@@ -404,6 +414,7 @@ def test_date_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
+ assert_cursor_description(cur, trino_type="date")
 
 
 def test_null_date(trino_connection):
@@ -413,6 +424,7 @@ def test_null_date(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="date")
 
 
 def test_unsupported_python_dates(trino_connection):
@@ -462,6 +474,16 @@ def test_supported_special_dates_query_param(trino_connection):
  assert rows[0][0] == params
 
 
+def test_char(trino_connection):
+ cur = trino_connection.cursor()
+
+ cur.execute("SELECT CHAR 'trino'")
+ rows = cur.fetchall()
+
+ assert rows[0][0] == 'trino'
+ assert_cursor_description(cur, trino_type="char(5)", size=5)
+
+
 def test_time_query_param(trino_connection):
  cur = trino_connection.cursor()
 
@@ -471,7 +493,7 @@ def test_time_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == params
- assert cur.description[0][1] == "time(6)"
+ assert_cursor_description(cur, trino_type="time(6)", precision=6)
 
 
 def test_time_with_named_time_zone_query_param(trino_connection):
@@ -501,7 +523,7 @@ def test_time(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == time(1, 2, 3, 456000)
- assert cur.description[0][1] == "time(3)"
+ assert_cursor_description(cur, trino_type="time(3)", precision=3)
 
 
 def test_null_time(trino_connection):
@@ -511,6 +533,7 @@ def test_null_time(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="time(3)", precision=3)
 
 
 def test_time_with_time_zone_negative_offset(trino_connection):
@@ -522,7 +545,7 @@ def test_time_with_time_zone_negative_offset(trino_connection):
  tz = timezone(-timedelta(hours=8, minutes=0))
 
  assert rows[0][0] == time(1, 2, 3, 456000, tzinfo=tz)
- assert cur.description[0][1] == "time(3) with time zone"
+ assert_cursor_description(cur, trino_type="time(3) with time zone", precision=3)
 
 
 def test_time_with_time_zone_positive_offset(trino_connection):
@@ -534,7 +557,7 @@ def test_time_with_time_zone_positive_offset(trino_connection):
  tz = timezone(timedelta(hours=8, minutes=0))
 
  assert rows[0][0] == time(1, 2, 3, 456000, tzinfo=tz)
- assert cur.description[0][1] == "time(3) with time zone"
+ assert_cursor_description(cur, trino_type="time(3) with time zone", precision=3)
 
 
 def test_null_date_with_time_zone(trino_connection):
@@ -544,6 +567,7 @@ def test_null_date_with_time_zone(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] is None
+ assert_cursor_description(cur, trino_type="time(3) with time zone", precision=3)
 
 
 @pytest.mark.parametrize(
@@ -717,7 +741,7 @@ def test_float_query_param(trino_connection):
  cur.execute("SELECT ?", params=(1.1,))
  rows = cur.fetchall()
 
- assert cur.description[0][1] == "double"
+ assert_cursor_description(cur, trino_type="double")
  assert rows[0][0] == 1.1
 
 
@@ -726,7 +750,7 @@ def test_float_nan_query_param(trino_connection):
  cur.execute("SELECT ?", params=(float("nan"),))
  rows = cur.fetchall()
 
- assert cur.description[0][1] == "double"
+ assert_cursor_description(cur, trino_type="double")
  assert isinstance(rows[0][0], float)
  assert math.isnan(rows[0][0])
 
@@ -736,6 +760,7 @@ def test_float_inf_query_param(trino_connection):
  cur.execute("SELECT ?", params=(float("inf"),))
  rows = cur.fetchall()
 
+ assert_cursor_description(cur, trino_type="double")
  assert rows[0][0] == float("inf")
 
  cur.execute("SELECT ?", params=(float("-inf"),))
@@ -750,13 +775,13 @@ def test_int_query_param(trino_connection):
  rows = cur.fetchall()
 
  assert rows[0][0] == 3
- assert cur.description[0][1] == "integer"
+ assert_cursor_description(cur, trino_type="integer")
 
  cur.execute("SELECT ?", params=(9223372036854775807,))
  rows = cur.fetchall()
 
  assert rows[0][0] == 9223372036854775807
- assert cur.description[0][1] == "bigint"
+ assert_cursor_description(cur, trino_type="bigint")
 
 
 @pytest.mark.parametrize('params', [
@@ -1239,3 +1264,12 @@ def test_describe_table_query(run_trino):
  aliased=False,
  )
  ]
+
+
+def assert_cursor_description(cur, trino_type, size=None, precision=None, scale=None):
+ assert cur.description[0][1] == trino_type
+ assert cur.description[0][2] is None
+ assert cur.description[0][3] is size
+ assert cur.description[0][4] is precision
+ assert cur.description[0][5] is scale
+ assert cur.description[0][6] is None

trino/constants.py

@@ -53,3 +53,7 @@
 HEADER_SET_CATALOG = "X-Trino-Set-Catalog"
 
 HEADER_CLIENT_CAPABILITIES = "X-Trino-Client-Capabilities"
+
+LENGTH_TYPES = ["char", "varchar"]
+PRECISION_TYPES = ["time", "time with time zone", "timestamp", "timestamp with time zone", "decimal"]
+SCALE_TYPES = ["decimal"]

trino/dbapi.py

@@ -28,6 +28,7 @@
 import trino.exceptions
 import trino.logging
 from trino import constants
+from trino.constants import LENGTH_TYPES, PRECISION_TYPES, SCALE_TYPES
 from trino.exceptions import (
  DatabaseError,
  DataError,
@@ -237,6 +238,31 @@ def from_row(cls, row: List[Any]):
  return cls(*row)
 
 
+class ColumnDescription(NamedTuple):
+ name: str
+ type_code: int
+ display_size: int
+ internal_size: int
+ precision: int
+ scale: int
+ null_ok: bool
+
+ @classmethod
+ def from_column(cls, column: Dict[str, Any]):
+ type_signature = column["typeSignature"]
+ raw_type = type_signature["rawType"]
+ arguments = type_signature["arguments"]
+ return cls(
+ column["name"], # name
+ column["type"], # type_code
+ None, # display_size
+ arguments[0]["value"] if raw_type in LENGTH_TYPES else None, # internal_size
+ arguments[0]["value"] if raw_type in PRECISION_TYPES else None, # precision
+ arguments[1]["value"] if raw_type in SCALE_TYPES else None, # scale
+ None # null_ok
+ )
+
+
 class Cursor(object):
  """Database cursor.
 
@@ -278,14 +304,13 @@ def update_type(self):
  return None
 
  @property
- def description(self):
+ def description(self) -> List[ColumnDescription]:
  if self._query.columns is None:
  return None
 
  # [ (name, type_code, display_size, internal_size, precision, scale, null_ok) ]
  return [
- (col["name"], col["type"], None, None, None, None, None)
- for col in self._query.columns
+ ColumnDescription.from_column(col) for col in self._query.columns
  ]
 
  @property