allow commit() within run_transaction() callback

[gpaul]

Running the following yields psycopg2.InternalError: SAVEPOINT not supported except for COCKROACH_RESTART

I believe this is due to https://github.com/cockroachdb/cockroachdb-python/blob/master/cockroachdb/sqlalchemy/transaction.py#L54.

Why do we need savepoint_state.cockroach_restart? Considering that CockroachDB supports no other savepoint name, why not hard-code it and drop the conditional in the following functions:
https://github.com/cockroachdb/cockroachdb-python/blob/master/cockroachdb/sqlalchemy/dialect.py#L158

Steps to reproduce:

docker run -p 26257:26257 cockroachdb/cockroach:v1.0 start --insecure

Then run the following code:
https://gist.github.com/gpaul/648189354975992a2c285143d08be88a

This gives the following output using cockroachdb-python v0.1:
https://gist.github.com/gpaul/14fa2967c42d4ccf6b223322a7967476

[bdarnell]

The callback passed to run_transaction should not try to commit the transaction; leave that to the wrapper.

We need savepoint_state.cockroach_restart because the semantics of the cockroach_restart psuedo-savepoint are not quite the same as regular savepoints (releasing the cockroach_restart savepoint commits the transaction). We don't support savepoints or nested transactions in general, only the specific form used by run_transaction. If we hard-coded the name then applications could try to use nested transactions in ways that would fail mysteriously (and while it would change the error message you see here, it wouldn't make it correct to call commit inside the run_transaction callback).

[gpaul]

Thanks for the explanation. Would you agree that cockroachdb forces one to use rather unnatural patterns in one's sqlalchemy code? Do you foresee the situation improving?

Consider the following example code that sets a config option in the database if not already set:

    def achieve_value_consensus(self, key, value):
        log.info('Achieve consensus on key `%s`.', key)
        dbsession.add(ConfigItem(key, value))
        try:
            dbsession.commit()
            log.info('Consensus: set the key `%s`.', key)
            return value
        except sqlalchemy.exc.IntegrityError as exc:
            log.info('Consensus: already set, read it.')
            dbsession.rollback()

        return dbsession.query(ConfigItem).filter_by(key=key).one().value

This code (and any other that calls commit() only once) could be trivially wrapped in a run_transaction callback from wherever it is called and the code base could be kept very portable. With the added constraint that commit() is not allowed with cockroachdb this code becomes:

    def achieve_value_consensus(self, key, value):
        log.info('Achieve consensus on key `%s`.', key)
        try:
            def callback(dbsession):
                dbsession.add(ConfigItem(key, value))
            run_transaction(_engine, callback)
            log.info('Consensus: I have set the key `%s`.', key)
            return value
        except sqlalchemy.exc.IntegrityError as exc:
            log.info('Consensus: already set, read it.')

        def callback(dbsession):
            return dbsession.query(ConfigItem).filter_by(key=key).one().value
        return run_transaction(_engine, callback)

That illustrates the kind of sweeping changes that would be required throughout our sqlalchemy codebase.

[bdarnell]

Yes, handling transaction restarts in this way requires intrusive changes to many apps. We've made some changes that reduce the likelihood of transaction restarts, so you might be able to get by without surrounding everything in a retry loop, but these restarts remain more common with CockroachDB than in other databases and I think you'll still want retry handling.

I appreciate the feedback about potential interfaces; I've never used the Session portions of sqlalchemy so I don't know what patterns are "natural" with it. Can you explain this achieve_value_consensus method a little more to me? Why does it use two transactions? We might be able to allow the pattern you're describing in the text of your comment, in which a single call to commit or rollback could be allowed within a run_transaction callback (and then any subsequent attempts to use the transaction would fail). However, the fact that you do a separate query after committing the initial transaction wouldn't work in this model - you'd need full nested transaction support to make this work as written.

My recommendation for now is to think of run_transaction (or something like it) as portable. It's not the way you're writing your app currently, but once you've made the change, you won't be locked in to cockroachdb. The same pattern will work for any database (and in fact could be desirable even on other databases to handle deadlocks).

