fix: Ensure SQLAlchemy sessions are closed

[john-bodley]

SUMMARY

At Airbnb we've been running into the infamous SQLAlchemy QueuePool issue which typically occurs when connections aren't being closed. Thankfully Flask-SQLAlchemy normally handles this for us, per here, when the app context is torn down.

Grokking the code it seems like there are instances (for right or wrong) where we create our own sessions which aren't then explicitly closed. This PR ensures that either i) these sessions are now closed, or ii) the Flask-SQLAlchemy session is used.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

CI.

ADDITIONAL INFORMATION

• Has associated issue:
• 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

[john-bodley]

I gather these standalone scoped sessions exist because tasks fall outside the jurisdiction of Flask. If this isn't the case then—to adhere to the KISS priniciple—we should simply reuse the Flask-SQLAlchemy session.

[villebro]

LGTM. However, at some point we might want to discuss if we may want to make it possible to not commit in the DAOs or similar methods. I've been writing some custom functionality in our security manager that wants to perform transaction-like operations, and I'd prefer to either have the whole chain of events succeed, or rollback everything. An option could be to add a commit=True parameter to these methods, or then just remove the explicit commit and assume callers do it (e.g. in the api.py) after performing the subtasks.

[john-bodley]

@villebro I was talking with @michael-s-molina earlier this week and plan to write a SIP about DAOs. The plan is—as you pointed out—that one shouldn't commit within a DAO but rather on a unit of work (typically a Flask request, though in our case it's likely a command). I hope to have the SIP drafted earlier next week which will also mention SQLAlchemy nested transactions, which as the name suggests, allows for one to nest operations (leveraging SAVEPOINTS) which ensures one of two states, the transaction is committed or rolled back.

This removes the need to pass around commit = True | False to the various methods. Here's a taste of what it looks like.

[villebro]

@john-bodley this sounds amazing - I really look forward to this reform! ❤️

[michael-s-molina]

Shouldn't we use the connection parameter here? According to the event definition:

connection – the Connection being used to emit INSERT statements for this instance. This provides a handle into the current transaction on the target database specific to this instance.

[john-bodley]

@michael-s-molina I've reverted the PR back to my initial implementation which leverages the connection parameter.

It states that it provides a handle into the current transaction, yet we're likely over committing, i.e., there should only be one commit per unit of work. I'm not sure what happens if you close the session without committing. Maybe we using the connection wrongly.

[john-bodley]

Oh it seems we should be using SQLAlchemy's object_session rather than creating a new one. This should remove the need to commit and explicitly close the session as this is handled by Flask-SQLAlchemy.

[john-bodley]

Regrettably I was running into this issue and had to revert my last commit. I think we likely will have to revisit how tagging is working at a later date.

[john-bodley]

It's unclear why a new session was instantiated (and never closed) here. On line #433 the Flask-SQLAlchemy session is used and thus it seems prudent (and hopefully) safe to use db.session.

[zephyring]

I wonder if we can just reuse the db.session everywhere so we don't create this arbitrary session

[zephyring]

if we really need to use this session, better to wrap into try/catch and close at final otherwise exceptions thrown in commit will skip session.close

[john-bodley]

@zephyring that is the ideal solution, using the same session (which is obtainable via inspect(target).session) however per #25031 (comment) this is currently problematic given the way our SQLAlchemy event listener callbacks are configured.

[michael-s-molina]

@john-bodley Could you apply the try/finally pattern to the places where you're closing the session? I also think we can remove the session.commit() on line 91.

[john-bodley]

BTW @michael-s-molina per our offline conversation you can see here that the session is rolled back before closing.

[eschutho]

This is great. We were just investigating this issue that we originally thought was a problem with ssh connections, but later found that that the root cause was data dbs hitting the max connections. cc @hughhhh