Is it possible to read and write to the TCP transport of asyncpg's connection class?

[jackwotherspoon]

Is there a public interface to read and write data to the TCP transport of asyncpg's Connection class?

In the past I have used pg8000 where I could take advantage of its Connection class' public methods sock_read and sock_write to read and write data to the underlying TCP socket of the database Connection class.

Example of pg8000:

import pg8000
connection = pg8000.dbapi.connect(
        user,
        database=db,
        password=passwd,
        host=ip_address,
        port=SERVER_PROXY_PORT,
        ssl_context=ctx,
        **kwargs,
    )

# write data to socket
connection.sock_write(buffer)

# read data from socket
resp = connection.sock_read(bytes)

Is there a way to do something similar in asyncpg? I am relatively new to asyncpg and asyncio so I am wondering what the best way to achieve a similar pattern would be?

[jackwotherspoon]

@elprans Would the suggested way to implement this be to create a custom connection_class to be returned by asyncpg.connect? Or is there an easier solution that I am missing 😄

[elprans]

You can try connection._transport.get_extra_info('socket'). The issue is that doing anything to the underlying socket is usually very dicey. What is your use case exactly?

[enocom]

A little context here: the Cloud SQL Python Connector uses asyncpg (among other drivers) to connect to Cloud SQL instances. Cloud SQL provides a server-side proxy that sits along side Postgres instances that the Python Connector connects to as part of establishing a connection.

We're exploring ways to improve IAM-level permission checks in the server side proxy and in one approach, we'd write an OAuth2 token into the TLS socket prior to the database protocol taking over. Once the write was complete (and the token authorized), we'd hand the socket over to the database driver for the usual Postgres connection. Hence the question for the best way to get ahold of the socket.

[elprans]

Hooking into connection after it is established but before any messages get sent would be fairly hairy. Consider sending the auth token as a server parameter in the startup message, e.g:

await asyncpg.connect(..., server_settings={..., "cloudsql.auth": "<token>"})

The proxy can then parse the startup message and extract the token from there.

[enocom]

That's a great idea. Thank you for the suggestion.