[Python] psycopg2 'execute block of code in transaction' context manager

pg-transaction.py

@contextmanager
def transaction(pool='common_write'):
    """Context manager.
    Executes block inside DB transaction. Returns cursor.
    At the end of the block, the connection is returned to pool.

    >>> with transaction() as cursor:
    ...     rows = cursor.execute(...).fetchall()
    ...     process(rows)
    ...     cursor.execute(...)
    """
    pool = random.choice(pool_map[pool])
    connection = pool.get()
    try:
        cursor = connection.cursor()
        cursor.execute("begin")
        try:
            yield cursor
        except Exception:
            cursor.execute("rollback")
            raise
        else:
            cursor.execute("commit")
    finally:
        connection.close() # return connection to pool