Unable to catch error within stream for loop

[comodore31]

Python 3.9
Psycopg 3.0.8

Unable to catch psycopg.OperationalError

import asyncio
import psycopg

async def main():
    async with await psycopg.AsyncConnection.connect(conninfo='postgresql://xxxxxxxxxxxxxx') as con:
        async with con.cursor() as cur:
            try:
                async for r in cur.stream('select 1 a'):
                    raise NameError('Error here!')
            except:
                pass

if __name__ == '__main__':
    asyncio.run(main())

and we get the error

  File "psycopg_binary\\pq/pgconn.pyx", line 223, in psycopg_binary.pq.PGconn.send_query
psycopg.OperationalError: sending query failed: another command is already in progress

[dvarrazzo]

I see the problem. It can be reproduced in the test suite with:

# tests/test_cursor_async.py

async def test_stream_error_python(aconn):
    await aconn.set_autocommit(True)
    cur = aconn.cursor()
    with pytest.raises(ZeroDivisionError):
        async for rec in cur.stream("select 1"):
            1 / 0

    assert aconn.info.transaction_status == aconn.TransactionStatus.INERROR

The test fails because the transaction status is found ACTIVE.

...it is not obvious to me how to fix it though :\ There doesn't seem to be a way to interrupt single-row mode from the client (as opposite as COPY, which as a way to inject an error in the server). Maybe throwing a cancel might work, but it requires testing.

[dvarrazzo]

Unfortunately stream() is a generator, not a context manager, so we don't really have a way to intercept an error in the consumer that makes the generator not entirely consumed.

Your script, changed this way, works:

        async with con.cursor() as cur:
            recs = cur.stream('select generate_series(1,10000)')
            try:
                async for r in recs:
                    raise NameError('Error here!')
            except:
                con.cancel()
                try:
                    async for r in recs:
                        pass
                except psycopg.errors.QueryCanceled:
                    pass

I don't think we have a way to fix it ourselves, unless we change the interface of stream(), to make it usable as:

        async with con.cursor() as cur:
            try:
                async with cur.stream('select 1 a') as stream:
                    async for rec in stream:
                        raise NameError('Error here!')
            except:
                pass

so it wouldn't be backward-compatible (I doubt we can write a cur.stream() able to identify whether it's being called as generator or as context manager). If any we could use a different method name.

[comodore31]

Ok thanks, maybe we could implement a new methode "stream_context"

[dvarrazzo]

With the current, iterator interface, the problem is in the hands of the user. Only consuming the iterator completely will restore the connection to an usable state. Will document here how to do that, but the idea is to implement a context interface sometimes in the future.

Failing to consume the iterator entirely will result in a connection left in transaction status ACTIVE: this connection will refuse to receive further commands.

try:
    for rec in cur.stream("select generate_series(1, 10)"):
        1 / 0
except ZeroDivisionError:
    pass

conn.execute("select 1")
# psycopg.OperationalError: sending query failed: another command is already in progress

You can restore the connection to a working state by consuming the generator entirely:

# Give the generator a name
stream = cur.stream("select generate_series(1, 10000)")
try:
    for rec in stream:
        1 / 0
except ZeroDivisionError:
    pass

# Make sure to consume all the rows.
for rec in stream:
    pass

# The connection is now usable.

If the generated rows are excessive or infinite and you wish to terminate the operation before consuming them all, you can cancel the query and recover from there.

# Send a cancel message to the server
conn.cancel()

# Consume the rows already produced
try:
    for rec in stream:
        pass

# The operation might or might not raise an exception for having been canceled,
# depending on when the server receives the signal
except psycopg.errors.QueryCanceled:
    pass

# Roll back the connection to restore the connection from an
# eventual QueryCanceled
conn.rollback()

# The connection is now usable.

[dvarrazzo]

Unhappily closing this, leaving recovery in the hands on the client. Proposals for a context interface for stream in #383.

[dlax]

If I understand correctly, that's something that might be resolved by leveraging the finalization mechanism for asynchronous generators. From user side, this implies calling the aclose() method of the generator (await gen.aclose()). On our side, this would probably imply implementing something in the finally: clause of the async generator function (there is none at the moment).

To avoid the need for calling aclose() (on user side), async_generators's aclosing() context manager might help.

---

Aside, that's something trio would warn about (through a ResourceWarning); and which I've worked around in fbd5dcc.

[dvarrazzo]

Thank you @dlax. There is also close() on regular generators. Maybe we can actually implement those to make easier for the client to recover.

[dlax]

Also, there's contextlib.aclosing() in stdlib from python 3.10 alongside closing().