If achieve_value_consensus used a single transaction, I would write it like this (note that reading before the write is going to be substantially faster than writing first and catching the IntegrityError in CockroachDB):

def achieve_value_consensus(self, key, value):
    log.info('Achieve consensus on key `%s`.', key)
    def callback(dbsession):
        try:
            return dbsession.query(ConfigItem).filter_by(key=key).one().value
        except sqlalchemy.orm.exc.NoResultFound:
            dbsession.add(ConfigItem(key, value))
            return value
    run_transaction(_engine, callback)

Since the entire function is now in one transaction, a decorator form of run_transaction might be more convenient.

[gpaul]

@bdarnell firstly, thanks for all this effort. The response from you and the rest of the cockroach team has been stellar.

We might be able to allow the pattern you're describing in the text of your comment, in which a single call to commit or rollback could be allowed within a run_transaction callback (and then any subsequent attempts to use the transaction would fail). 

This sounds perfect. I don't want to dwell on the example I gave because honestly that is the only function in our entire codebase which performs more than one transaction, and as you elegantly showed, we can trivially rewrite it as a single transaction.

Would you consider reopening this issue in the light of your suggested patch? Alternatively, one could open a new issue to track the implementation thereof.

[bdarnell]

I'll reopen, but it's generous to call what I said above a "patch". I don't know exactly what the change would look like.

[gpaul]

I'd like to come back to this now that I have a bit more experience with cockroachdb-python and SQLAlchemy.

Firstly, what you have right now is one of two correct ways to go about it.

The cockroachdb-python package' run_transaction function uses a nested transaction.

The SQLAlchemy's nested transaction API has two different styles in which it can be used.

The first style is as a context manager:

with session.begin_nested():
    ...issues a SAVEPOINT statement...
    # execute statements in the nested transaction
    session.add(...)
    session.remove(...)
    ...__exit__() issues a commit() or rollback() for you...

The second style is more manual. It requires one to explicitly issue exactly one commit() or rollback() instruction for every begin_nested() instruction.

session.begin_nested() # issues a SAVEPOINT statement
try:
    # execute statements in the nested transaction
    session.add(...)
    session.remove(...)
    session.commit()
except Exception:
    session.rollback()

The relevant quote from the SQLAlchemy docs is:

For each begin_nested() call, a corresponding rollback() or commit() must be issued. (But note that if the return value is used as a context manager, i.e. in a with-statement, then this rollback/commit is issued by the context manager upon exiting the context, and so should not be added explicitly.)

~ http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#using-savepoint

In my case, I would like to call commit() and rollback() explicitly. As such, I have implemented a version of run_transaction that does not use begin_nested() as a context manager, but as a simple method call. The caller is then responsible for calling commit() or rollback() in the callback that is provided.

The exact code is quite tailored to our code base but I've reproduced it below anyway as it might help further the discussion.

def run_transaction(callback):
    # `dbsession` is a SQLAlchemy `scoped_session` object: a thread-local
    # session registry.
    # `callback` is a callable that operates directly on the same `dbsession`
    # and therefore does not take a `conn` argument the way the original
    # `run_transaction`'s callback does.
    # The `callback` is expected to call `commit()`. If a database exception is
    # raised, `rollback()` will be called by `run_transaction`.
    return _txn_retry_loop(dbsession, callback)


class _NestedTransaction(object):

    def __init__(self, session):
        self.session = session

    def __enter__(self):
        self.session.begin_nested()
        # Sessions are lazy and don't execute the savepoint
        # statement until you ask for the connection.
        self.session.connection()
        return self

    def __exit__(self, typ, value, tb):
        # Ignore any non-SQLAlchemy exceptions such as those used to trigger redirects.
        if typ is not None and issubclass(typ, sqlalchemy.exc.SQLAlchemyError):
            self.conn.rollback()


