fix(sqllab): Force trino client async execution

[giftig]

SUMMARY

We are currently unable to stop trino queries, because the underlying trino client blocks until the query completes, and doesn't make the query ID or any other info available in the meantime. Unfortunately it doesn't look like they plan to change that any time soon, either.

Make the following changes:

• Add a new method execute_with_cursor to db_engine_spec which combines execute with handle_cursor, factoring it out of the one place it's used, deep in the execute query logic.
• Make handle_cursor poll the cursor for query ID, as it's going to be populated asynchronously. Add warnings that the trino impl will require using execute_with_cursor. Currently nothing is directly calling handle_cursor, with the one original call eliminated.
• Override execute_with_cursor for the trino engine and execute the two tasks in parallel to allow us to poll for the query ID while the query is still blocking.

TESTING INSTRUCTIONS

Test that trino queries can be stopped in SQL Lab, that database errors from trino are still reported correctly, and that queries in other engines still run / stop correctly.

ADDITIONAL INFORMATION

• Has associated issue: Trino queries cannot be stopped in SQL Lab #24858
• Required feature flags:
• Changes UI
• Includes DB Migration (follow approval process in SIP-59)
  • Migration is atomic, supports rollback & is backwards-compatible
  • Confirm DB migration upgrade and downgrade tested
  • Runtime estimates and downtime expectations provided
• Introduces new feature or API
• Removes existing feature or API

[codecov]

Codecov Report

Merging #24859 (b87caf6) into master (7f9b038) will decrease coverage by 0.08%.
Report is 3 commits behind head on master.
The diff coverage is 83.72%.

❗ Current head b87caf6 differs from pull request most recent head 4919b58. Consider uploading reports for the commit 4919b58 to get more accurate results

@@            Coverage Diff             @@
##           master   #24859      +/-   ##
==========================================
- Coverage   69.00%   68.93%   -0.08%     
==========================================
  Files        1904     1904              
  Lines       74105    74136      +31     
  Branches     8193     8194       +1     
==========================================
- Hits        51137    51105      -32     
- Misses      20847    20912      +65     
+ Partials     2121     2119       -2     

Flag	Coverage Δ
hive	54.11% <32.43%> (-0.03%)	⬇️
mysql	79.17% <32.43%> (-0.05%)	⬇️
postgres	?
presto	54.01% <32.43%> (-0.03%)	⬇️
python	83.20% <83.78%> (-0.17%)	⬇️
sqlite	?
unit	55.02% <83.78%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files Changed	Coverage Δ
...ins/legacy-preset-chart-nvd3/src/transformProps.js	7.14% <0.00%> (-0.27%)	⬇️
superset/db_engine_specs/trino.py	82.55% <79.31%> (-1.19%)	⬇️
...dashboard/components/AddSliceCard/AddSliceCard.tsx	68.88% <100.00%> (+5.25%)	⬆️
superset/db_engine_specs/base.py	89.40% <100.00%> (-1.62%)	⬇️
superset/sql_lab.py	86.25% <100.00%> (-0.16%)	⬇️

... and 17 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[giftig]

Tests are passing with a small modification to the way the mock is used, and I've just pushed a fix for the small formatting issue black corrected in the tests I modified, so build should now pass.

I'm looking at adding a unit test for this scenario in test_trino.py but the test scenario here will be quite fiddly / timing-based because I need to call execute_with_cursor and then prove that query_id has been found and populated prior to that call returning, which means my test will need a threaded solution similar to what I've done in the actual code, and my execute mock call will have to wait for an arbitrary period of time; there's potential to slow down tests and introduce flakiness based on timing issues with that approach, so I'm not keen.

Existing test coverage covers trino's handle_cursor, along with the new call to handle_with_cursor in sql_lab, but there's no coverage for trino's handle_with_cursor specifically. Perhaps simpler coverage of the trino method will suffice; i.e. make sure that we at least make the expected calls on and fetch the appropriate query_id from the cursor eventually, and not worrying about the timing of that. That won't fully test the scenario I've fixed here, but will at least provide basic coverage of the method.

[giftig]

Added a basic coverage unit test as described above.

[john-bodley]

@nytai @villebro would you mind reviewing this? I know @villebro was working with the cancel query logic in the past and I suspect this was previously tested with Trino—though I may be wrong.

[dungdm93]

After discuss with @giftig and find out this issues is caused by
https://github.com/trinodb/trino-python-client/blob/d6047a80c1a772cb8240f4a81df810b0d48af4df/trino/client.py#L819-L821

I'm gonna ask Trino community about this to figure out it's should be fix in Trino or Superset side.

[giftig]

To share the feedback from the trino team via @dungdm93:

Because before this change lot of people didn't use the API correctly and assumed that DDL queries for example didn't need fetch methods to be called.
One option is to have a thread pool within the DB-API client itself on which the queries execute so that Cursor.execute can return eagerly but that's not something we have considered since no one has complained about the change until now.
Also it was an easy solution to fix a bug we had about not implementing protocol correctly - see trinodb/trino-python-client#220

So it doesn't sound like there are immediate plans to fix the issue in the client, though there's a suggestion how it might be done.

Currently the lack of that feature in the client causes a bug in superset so imo the way forward it to merge this PR to fix the issue and to follow up and try to get a more direct fix merged into the trino client. We'll then have to upgrade the trino client in superset, and my understanding from @villebro is we'd prefer to stay a couple of versions behind the latest, so we can assume some lead time to catch up with changes we make there. Once we've upgraded we can revert this workaround.

Thoughts @dungdm93 @villebro ?

[giftig]

Sorry to nag @villebro @dungdm93 but can we discuss any remaining concerns with this PR and try to get this moving? I understand there's an alternate fix being proposed by @dungdm93 in parallel, but I think we all agree that will take quite a while to be ready to go into master, and in the meantime this remains a fairly impactful bug. I'd like to get this workaround merged into master in the meantime, with the understanding it can be reverted / dropped from future releases once we've managed to upgrade to a trino client release containing @dungdm93 's fix.

Happy to address any comments which come up on this PR in the meantime, and talk through any reasoning / the testing I did as needed.

[villebro]

LGTM - we can revisit this once the Trino client adds non-blocking functionality.