Add Cursor.describe to retrieve the schema of a query

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -20,6 +20,7 @@
 import trino
 from tests.integration.conftest import trino_version
 from trino import constants
+from trino.dbapi import DescribeOutput
 from trino.exceptions import NotSupportedError, TrinoQueryError, TrinoUserError
 from trino.transaction import IsolationLevel
 
@@ -1107,3 +1108,69 @@ def test_prepared_statements(run_trino):
     cur.execute('DEALLOCATE PREPARE test_prepared_statements')
     cur.fetchall()
     assert cur._request._client_session.prepared_statements == {}
+
+
+def test_describe(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch",
+    )
+    cur = trino_connection.cursor()
+
+    result = cur.describe("SELECT 1, DECIMAL '1.0' as a")
+
+    assert result == [
+        DescribeOutput(name='_col0', catalog='', schema='', table='', type='integer', type_size=4, aliased=False),
+        DescribeOutput(name='a', catalog='', schema='', table='', type='decimal(2,1)', type_size=8, aliased=True)
+    ]
+
+
+def test_describe_table_query(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", catalog="tpch",
+    )
+    cur = trino_connection.cursor()
+
+    result = cur.describe("SELECT * from tpch.tiny.nation")
+
+    assert result == [
+        DescribeOutput(
+            name='nationkey',
+            catalog='tpch',
+            schema='tiny',
+            table='nation',
+            type='bigint',
+            type_size=8,
+            aliased=False,
+        ),
+        DescribeOutput(
+            name='name',
+            catalog='tpch',
+            schema='tiny',
+            table='nation',
+            type='varchar(25)',
+            type_size=0,
+            aliased=False,
+        ),
+        DescribeOutput(
+            name='regionkey',
+            catalog='tpch',
+            schema='tiny',
+            table='nation',
+            type='bigint',
+            type_size=8,
+            aliased=False,
+        ),
+        DescribeOutput(
+            name='comment',
+            catalog='tpch',
+            schema='tiny',
+            table='nation',
+            type='varchar(152)',
+            type_size=0,
+            aliased=False,
+        )
+    ]

trino/dbapi.py

@@ -21,7 +21,7 @@
 import math
 import uuid
 from decimal import Decimal
-from typing import Any, Dict, List, Optional  # NOQA for mypy types
+from typing import Any, Dict, List, NamedTuple, Optional  # NOQA for mypy types
 
 import trino.client
 import trino.exceptions
@@ -220,6 +220,20 @@ def cursor(self, experimental_python_types: bool = None):
         )
 
 
+class DescribeOutput(NamedTuple):
+    name: str
+    catalog: str
+    schema: str
+    table: str
+    type: str
+    type_size: int
+    aliased: bool
+
+    @classmethod
+    def from_row(cls, row: List[Any]):
+        return cls(*row)
+
+
 class Cursor(object):
     """Database cursor.
 
@@ -517,6 +531,27 @@ def fetchmany(self, size=None) -> List[List[Any]]:
 
         return result
 
+    def describe(self, sql: str) -> List[DescribeOutput]:
+        """
+        Returns the schema of a given SQL statement
+
+        :param sql: SQL statement
+        """
+        statement_name = self._generate_unique_statement_name()
+        self._prepare_statement(sql, statement_name)
+        try:
+            sql = f"DESCRIBE OUTPUT {statement_name}"
+            self._query = trino.client.TrinoQuery(
+                self._request,
+                sql=sql,
+                experimental_python_types=self._experimental_pyton_types,
+            )
+            result = self._query.execute()
+        finally:
+            self._deallocate_prepared_statement(statement_name)
+
+        return list(map(lambda x: DescribeOutput.from_row(x), result))
+
     def genall(self):
         return self._query.result