How can I get affected row count from psycopg2 connection.commit()?

amphibient picture amphibient · Jun 27, 2014 · Viewed 34.5k times · Source

Currently, I have the following method to execute INSERT/UPDATE/DELETE statements using psycopg2 in Python:

def exec_statement(_cxn, _stmt):
    try:
        db_crsr = _cxn.cursor()
        db_crsr.execute(_stmt)
        _cxn.commit()
        db_crsr.close()
        return True
    except:
        return False

But what I would really like it to do, instead of bool, is return the row count affected by the transaction or -1 if the operation fails.

Is there a way to get a number of rows affected by _cxn.commit()? E.g. for a single INSERT it would be always 1, for a DELETE or UPDATE, the number of rows affected by the statement etc.?

Answer

khampson picture khampson · Jun 27, 2014

commit() can't be used to get the row count, but you can use the cursor to get that information after each execute call. You can use its rowcount attribute to get the number of rows affected for SELECT, INSERT, UPDATE and DELETE.

i.e.

    db_crsr = _cxn.cursor()
    db_crsr.execute(_stmt)

    rowcount = db_crsr.rowcount

    _cxn.commit()
    db_crsr.close()

    return rowcount

If you want to return the number of affected rows, I would recommend not catching any exceptions, since if the operation truly failed (say the query was malformed, or there was a FK constraint violation, etc.), an exception should be raised, and in that case the caller could catch that and behave as desired. (Or, if you want to centralize the exception handling, perhaps raise a custom MyPostgresException, or similar.)

-1 can be returned in a non-failure case in certain situations (http://initd.org/psycopg/docs/cursor.html#cursor.rowcount), so I would recommend against using that value as the failure indicator. If you really want to return a numerical value in the case of failure, perhaps returning a number like -10 would work (in the except block), since rowcount shouldn't ever return that.