fix(trino): db session error in handle cursor #31265
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #31265 +/- ##
===========================================
+ Coverage 60.48% 83.80% +23.31%
===========================================
Files 1931 536 -1395
Lines 76236 38943 -37293
Branches 8568 0 -8568
===========================================
- Hits 46114 32637 -13477
+ Misses 28017 6306 -21711
+ Partials 2105 0 -2105
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
superset/db_engine_specs/trino.py
Outdated
| @@ -216,6 +216,8 @@ def handle_cursor(cls, cursor: Cursor, query: Query) -> None: | |||
| if tracking_url := cls.get_tracking_url(cursor): | |||
| query.tracking_url = tracking_url | |||
|
|
|||
| db.session.commit() | |||
There was a problem hiding this comment.
Maybe move the global override to the line?
| db.session.commit() | |
| db.session.commit() # pylint: disable=consider-using-transaction |
| @@ -283,6 +286,8 @@ def _execute( | |||
| ) | |||
| execute_thread.start() | |||
|
|
|||
| # Wait for the thread to start before continuing | |||
| time.sleep(0.1) | |||
There was a problem hiding this comment.
@villebro @betodealmeida this is being added to make the thread run once but I think it's not guaranteed. Do you think we could change the initial condition while not cursor.query_id to avoid having 2 sleeps?
There was a problem hiding this comment.
@michael-s-molina I don't think that would work, as .set() is only called once the thread is completed. Should we perhaps be checking that execute_thread.is_alive() is returning True before proceeding? I think that's what we're trying to do here.
There was a problem hiding this comment.
@justinpark Could you try @villebro's suggestion?
(cherry picked from commit 1e0c04f)
SUMMARY
In the thread execution function, the query is referenced, and since it occupies the session within the thread scope, it became impossible to perform session transactions related to the query in the handle_cursor that runs separately.
This resulted in errors for multi-statement queries. This commit changes the code to ensure that the database object is only referenced in the thread, thereby separating the session scope.
Additionally, it resolves the issue where handle_cursor would skip sleep when a multi-statement is executing, allowing access before execution.
BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
No error thrown
TESTING INSTRUCTIONS
Run a multi-statement query in trino db, such like
select 1; select 2;ADDITIONAL INFORMATION