chore: Cleanup database sessions

[john-bodley]

SUMMARY

At Airbnb we've been running into issues with PyMySQL on Aurora MySQL 5.7. A Google search suggests the problem may be related to thread safety and possibly that Apache Superset may not be leveraging SQLAlchemy's scoped sessions correctly.

After grokking the code I realized that we were somewhat inconsistent with how we handle sessions with the metadata database. Ideally we should be leaning on a single instance of Flask-SQLAlchemy's scope session which is safely removed on teardown. I noticed examples of where we created additional new local sessions (rather than the global session) which were only sometimes closed and/or removed. Note locally scoped session refer to the same global scope session per here.

This PR simply ensures that all ORM interactions with the metadata database use the db.session construct where db is the Flask-SQLAlchemy object provided by FAB. Also it seems redundant and adds unnecessary obfuscation to pass around the db.session and thus a number of function signatures have been updated.

Note sadly this change did not remedy the issue we were seeing, and thus I wonder if the probably lies with a third party library or something else within Apache Superset. I've been able to successfully use the SQLAlchemy ORM using scoped sessions without issue via a toy example. That said it could help with any connection pool issues we sporadically see, i.e., using the global Flask-SQLAlchemy scoped session ensures that any existing transactional/connection resources will be released (which wasn't the case previously).

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TEST PLAN

CI.

ADDITIONAL INFORMATION

• Has associated issue:
• Changes UI
• Requires DB Migration.
• Confirm DB Migration upgrade and downgrade tested.
• Introduces new feature or API
• Removes existing feature or API

[codecov-commenter]

Codecov Report

Merging #10427 into master will increase coverage by 5.39%.
The diff coverage is n/a.

@@            Coverage Diff             @@
##           master   #10427      +/-   ##
==========================================
+ Coverage   65.98%   71.38%   +5.39%     
==========================================
  Files         604      408     -196     
  Lines       32376    13339   -19037     
  Branches     3282     3282              
==========================================
- Hits        21364     9522   -11842     
+ Misses      10824     3705    -7119     
+ Partials      188      112      -76     

Flag	Coverage Δ
#cypress	54.48% <ø> (?)
#javascript	59.72% <ø> (-0.01%)	⬇️
#python	?

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

Impacted Files	Coverage Δ
...src/components/FilterableTable/FilterableTable.tsx	81.77% <0.00%> (-0.41%)	⬇️
superset/views/database/mixins.py
superset/examples/paris.py
superset/charts/commands/update.py
superset/charts/commands/exceptions.py
superset/dashboards/api.py
superset/models/annotations.py
superset/dashboards/commands/exceptions.py
superset/datasets/commands/refresh.py
superset/db_engine_specs/athena.py
... and 315 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update edaf785...c4618ba. Read the comment docs.

[john-bodley]

Note the async SQL Lab queries are failing but only for MySQL and I'm not entirely sure why. I speculate it may be due to the query status update not being committed.

[villebro]

@john-bodley I have a similar problem on Vertica connections that I noticed recently. Not sure if it's related, was planning to look more closely at the problem but didn't yet have the time. I should be able to review this and the other PR in the coming days.

[john-bodley]

Note in terms of the failing tests I've reverted some of my changes in tests/celery_tests.py which removed some session.close() statements. Per the Session.close documentation,

This clears all items and ends any transaction in progress.

which may be necessary as some of the Celery tasks first try to drop a table and may throw an exception if the table exists. Ending the transaction may be necessary prior to running the async query.

[dpgaspar]

Excellent work @john-bodley, but I would prefer the db.session imports to be:

from superset.extensions import db

instead of

from superset import db

[john-bodley]

@dpgaspar agreed especially given circular import concerns. I'll make the change.

[john-bodley]

This is the same session per here and added unnecessary obfuscation.

[john-bodley]

@dpgaspar I've addressed your comments.

[john-bodley]

@dpgaspar I've addressed your comments.

[villebro]

Uhh, kudos for the cleanup effort..! I wish we had a linter that could ban passing a Session to avoid this building up post this PR. Just one comment, happy to approve if you don't see any risk with removing the scoped session.

[villebro]

This has me wondering, is there risk that this gets torn down prematurely by the initiating session's closure? My intuition says it shouldn't, but it seems strange that this would have been put here by accident.

[john-bodley]

@villebro I'm unsure though I believe the scoped session is actually the same instance as db.session per here, i.e.,

>>> Session = scoped_session(session_factory)
>>> some_session = Session()
>>> some_other_session = Session()
>>> some_session is some_other_session
True

[villebro]

I threw this into the codebase, and the result was False:

>>> scoped_session = db.create_scoped_session()
>>> print(scoped_session is db.session)
False

Perhaps we should leave these out for now and address them in a follow-up PR?

[john-bodley]

@villebro possibly, though the scoped session is merely running a query as far as I can tell. There’s no complex logic, transaction, etc. being used.

[villebro]

@john-bodley good point.

[villebro]

LGTM (needs a rebase).