Skip to content

[Trino] Adding explain query method #3645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 15, 2024
Merged

[Trino] Adding explain query method #3645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions desktop/libs/notebook/src/notebook/connectors/trino.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -321,6 +321,26 @@ def _get_columns(self, database, table):
]


@query_error_handler
def explain(self, notebook, snippet):
statement = snippet['statement'].rstrip(';')
explanation = ''

if statement:
try:
TrinoQuery(self.trino_request, 'USE ' + snippet['database']).execute()
result = TrinoQuery(self.trino_request, 'EXPLAIN ' + statement).execute()
explanation = result.rows
except Exception as e:
explanation = str(e)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding error in explanation?

Copy link
Collaborator Author

@agl29 agl29 Mar 15, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that's what we are doing in the Hive also
https://github.com/cloudera/hue/blob/master/desktop/libs/notebook/src/notebook/connectors/hiveserver2.py#L608

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JohanAhlen Is UI expecting this for both editor v1 and editor v2?
Just a heads up, we need to improve all this flow then once we pick up this roadmap item. cc @agl29

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's a bit odd indeed, @agl29 did you try how the UI behave if you would return a 500 request with the error message instead?


return {
'status': 0,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will the TrinoQuery.execute() throws any exception?
If any exception, should the status be updated to -1?
Take a look existing beeswax API examples: response = {'status': -1, 'message': ''} are defined at the beginning of menthod.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the try catch.
We need to send the {'status': 0} to capture the output in UI.
For the ref https://github.com/cloudera/hue/blob/master/desktop/libs/notebook/src/notebook/connectors/hiveserver2.py#L591

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@wing2fly we should try to move away from using the status flag, today API users expect the proper http error code when something goes bad. That being said the UI might have some old code that checks for status == 0.

'explanation': explanation,
'statement': statement
}


def download(self, notebook, snippet, file_format='csv'):
result_wrapper = TrinoExecutionWrapper(self, notebook, snippet)

Expand Down
42 changes: 42 additions & 0 deletions desktop/libs/notebook/src/notebook/connectors/trino_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,3 +253,45 @@ def test_get_select_query(self):
self.trino_api._get_select_query(database, table),
expected_statement
)


def test_explain(self):
with patch('notebook.connectors.trino.TrinoQuery') as TrinoQuery:
snippet = {'statement': 'SELECT * FROM tpch.sf1.partsupp LIMIT 100;', 'database': 'tpch.sf1'}
output = [['Trino version: 432\nFragment 0 [SINGLE]\n Output layout: [partkey, suppkey, availqty, supplycost, comment]\n '\
'Output partitioning: SINGLE []\n Output[columnNames = [partkey, suppkey, availqty, supplycost, comment]]\n │ '\
'Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, comment:varchar(199)]\n │ '\
'Estimates: {rows: 100 (15.67kB), cpu: 0, memory: 0B, network: 0B}\n └─ Limit[count = 100]\n │ '\
'Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, comment:varchar(199)]\n │ '\
'Estimates: {rows: 100 (15.67kB), cpu: 15.67k, memory: 0B, network: 0B}\n └─ LocalExchange[partitioning = SINGLE]\n '\
'│ Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, comment:varchar(199)]\n │ '\
'Estimates: {rows: 100 (15.67kB), cpu: 0, memory: 0B, network: 0B}\n └─ RemoteSource[sourceFragmentIds = [1]]\n'\
' Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, comment:varchar(199)]\n\n'\
'Fragment 1 [SOURCE]\n Output layout: [partkey, suppkey, availqty, supplycost, comment]\n Output partitioning: SINGLE []\n'\
' LimitPartial[count = 100]\n │ Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, '\
'comment:varchar(199)]\n │ Estimates: {rows: 100 (15.67kB), cpu: 15.67k, memory: 0B, network: 0B}\n └─ '\
'TableScan[table = tpch:sf1:partsupp]\n Layout: [partkey:bigint, suppkey:bigint, availqty:integer, supplycost:double, '\
'comment:varchar(199)]\n Estimates: {rows: 800000 (122.44MB), cpu: 122.44M, memory: 0B, network: 0B}\n '\
'partkey := tpch:partkey\n availqty := tpch:availqty\n supplycost := tpch:supplycost\n '\
'comment := tpch:comment\n suppkey := tpch:suppkey\n\n']]
# Mock TrinoQuery object and its execute method
query_instance = TrinoQuery.return_value
query_instance.execute.return_value = MagicMock(next_uri=None, id='123', rows=output, columns=[])

# Call the explain method
result = self.trino_api.explain(notebook=None, snippet=snippet)

# Assert the result
assert_equal(result['status'], 0)
assert_equal(result['explanation'], output)
assert_equal(result['statement'], 'SELECT * FROM tpch.sf1.partsupp LIMIT 100')

query_instance = TrinoQuery.return_value
query_instance.execute.side_effect = Exception('Mocked exception')

# Call the explain method
result = self.trino_api.explain(notebook=None, snippet=snippet)

# Assert the exception message
assert_equal(result['explanation'], 'Mocked exception')