CODEWITHSUNDEEP
pythonpostgresqlpsycopg2
Geoff
Geoff

Reputation: 1634

Psycopg2 rowcount for server side cursor

I am querying a Postgres database for a large number of results and want to use server side cursors to stream the results to my client. It looks like when I do this, the rowcount attribute of the cursor is now set to -1 after I execute the query. I'm creating the cursor like so:

with db.cursor('cursor_name') as cursor:

Is there a way to find the number of results of my query while streaming results from the database? (I could do a SELECT COUNT(*), but I'd like to avoid that because I'm trying to abstract away the code around the query and that would complicate the API).

Upvotes: 1

Views: 2824

Answers (1)

mhawke
mhawke

Reputation: 87154

In the case of a server-side cursor, although cursor.execute() returns, the query has not necessarily been executed by the server at that point, and so the row count is not available to psycopg2. This is consistent with the DBAPI 2.0 spec which states that rowcount should be -1 if the row count of the last operation is indeterminate.

Attempts to coerce it with cursor.fetchone(), for example, updates cursor.rowcount, but only by the number of items retrieved, so that is not useful. cursor.fetchall() will result in rowcount being correctly set, however, that performs the full query and transfer of data that you seek to avoid.

A possible workaround that avoids a completely separate query to get the count, and which should give accurate results is:

select *, (select count(*) from test) from test;

This will result in each row having the table row count appended as the final column. You can then get the table row count using cursor.fetchone() and then taking the final column:

with db.cursor('cursor_name') as cursor:
    cursor.execute('select *, (select count(*) from test) from test')
    row = cursor.fetchone()
    data, count = row[:-1], row[-1]

Now count will contain the number of rows in the table. You can use row[:-1] to refer to the row data.

This might slow down the query because a possibly expensive SELECT COUNT(*) will be performed, but once done retrieving the data should be fast.

Upvotes: 2

Related Questions