def _txn_retry_loop(session, callback):
    # Using cockroachdb it is never legal to start or end a nested
    # transaction while `savepoint_state.cockroach_restart is `False`.
    # Doing so would cause a `RELEASE SAVEPOINT {name}` or
    # `ROLLBACK TO SAVEPOINT {name}` instruction
    # where `name != 'cockroach_restart'`.
    #
    # As cockroach_restart is a thread-local variable it is not enough
    # to set it globally. It must be set by every thread. As such,
    # we set it here, before it gets used by any thread.
    savepoint_state.cockroach_restart = True
    with session.begin():
        while True:
            try:
                with _NestedTransaction(session):
                    return callback()
            except sqlalchemy.exc.DatabaseError as e:
                if isinstance(e.orig, psycopg2.OperationalError):
                    if e.orig.pgcode == psycopg2.errorcodes.SERIALIZATION_FAILURE:
                        continue
                raise

[bdarnell]

It seems error prone to require that the callback run commit or rollback exactly once; if you don't, then you've left the transaction nesting imbalanced. Is there some way we can ask the session for it's nesting depth and do the rollback whenever the transaction we started is still open? It looks like we might be able to save self.session.transaction after calling begin_nested, then do the rollback if self.session.transaction is the same object in __exit__.

What happens if you end up calling commit twice in a callback?

[gpaul]

Interesting and insightful questions, thank you.

you've left the transaction nesting imbalanced. Is there some way we can ask the session for it's nesting depth and do the rollback whenever the transaction we started is still open?

Right now I'm leaning towards saying "cockroachdb doesn't support nested transactions inside other nested transactions" and punt on the nesting depth issue until the database supports depth of greater than 2 (ie., the outer transaction and the nested transaction therein.)

If you mean "rollback unless commit() / rollback() was explicitly called", I think that's the default behaviour. However, I'd prefer an exception in that case: it is unlikely that the author uses my variation of run_transaction knowing that an explicit commit() or rollback() is required and intentionally omitted both.

What happens if you end up calling commit twice in a callback?

That raises an exception along the lines of "Transaction is already closed."

It seems error prone to require that the callback run commit or rollback exactly once;

In my case I opted to require a single explicit commit() but relaxed the rollback() requirement by automatically executing rollback() if an exception is raised.

I would say it is more error prone, yes. In our project I discovered a couple of functions that called commit() twice! Without solid unit tests and good test coverage these would have been errors in production.

That said, this approach complies with the SQLAlchemy docs and the explicit commit() comes with the benefit of being able to know that the transaction was successfully committed from within the callback and log or branch on that fact.

I don't think one way is better than the other: each have trade-offs to consider.

[bdarnell]

Right now I'm leaning towards saying "cockroachdb doesn't support nested transactions inside other nested transactions" and punt on the nesting depth issue until the database supports depth of greater than 2 (ie., the outer transaction and the nested transaction therein.)

We definitely don't support more than 1 nested transaction. My point is that with your version of run_transaction, if you forget to call commit, then you've left your connection in a weird state and the next thing you do will potentially get tied into this still-open transaction (or maybe it'll just get a confusing error message when it tries to start another nested layer).

If you mean "rollback unless commit() / rollback() was explicitly called", I think that's the default behaviour. However, I'd prefer an exception in that case

Raising an exception (after rolling back to put the connection back in a clean state!) would be fine too (but you're not doing it here).

the benefit of being able to know that the transaction was successfully committed from within the callback and log or branch on that fact.

Why does it matter that you know that inside the callback instead of in the function that called run_transaction?

In a clean-slate application, I'd argue that the design in which run_transaction is responsible for committing the transaction is superior because it is less error-prone. However, I can see that since most existing sqlalchemy code is built around explicit commits, an interface like this can be useful (as a transitional strategy if nothing else), so I could see offering it as an alternative. Now for the tricky part: how to name it to convey the difference from regular run_transaction?

[alanhamlett]

I'm also interested in nested transaction support, since my application uses the decorator form of begin_nested often. Any plans to support nested transactions in the future?