Psycopg2: Allow to start transactions in autocommit mode

[anarazel]

Hi,

using autocommit mode can be quite useful for some long running applications, to make it less likely that transactions are accidentally left open, and to perform operations like VACUUM etc. psycopg provides for that by allowing to set autocommit to true.

Unfortunately I think that's too large a hammer in a lot of cases. What I think is commonly desired is being able to granularly choose the desired behaviour, but psycopg doesn't allow that currently.

E.g. even when in autocommit mode, one might want to explicitly use a transaction. But that doesn't really work well right now. Approaches I considered:
1) Use con.begin() - which doesn't exist
2) Just start a transaction using .execute() - but that breaks con.commit(), because:
https://github.com/psycopg/psycopg2/blob/master/psycopg/pqpath.c#L396

    if (conn->autocommit || conn->status != CONN_STATUS_BEGIN) {
        Dprintf("pq_commit: no transaction to commit");
        retvalue = 0;
}

even though the connection knows that it's in a transaction, it ignores that fact.

While I obviously just could commit using .execute('COMMIT') that seems too likely to break assumptions of other code.

3) Disable autocommit for a single transaction:

psycopg2.ProgrammingError: set_session cannot be used inside a transaction

While I could work around that by resetting autocommit later, that requires exception handling etc.

What I think a sensible approach would be:
1) Have pq_commit() actually commit, even if autocommit mode is on, if there's a transaction in progress - the transaction status is known at the client side
2) Provide .begin(), to explicitly start a transaction (probably erroring out if there already is one in progress)
3) Have a context manager for transactions (rather than just connections and cursors). Probably just return one from con.begin().

Regards

Andres

[dvarrazzo]

Hello Andres.

I think there are a few things that can be done, yes. First a correction: it's a bit surprising but the connection context manager doesn't exit with close, but it terminates a transaction.

I don't have code handy but I guess entering the CM with an autocommit transaction is a no-op.

So I see two easy improvements:

• change __enter__ to begin a transaction even if the connection is autocommit. It is a change in behaviour, but I don't think anyone really does it. We could introduce a warning in the current maintenance branch and change behaviour in the next major release.

• change commit() and rollback() to terminate the transaction on autocommit too, so that the connection would behave as expected even if someone issued a manual BEGIN (I don't think there are breaking assumptions around in the code).

In case of autocommit it would be better tracking the status using libpq PQtransactionStatus instead of conn->status (even betterer would be to always use the libpq, probably...)

Would you like to work on these changes in a MR?

[and-semakin]

My current workaround (that relies on 3rd approach of @anarazel) to this issue is following:

from contextlib import contextmanager
from typing import Iterator

from psycopg2.extensions import connection as PgConnection


@contextmanager
def transaction_in_autocommit_mode(connection: PgConnection) -> Iterator[None]:
    connection.rollback()
    connection.autocommit = False
    try:
        with connection:
            yield
    finally:
        connection.autocommit = True

Seems to be working as expected:

In [43]: def get_txid(cur):
    ...:     cur.execute("SELECT txid_current()")
    ...:     txid, = cur.fetchone()
    ...:     return txid
    ...:

In [44]: with transaction_in_autocommit_mode(conn):
    ...:     with conn.cursor() as cur:
    ...:         print(get_txid(cur))
    ...:         print(get_txid(cur))
    ...:
654
654

[tometzky]

Based on @and-semakin's code I've created an elegant (IMHO) connection factory with autocommit and explicit transactions:

"""AutocommitConnection"""

import contextlib
import sys
import typing

import psycopg2  # type: ignore


class AutocommitConnection(
    psycopg2.extensions.connection
):  # pylint: disable=too-few-public-methods
    """Connection factory with autocommit and explicit transactions"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.set_session(autocommit=True)

    @contextlib.contextmanager
    def transaction(self) -> typing.Iterator[None]:
        """Context manager to encapsulate a transaction"""
        self.set_session(autocommit=False)
        try:
            with self:
                yield
        finally:
            self.set_session(autocommit=True)

# Example usage:
# Uses default connection or environmental variables like PGHOST PGDATABASE PGUSER PGPASSWORD
conn: AutocommitConnection = psycopg2.connect(
    "", connection_factory=AutocommitConnection
)
with conn.cursor() as cur:
    cur.execute("create temporary table test_autocommit(id int)")
    cur.execute("insert into test_autocommit values (1)")
    try:
        cur.execute("insert into test_autocommit values ('A')")
    except psycopg2.Error as err:
        print("Error:", err, file=sys.stderr)
    with conn.transaction():
        cur.execute("insert into test_autocommit values (2)")
        cur.execute("insert into test_autocommit values (3)")
        # Automatic commit
    with conn.transaction():
        cur.execute("insert into test_autocommit values (4)")
        try:
            cur.execute("insert into test_autocommit values ('A')")
        except psycopg2.Error as err:
            print("Error:", err, file=sys.stderr)
        # Automatic rollback
    with conn.transaction():
        cur.execute("insert into test_autocommit values (5)")
        conn.rollback()
        # Explicit rollback

# Also works in one line:
with conn.cursor() as cur, conn.transaction():
    cur.execute("insert into test_autocommit values (6)")
    cur.execute("insert into test_autocommit values (7)")
with conn.cursor() as cur, conn.transaction():
    cur.execute("insert into test_autocommit values (8)")
    cur.execute("insert into test_autocommit values (9)")
    conn.rollback()

with conn.cursor() as cur:
    cur.execute("select id from test_autocommit order by id")
    for (val,) in cur:
        print(val)
    # It should print 1, 2, 3, 6, 7
    # It should not print 4, 5, 8, nor 9 as they were rolled-back