[SIP-99B] Proposal for (re)defining a "unit of work"

[john-bodley]

[SIP-99B] Proposal for (re)defining a "unit of work"

This SIP is part of the [SIP-99] Proposal for correctly handling business logic series. Specifically, it proposes formalizing the "unit of work" construct. The topics outlined here are highly coupled with [SIP-99A] Primer for managing SQLAlchemy sessions as they both relate to managing database transactions.

Unit of Work

The "unit of work" is used to group multiple service operations into a single atomic unit. It materializes as a database transaction which—with the support of a SQLAlchemy session—allows us to create a block of code within which database atomicity is guaranteed. If the block of code is successfully completed, the changes are committed to the database. If there is an exception, the changes are rolled back.

SQLAlchemy also supports this feature:

The Unit Of Work system, a central part of SQLAlchemy's Object Relational Mapper (ORM), organizes pending INSERT/UPDATE/DELETE operations into queues and flushes them all in one batch.

Historically, Superset’s "unit of work" has been ill-defined, mismanaged, and/or misconstrued. It has resulted in us over (partial) committing*—a somewhat typical SQLAlchemy mistake (see SIP-99A for more details)—which violates the atomicity of the operation and leads to unnecessary complexity.

* The atomic unit vagueness has led to code inconsistencies which have resulted in us often adopting the “when in doubt, commit” mentality.

[SIP-35] Proposal for Improving Superset’s Python Code Organization introduced the concept of a Commands and Data Access Objects (DAOs) which (clearly and rightfully) stated,

Commands perform a single cohesive business operation, have a very small public API, and either entirely succeed or entirely fail. In practice, this requires that database transaction behavior be controlled at the Command level. When possible, commands should be modeled such that they perform a single database transaction. In the case of failure, they raise an exception or return an error code, and it is the responsibility of the caller to translate that into a correct response for a user.

however, sadly the code examples (where the database logic is defined entirely within the DAO which—via the commit and rollback operations—ends the transaction) likely led people astray—violating the "unit of work" concept.

Examples

The following examples illustrate where the atomic business operation has been violated:

• The BaseDAO.create() method auto commits by default, which means that the DAO, as opposed to the Command, acts as the “unit of work”.
• The QueryDAO.stop_query() method over commits.
• The ObjectUpdater.after_insert() SQLAlchemy event handler instantiates a session (and thus transaction) outside of the Flask-SQLAlchemy session and explicitly commits*.
• The UpdateDashboardCommand.run() method prevents command chaining due to the explicit commit.
• The EmbeddedDashboardDAO.upsert() method explicitly commits.

* The reason for this behavior is, per the SQLAlchemy event documentation,

Mapper-level flush events only allow very limited operations, on attributes local to the row being operated upon only, as well as allowing any SQL to be emitted on the given Connection.

The following examples illustrate where preserving the atomic business operation has added complexity or inconsistency which is difficult to grok:

• The DatasetDao.update() method has complex commit chaining.
• The DashboardDAO.update_chart_owners() and DashboardDAO.set_dash_metadata() methods are inconsistent, i.e., (by default) the former and latter commit and do not commit respectively—which leads to ambiguity within the same DAO in terms of the meaning of set and update.

Proposed Change

Unit of Work

In the context of Flask a typical "unit of work" is a request, however in addition to the RESTful API, Superset provides a CLI, and thus not all operations occur within the confines of a Flask request.

Per SIP-35 Commands really are the best representation of a “unit of work” as they provide a single cohesive business operation. Note not all RESTful API endpoints interface with a Command and thus, under these circumstances, the RESTful API can serve as the de facto atomic unit.

This proposal is inline with other recommendations that state one should

... decouple your business layer from any DAO technologies.

and

... keep the control of the transaction in the business layer because [the] transaction is a business problem.

Nested Transactions

Using a Command (as opposed to a DAO) as the "unit of work" does not mitigate the over commit problem completely as it is conceivable that a business transaction may not be encapsulated by a single Command.

To address this—without adding code complexity—we recommend leveraging SQLAlchemy's nested transaction which is both viable as all our metadata engines—MySQL, PostgreSQL, and SQLite— support the SAVEPOINT construct and conducive to our design pattern.

Used in conjunction with a context manager, upon exit, the outermost returned transaction object is either committed or rolled back to the SAVEPOINT. See here for the pseudo implementation.

The merits of this approach are:

1. Ensures that the "unit of work" is indeed a single atomic unit.
2. Does not require an explicit Session.flush(), Session.commit(), or Session.rollback() as this is provided by the context manager.
3. Does not require a priori knowledge of what constitutes a "unit of work" which is highly context-specific, i.e., the Command can be somewhat agnostic with regards to which nested level it represents—as long as it adheres to (2).

Event Handlers

Apart from asynchronous Celery tasks, only the Flask-SQLAlchemy singleton session should be used. Event handlers should share the same session of the mapped instance being persisted, i.e.,

@event.listens_for(Model, "after_insert")
def after_insert(mapper: Mapper, connection: Connection, target: Model) -> None:
    session = inspect(target).session
    ...

SQLAlchemy event handlers typically either call Connection.execute() (which auto-commit) or instantiate a new session (and corresponding transaction)—both of which violate the atomic unit construct. Invoking additional database operations can be problematic, i.e., the transaction may be closed.

In general we should discourage the use of the SQLAlchemy event handlers due to their complexity. Furthermore SQLAlchemy's bulk operations do not dispatch the corresponding ORM event callbacks which historically are required to augment additional records loosely defined via implicit relationships.

The proposed resolution would be to migrating the various operations—which typically represent business logic—to either the database (if appropriate) or Command level.

