Make names of ROW datatype available in result set

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -745,6 +745,49 @@ def test_named_row(trino_connection):
     rows = cur.fetchall()
 
     assert rows[0][0] == (1, 2.0)
+    assert rows[0][0][0] == 1
+    assert rows[0][0][1] == 2.0
+    assert rows[0][0].x == 1
+    assert rows[0][0].y == 2.0
+
+    assert rows[0][0].__annotations__["names"] == ['x', 'y']
+    assert rows[0][0].__annotations__["types"] == ['bigint', 'double']
+
+
+def test_named_row_duplicate_names(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute("SELECT CAST(ROW(1, 2e0) AS ROW(x BIGINT, x DOUBLE))")
+    rows = cur.fetchall()
+
+    assert rows[0][0] == (1, 2.0)
+    with pytest.raises(ValueError, match="Ambiguous row field reference: x"):
+        rows[0][0].x
+
+    assert rows[0][0].__annotations__["names"] == ['x', 'x']
+    assert rows[0][0].__annotations__["types"] == ['bigint', 'double']
+    assert str(rows[0][0]) == "(1, 2.0)"
+
+
+def test_nested_named_row(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute("SELECT CAST(ROW(DECIMAL '2.3', ROW(1, 'test')) AS ROW(x DECIMAL(3,2), y ROW(x BIGINT, y VARCHAR)))")
+    rows = cur.fetchall()
+
+    assert rows[0][0] == (Decimal('2.3'), (1, 'test'))
+    assert rows[0][0][0] == Decimal('2.3')
+    assert rows[0][0][1] == (1, 'test')
+    assert rows[0][0][1][0] == 1
+    assert rows[0][0][1][1] == 'test'
+    assert rows[0][0].x == Decimal('2.3')
+    assert rows[0][0].y.x == 1
+    assert rows[0][0].y.y == 'test'
+
+    assert rows[0][0].__annotations__["names"] == ['x', 'y']
+    assert rows[0][0].__annotations__["types"] == ['decimal', 'row']
+
+    assert rows[0][0].y.__annotations__["names"] == ['x', 'y']
+    assert rows[0][0].y.__annotations__["types"] == ['bigint', 'varchar']
+    assert str(rows[0][0]) == "(x: Decimal('2.30'), y: (x: 1, y: 'test'))"
 
 
 def test_float_query_param(trino_connection):

trino/client.py

@@ -1128,14 +1128,48 @@ def map(self, values: List[Any]) -> Optional[List[Any]]:
         return [self.mapper.map(value) for value in values]
 
 
+class NamedRowTuple(tuple):
+    """Custom tuple class as namedtuple doesn't support missing or duplicate names"""
+    def __new__(cls, values, names: List[str], types: List[str]):
+        return super().__new__(cls, values)
+
+    def __init__(self, values, names: List[str], types: List[str]):
+        self._names = names
+        # With names and types users can retrieve the name and Trino data type of a row
+        self.__annotations__ = dict()
+        self.__annotations__["names"] = names
+        self.__annotations__["types"] = types
+        elements: List[Any] = []
+        for name, value in zip(names, values):
+            if names.count(name) == 1:
+                setattr(self, name, value)
+                elements.append(f"{name}: {repr(value)}")
+            else:
+                elements.append(repr(value))
+        self._repr = "(" + ", ".join(elements) + ")"
+
+    def __getattr__(self, name):
+        if self._names.count(name):
+            raise ValueError("Ambiguous row field reference: " + name)
+
+    def __repr__(self):
+        return self._repr
+
+
 class RowValueMapper(ValueMapper[Tuple[Optional[Any], ...]]):
-    def __init__(self, mappers: List[ValueMapper[Any]]):
+    def __init__(self, mappers: List[ValueMapper[Any]], names: List[str], types: List[str]):
         self.mappers = mappers
+        self.names = names
+        self.types = types
 
     def map(self, values: List[Any]) -> Optional[Tuple[Optional[Any], ...]]:
         if values is None:
             return None
-        return tuple(self.mappers[index].map(value) for index, value in enumerate(values))
+        return NamedRowTuple(
+            list(self.mappers[index].map(value) for index, value in enumerate(values)),
+            self.names,
+            self.types
+        )
 
 
 class MapValueMapper(ValueMapper[Dict[Any, Optional[Any]]]):
@@ -1183,8 +1217,14 @@ def _create_value_mapper(self, column) -> ValueMapper:
             value_mapper = self._create_value_mapper(column['arguments'][0]['value'])
             return ArrayValueMapper(value_mapper)
         elif col_type == 'row':
-            mappers = [self._create_value_mapper(arg['value']['typeSignature']) for arg in column['arguments']]
-            return RowValueMapper(mappers)
+            mappers = []
+            names = []
+            types = []
+            for arg in column['arguments']:
+                mappers.append(self._create_value_mapper(arg['value']['typeSignature']))
+                names.append(arg['value']['fieldName']['name'] if "fieldName" in arg['value'] else None)
+                types.append(arg['value']['typeSignature']['rawType'])
+            return RowValueMapper(mappers, names, types)
         elif col_type == 'map':
             key_mapper = self._create_value_mapper(column['arguments'][0]['value'])
             value_mapper = self._create_value_mapper(column['arguments'][1]['value'])