sql: support asyncpg

[jordanlewis]

This is a tracking issue for supporting asyncpg.

• MOVE is expected to be able to interact with pgwire-level portals. It currently does not, which results in errors like cursor "__asyncpg_portal_3__" does not exist when trying to use the asyncpg API cursor.forward.
• Multiple Active Portals support pgwire: multiple active result sets (portals) not supported #40195. In particular, if you open a cursor, don't close it, and then run a statement, you run into the MARS issue.
• decorrelate subqueries involving CASE: sql: query decorrelation error under recursive CTE in asyncpg #71908

Jira issue: CRDB-16803

[jordanlewis]

A very simple test case that reproduces the MARS issue is as follows:

#!/usr/bin/python

import asyncio
import asyncpg

async def main():
    # Establish a connection to an existing database named "test"
    # as a "postgres" user.
    conn = await asyncpg.connect('postgresql://root:26257@localhost/defaultdb')

    async with conn.transaction():

        cur = await conn.cursor('SELECT generate_series(1,10000)')

        print(await cur.fetch(5))

        print("hiii")
        await cur.close()
        print("hi")
        await conn.execute('''
            INSERT INTO users(name, dob) VALUES($1, $2)
        ''', 'Bob', datetime.date(1984, 3, 1))

    # Close the connection.
    await conn.close()

asyncio.get_event_loop().run_until_complete(main())

Interestingly, asyncpg doesn't have an explicit close method on its cursor - the close method seems to no-op or delegates to something else.

I think the root cause might actually be that asyncpg needs to be upgraded to support being able to close a cursor. Without that, it seems like it'd just be leaking cursors all over the place since they're uncloseable without scrolling to the end.

[jordanlewis]

Here's a more minimal example:

#!/usr/bin/python

import asyncio
import asyncpg

async def main():
    conn = await asyncpg.connect('postgresql://root:26257@localhost/defaultdb')

    async with conn.transaction():

        cur = await conn.cursor('SELECT generate_series(1,10)')

        print(await cur.fetch(10))

        # Without the following line, an error is returned later:
        print(await cur.fetch(1))

        await conn.execute('''SELECT 1''')

    # Close the connection.
    await conn.close()

asyncio.get_event_loop().run_until_complete(main())

If the cursor is properly drained, no error is returned.

This makes me think that it's likely that asyncpg errors that we see in the wild are more likely caused by improperly closed cursors than true multiple-active-cursor support being required...

[jordanlewis]

One very interesting thing we could do, assuming that people "abandoning old cursors" is the common case (and that we aren't expecting people to actually be interleaving cursors), is the following diff (possibly gated behind a session flag?). This does cause the test cases above to work, but I haven't played around with what else it would break.

diff --git a/pkg/sql/pgwire/command_result.go b/pkg/sql/pgwire/command_result.go
index 5180a3069b..778406bfbd 100644
--- a/pkg/sql/pgwire/command_result.go
+++ b/pkg/sql/pgwire/command_result.go

@@ -508,17 +507,9 @@ func (r *limitedCommandResult) moreResultsNeeded(ctx context.Context) error {
                                return err
                        }
                default:
-                       // If the portal is immediately followed by a COMMIT, we can proceed and
-                       // let the portal be destroyed at the end of the transaction.
-                       if isCommit, err := r.isCommit(); err != nil {
-                               return err
-                       } else if isCommit {
-                               return r.rewindAndClosePortal(ctx, prevPos)
-                       }
-                       // We got some other message, but we only support executing to completion.
-                       telemetry.Inc(sqltelemetry.InterleavedPortalRequestCounter)
-                       return errors.WithDetail(sql.ErrLimitedResultNotSupported,
-                               fmt.Sprintf("cannot perform operation %T while a different portal is open", c))
+                       // If we're not trying to do anything more with the portal, simply destroy
+                       // it and let the transaction continue along.
+                       return r.rewindAndClosePortal(ctx, prevPos)

[jordanlewis]

Another very interesting thing: asyncpg doesn't even expect that it's possible to run multiple active cursors, even though Postgres does actually permit it based on the protocol.

This says to me that any MARS activity generated by asyncpg programs is actually not expected, and we should solve it by explicitly closing any open cursors on the asyncpg side when moving to the next operation.

MagicStack/asyncpg#738

it's possible that the linked issue is only about true concurrency rather than interleaved cursors, but I don't really think that the library is expecting people to interleave cursors either.

[pablo-sumup]

Perhaps you want to test those cases using PostgreSQL.
We are evaluating CockroachDB and we have the same problem with a SELECT * FROM table. It works fine in Postgresql 14.

We are using databases package to manage a pool with async function on python 3.10 and asyncpg 0.25.0

Do you have an ETA for this issue?