[Python] How to use role based authentication in SQLAlchemy without blocking the event loop

I am using SQLAlchemy to connect to AWS RDS Postgres and I am trying to use role based authentication to make the connections. In that case, the password token is temporary, meaning a connection may fail to access the database at any point in time and a new one has to be added to the connection pool with a fresh password token.

I am using SQLAlchemy taking advantage of its asynchronous capabilities in the 2.0 version and all code I have is assynchronous. However, when it comes to creating a new connection with the a new password token, the only way I found is to have a listener to the do_connect event. The problem here is that the do_connect is called as a synchronous function which means the code blocks while it waits for the call to get a new token from AWS.

The code I am currently using is the one below. Is there a way to avoid blocking on each new connection? I went through SQLAlchemy documentation but couldn't find anything about it.

import asyncio
import threading

from psycopg._adapters_map import AdaptersMap
from sqlalchemy import Dialect
from sqlalchemy.event import listen
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.pool import ConnectionPoolEntry


async def get_aws_token() -> str:
await asyncio.sleep(1) # This would be an async call using boto3
return "token"


def create_engine():
engine = create_async_engine(
# No password is passed in the url as it changes frequently
url="postgresql+asyncpg://user@localhost/test",
pool_size=5,
pool_pre_ping=True,
)
loop = asyncio.new_event_loop()
threading.Thread(target=loop.run_forever, daemon=True).start()

def on_connect(
dialect: Dialect,
conn_rec: ConnectionPoolEntry,
cargs: list[str | int],
cparams: dict[str, str | int | AdaptersMap],
) -> None:
future = asyncio.run_coroutine_threadsafe(
get_aws_token(), loop
)
cparams["password"] = future.result()

listen(engine.sync_engine, "do_connect", on_connect)

return engine


async def main():
engine = create_engine()
async with engine.begin() as conn:
await conn.execute("SELECT 1")


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

Continue reading...