Testing

Tests should leverage the nested transaction (with a twist). Typically tests should upon startup:

• Create a new session
• Use a nested transaction

And upon teardown:

• Rollback (as opposed to committing) the transaction (if possible)
• Expunge all the objects from the session
• Close the session

A pytest fixture similar to this (with function scope) achieves the desired functionality. Note the fixture uses the nested transcript construct sans context manager to ensure that the transaction is never explicitly committed.

Examples

The following code illustrates how the DAO and Command should be constructed, where it is evident that the Command (business layer) has control of the transaction.

class BaseDAO:
    @classmethod
    def update(cls, item: T, attributes: dict[str, Any]) -> T:    
        for key, value in attributes.items():
            setattr(item key, value)
    
        db.session.merge(item)

class UpdateDashboardCommand:
    def run(self, item: Dashboard, attributes: dict[str, Any]) -> Dashboard:
        try:
            with db.session.begin_nested():
                DashboardDAO.update(item, attributes)
                DashboardDAO.update_owners(item)
        except SQLAlchemyError as ex:
            raise DashboardUpdateFailedError() from ex 

Flask-AppBuilder

Superset relies heavily on Flask-AppBuilder (FAB) which has a tendency to explicitly commit, i.e., SecurityManager.add_user(), thus violating our definition of "unit of work", i.e., the following,

with db.session.begin_nested():
    security_manager.add_user(...)
    db.session.rollback()

will not roll back to the savepoint at the beginning of the nested transaction, as a rollback or commit updates the savepoint.

There are a couple of options available to address this issue:

1. Override the SecurityManager.get_session() property by providing a monkey patched session (associated with a nested transaction) where the commit operation merely flushes* and the rollback operation is a no-op.
2. Have FAB use nested transactions removing the need to explicit commits.

Though (2) is preferable, (1) can be implemented as follows:

class SupersetSecurityManager(SecurityManager):
    @property
    def get_session(self) -> Session:
        if db.session._proxied._nested_transaction:
            with db.session.begin_nested() as transaction:
                transaction.session.commit = transaction.session.flush
                transaction.session.rollback = lambda: None
                return transaction.session

        return db.session

* The reason the commit flushes as opposed to being a no-op (which SQLAlchemy invokes unconditionally prior to a commit) is this helps preserve existing commit workflows where a series of operations (insert, update, delete) are commicated to the database .

New or Changed Public Interfaces

None.

New Dependencies

May require an update to FAB to ensure that the atomic unit remains intact.

Migration Plan and Compatibility

• A wiki page will be created to specify the rules which should be adhered to.
• A series of PRs to address existing rule violations.
• New PRs (where possible) should adhere to the rules.

Rejected Alternatives

None.

[craig-rueda]

This is great! Have you thought about introducing an external (to SQLA) "manager" of sorts that tracks transaction lifecycles and is responsible for proxying/postponing commit() calls in the event that some third-party lib such as FAB attempts to perform a commit() out of turn. I'm thinking that something like a custom "session provider" might be a potential solution that patches the session.commit() method and instead only actually commits when the unit of work is complete.

Adding such a construct could also allow us to build unit tests that "automatically" rollback after running against their target DB, which would remove the need to perform cleanup after each test. (Ex: https://docs.spring.io/spring-framework/reference/testing/annotations/integration-spring/annotation-rollback.html)

One other potential case that I've seen in the past is when you may actually want to nest transactions (i.e. you need to create / commit some object within the context of a unit of work that doesn't necessarily depend on the outer scope/transaction)

Lastly, support for "read-only" transactions would be awesome. This would remove the need to even BEGIN a transaction, let alone to ROLLBACK transactions, which is the current behavior of Flask-SQLA.

[michael-s-molina]

Thanks for writing this SIP @john-bodley. I already review it internally so I only have one more addition: it's important that we are able to track the SIP progress given that it will generate a series of PRs. For SIP-61, I created a SIP-61 milestones project and we could do something similar here. Planning the SIP execution also have the additional benefit of enabling contributions from the community.

[john-bodley]

Thanks @craig-rueda for your feedback. In answer to your questions:

• I didn't consider an external SQLAlchemy manager outside of Flask-SQLAlchemy which is the defacto session manager for Superset, Flask-AppBuilder, and Flask-Migrate (as defined here). In terms of proxying/postponing the commit we could handle this on a per dependency basis as monkey patching this globally could be difficult. There's a section above which mentions how to deal with Flask-AppBuilder, and (outside of testing and non-Flask templates) it seems like Flask-Migrate does not explicitly commit. Not doubt there are pros/cons with having an external SQLAlchemy manger, but one significant pro of the Flask-SQLAlchemy approach is nested transactions can be readily adopted by Flask-AppBuilder etc. and seamlessly integrated with Superset given the both share the the same instance of the of the preconfigured scoped session (db.session) as described in SIP-99A.

• I did share a pattern which we could use/augment for testing. There are likely instances where we do need to commit the transaction (due to a subsequent API call) whereas other times having changes persisted to the SQLAlchemy session and simply rolling back is suffice.

• Sadly I don't believe that SQLAlchemy supports a "read-only" transaction given that the “virtual” transaction is created automatically when needed, though I like the idea from a readability etc. standpoint. Per a Google search on this topic, one could define a read-only user or configure this at the engine level (example).

[john-bodley]

cc @dpgaspar as this SIP references FAB.

[dpgaspar]

@john-bodley great work!
More then happy to implement optional non explicit commits to flask-appbuilder, let's sync on this and coordinate.

[rusackas]

Approved!

[rusackas]

If anyone can give us all an update on implementation status or plans, it would be appreciated :D