Skip to content

Commit

Permalink
Expose info_uri from query status in Cursor
Browse files Browse the repository at this point in the history
  • Loading branch information
chyzzqo2 authored and hashhar committed Nov 19, 2021
1 parent 5cd6523 commit c8144de
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 0 deletions.
11 changes: 11 additions & 0 deletions tests/integration/test_dbapi_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 6 additions & 0 deletions trino/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down
6 changes: 6 additions & 0 deletions trino/dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit c8144de

Please sign in to comment.