Creating transactions with with statements in psycopg2

I am trying to use psycopg2 to add some new columns to a table. PostgreSQL lacks a ALTER TABLE table ADD COLUMN IF NOT EXISTS, so I am adding each column in it’s own transaction. If the column exists, there will be a python & postgres error, that’s OK, I want my programme to just continue and try to add the next column. The goal is for this to be idempotent, so it can be run many times in a row.

It currently looks like this:

def main():
    # <snip>
    with psycopg2.connect("") as connection:
        create_columns(connection, args.table)

def create_columns(connection, table_name):
    def sql(sql):
        with connection.cursor() as cursor:
            cursor.execute(sql.format(table_name=table_name))

    sql("ALTER TABLE {table_name} ADD COLUMN my_new_col numeric(10,0);")
    sql("ALTER TABLE {table_name} ADD COLUMN another_new_col INTEGER NOT NULL;")

However, if my_new_col exists, there is an exception ProgrammingError('column "parent_osm_id" of relation "relations" already existsn',), which is to be expected, but when it tried to add another_new_col, there is the exception InternalError('current transaction is aborted, commands ignored until end of transaction blockn',).

The psycogpg2 document for the with statement implies that the with connection.cursor() as cursor: will wrap that code in a transaction. This is clearly not happening. Experimentation has shown me that I need 2 levels of with statements, to including the pscyopg2.connect call, and then I get a transaction.

How can I pass a connection object around and have queries run in their own transaction to allow this sort of “graceful error handling”? I would like to keep the postgres connection code separate, in a “clean architecture” style. Is this possible?

Answer

The psycogpg2 document for the with statement implies that the with connection.cursor() as cursor: will wrap that code in a transaction.

this is actually not true it says:

with psycopg2.connect(DSN) as conn:
    with conn.cursor() as curs:
       curs.execute(SQL)

When a connection exits the with block, if no exception has been raised by the block, the transaction is committed. In case of exception the transaction is rolled back. In no case the connection is closed: a connection can be used in more than a with statement and each with block is effectively wrapped in a transaction.

So it’s not about cursor object being handled by with but the connection object

Also worth noting that all resource held by cursor will be released when we leave the with clause

When a cursor exits the with block it is closed, releasing any resource eventually associated with it. The state of the transaction is not affected.

So back to your code you could probably rewrite it to be more like:

def main():
    # <snip>
    with psycopg2.connect("") as connection:
        create_columns(connection, args.table)

def create_columns(con, table_name):
    def sql(connection, sql):
        with connection:
            with connection.cursor() as cursor:
                cursor.execute(sql.format(table_name=table_name))

    sql(con, "ALTER TABLE {table_name} ADD COLUMN my_new_col numeric(10,0);")
    sql(con, "ALTER TABLE {table_name} ADD COLUMN another_new_col INTEGER NOT NULL;")

ensuring your connection is wrapped in with for each query you execute, so if it fails connection context manager will revert the transaction