Expose info_uri from query status in Cursor

Author: chyzzqo2

tests/integration/test_dbapi_integration.py

@@ -393,3 +393,14 @@ def test_eager_loading_cursor_description(trino_connection):
     assert description_after is not None
     assert len(description_after) == len(description_expected)
     assert all([a == e] for a, e in zip(description_after, description_expected))
+
+
+def test_info_uri(trino_connection):
+    cur = trino_connection.cursor()
+    assert cur.info_uri is None
+    cur.execute('SELECT * FROM system.runtime.nodes')
+    assert cur.info_uri is not None
+    assert cur._query.query_id in cur.info_uri
+    cur.fetchall()
+    assert cur.info_uri is not None
+    assert cur._query.query_id in cur.info_uri

trino/client.py

@@ -477,6 +477,7 @@ def __init__(
         self.query_id: Optional[str] = None
 
         self._stats: Dict[Any, Any] = {}
+        self._info_uri: Optional[str] = None
         self._warnings: List[Dict[Any, Any]] = []
         self._columns: Optional[List[str]] = None
         self._finished = False
@@ -507,6 +508,10 @@ def warnings(self):
     def result(self):
         return self._result
 
+    @property
+    def info_uri(self):
+        return self._info_uri
+
     def execute(self, additional_http_headers=None) -> TrinoResult:
         """Initiate a Trino query by sending the SQL statement
 
@@ -520,6 +525,7 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
 
         response = self._request.post(self._sql, additional_http_headers)
         status = self._request.process(response)
+        self._info_uri = status.info_uri
         self.query_id = status.id
         self._stats.update({"queryId": self.query_id})
         self._update_state(status)

trino/dbapi.py

@@ -201,6 +201,12 @@ def __iter__(self):
     def connection(self):
         return self._connection
 
+    @property
+    def info_uri(self):
+        if self._query is not None:
+            return self._query.info_uri
+        return None
+
     @property
     def description(self):
         if self._query.columns is None: