pythonpsycopg3

How to change cursor type in with clause


I have the following statement

    statement = """
    INSERT INTO rag.chunks (project_id, document_id, model_id, text, metadata)
    VALUES {}
    RETURNING *;
    """

    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            values = ', '.join(
                cur.mogrify("(%s, %s, %s, %s, %s)",
                            (c.project_id, c.document_id, c.model_id, c.text, Jsonb(c.metadata)))
                for c in chunks
            )

            statement = statement.format(values)
            await cur.execute(statement)
            chunks = await cur.fetchall()
            return chunks

Am using cur.mogrify because execute_many doesnt' return all the rows, just the last row (discussion).

However, because this is an AsyncConnectionPool, the cursor being returned is AsyncCursor, which does not have mogrify method.

AttributeError: 'AsyncCursor' object has no attribute 'mogrify'

However, the AsyncClientCursor class does have mogrify.

How can I specify I want to use an AsyncClientCursor? For example:

async with pool.connection() as conn:
    async with conn.cursor(cursorClass=AsyncClientCursor) as cur:    

But this fails with

TypeError: AsyncConnection.cursor() got an unexpected keyword argument 'cursorClass'

Solution

  • Figured it out - not sure if the ideal way but it can be set like this

    async with pool.connection() as conn:
        conn.cursor_factory = AsyncClientCursor
        async with conn.cursor() as cur:
            ...