python - How to use UPDATE/DELETE WHERE CURRENT OF in psycopg2?

Question

I'm using server-side cursor in PostgreSQL with psycopg2, based on this well-explained answer.

with conn.cursor(name='name_of_cursor') as cursor:
    query = "SELECT * FROM tbl FOR UPDATE"
    cursor.execute(query)
    for row in cursor:
        # process row

In processing each row, I'd like to update a few fields in the row using PostgreSQL's UPDATE tbl SET ... WHERE CURRENT OF name_of_cursor (docs), but it seems that, when the for loop enters and row is set, the position of the server-side cursor is in a different record, so while I can run the command, the wrong record is updated.

How can I make sure the result iterator is in the same position as the cursor? (also preferably in a way that won't make the loop slower than updating using an ID)

标签： pythonpostgresqlpsycopg2

Answer 1

The reason why a different record was being updated was because internally psycopg2 does a FETCH FORWARD 1000 (or whatever the default chunk size is), positioning the cursor at the end of the block. You can override this by fetching one record at a time:

updcursor = conn.cursor()
with conn.cursor(name='name_of_cursor') as cursor:
    cursor.itersize = 1  # to make server-side cursor be in the same position as the iterator
    cursor.execute('SELECT * FROM tbl FOR UPDATE')
    for row in cursor:
        # process row...
        updcursor.execute('UPDATE tbl SET fld1 = %s WHERE CURRENT OF name_of_cursor', [val])

The snippet above will update the correct record. Note that you cannot use the same cursor for selecting and updating, they must be different cursors.

Performance

Reducing the FETCH size to 1 reduces the retrieval performance by a lot. I definitely wouldn't recommend using this technique if you're iterating a large dataset (which is probably the case you're searching for server-side cursors) from a different host than the PostgreSQL server.

I ended up using a combination of exporting records to CSV, then importing them later using COPY FROM (with the function copy_expert).