Can't access rowcount from within custom cursorfactory

[sander1234567890]

Hi,

I have a custom cursorfactory based on the psycopg2 RealDictCursor extension. That all works fine both on psycopg2 as with momoko, except when i try to access the rowcount of a query from within the custom cursorfactory class itself.

It works when i use the cursorfactory with psycopg2, but when i use the same cursorfactory with momoko it always returns -1 from within the cursorfactory.

Example code and results are below

Sander

Test example returns this:

PSYCHOPG2

rows within cursorfactory TestCursor | 10, 10, 10 rows before fetch: 10 rows after fetch: 10 10

MOMOKO

rows within cursorfactory TestCursor | -1, -1, -1 rows before fetch: 10 rows after fetch: 10 10

Code:

import psycopg2
import psycopg2.extras
import psycopg2.extensions
import momoko
import tornado.ioloop

class TestCursor(psycopg2.extras.RealDictCursor, psycopg2.extensions.cursor):

    def __init__(self, *args, **kwargs):

        super(TestCursor, self).__init__(*args, **kwargs)

    def execute(self, sql_query, args=None):
        """
        :param sql_query:
        :param args:
        """

        try:
            super(TestCursor, self).execute(sql_query, args)

        except (Exception, psycopg2.Warning, psycopg2.Error) as e:
            print("%s | Database query failed: %s" % (e.__class__.__name__, e))

        finally:
            print(
                "rows within cursorfactory %s | %s, %s, %s" % (
                    self.__class__.__name__,
                    self.rowcount, super().__getattribute__("rowcount"),
                    super(TestCursor, self).rowcount
                )
            )

def test_momoko():

    print("############# MOMOKO")

    ioloop = tornado.ioloop.IOLoop.instance()

    conn = momoko.Connection(
        dsn="dbname=test user=test password=test host=localhost port=5432",
        cursor_factory=TestCursor
    )

    ioloop.add_future(conn.connect(), lambda x: ioloop.stop())
    ioloop.start()

    future = conn.execute(
        """
        SELECT *
        FROM test_table
        LIMIT 10;
        """
    )
    ioloop.add_future(future, lambda x: ioloop.stop())
    ioloop.start()
    cursor = future.result()

    print("rows before fetch: %s" % (cursor.rowcount,))
    results = cursor.fetchall()
    print("rows after fetch: %s %s" % (cursor.rowcount, len(results)))

def test_psychopg2():
    print("############# PSYCHOPG2")

    with psycopg2.connect(
        dsn="dbname=test user=test password=test host=localhost port=5432",
        cursor_factory=TestCursor
    ) as conn:

        with conn.cursor() as cursor:

            cursor.execute(
                """
                SELECT *
                FROM test_table
                LIMIT 10;
                """
            )

            print("rows before fetch: %s" % (cursor.rowcount,))
            results = cursor.fetchall()
            print("rows after fetch: %s %s" % (cursor.rowcount, len(results)))

test_psychopg2()
test_momoko()

[haizaar]

What Python version? What momoko version? Linux/Windows?

[haizaar]

Actually, you can not expect rowcount property to be updated in execute method when in async mode - execute method finishes even before the query is sent to the server.

I suggest you consult with psycopg2 guys how to hook on execution completion when implementing custom cursors in async mode.

On the other hand, you can extend momoko and add "done_callback" to future it returns to do whatever you originally wanted to do.

Feel free to reopen if you think otherwise.

Zaar