How to DRY up this psycopg connection pool boilerplate code with a reusable async function or generator?

Question

I’m using psycopg to connect to a PostgreSQL database using a connection pool. It works great, but any function that needs to run SQL in a transaction gets three extra layers of nesting:

/app/db.py

from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))

/app/foo.py

from db import pool
from psycopg.rows import dict_row


async def create_foo(**kwargs):
    foo = {}
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(row_factory=dict_row) as cur:
                # use cursor to execute SQL queries
    return foo


async def update_foo(foo_id, **kwargs):
    foo = {}
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(row_factory=dict_row) as cur:
                # use cursor to execute SQL queries
    return foo

I wanted to abstract that away into a helper function, so I tried refactoring it:

/app/db.py

from contextlib import asynccontextmanager
from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))


@asynccontextmanager
async def get_tx_cursor(**kwargs):
    async with pool.connection() as conn:
        conn.transaction()
        cur = conn.cursor(**kwargs)
        yield cur

…and calling it like this:

/app/foo.py

from db import get_tx_cursor
from psycopg.rows import dict_row


async def create_foo(**kwargs):
    foo = {}
    async with get_tx_cursor(row_factory=dict_row) as cur:
        # use cursor to execute SQL queries
    return foo

…but that resulted in an error:

TypeError: '_AsyncGeneratorContextManager' object does not support the context manager protocol

I also tried variations of the above, like this:

async def get_tx_cursor(**kwargs):
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(**kwargs) as cur:
                yield cur

…but got similar results, so it appears using a generator is not possible.

Does anyone know of a clean and simple way to expose the cursor to a calling function, without using another library?

Here are the versions I’m using:

python: 3.11
psycopg: 3.1.8
psycopg-pool: 3.1.6

Asked By: Shaun Scovil

||

Source

Answer 1

You were on the right track with the asynccontextmanager decorator, but you forgot to use the async with statement for the transaction and the cursor inside the get_tx_cursor function.

Here’s the corrected version of your /app/db.py:

from contextlib import asynccontextmanager
from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))


@asynccontextmanager
async def get_tx_cursor(**kwargs):
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(**kwargs) as cur:
                yield cur

Now you can use the get_tx_cursor function in your /app/foo.py as you intended:

from db import get_tx_cursor
from psycopg.rows import dict_row

async def create_foo(**kwargs):
    foo = {}
    async with get_tx_cursor(row_factory=dict_row) as cur:
        # use cursor to execute SQL queries
    return foo

This should work without any errors, and you can reuse the get_tx_cursor function in other parts of your application as well.

Answered By: PCDSandwichMan

How to DRY up this psycopg connection pool boilerplate code with a reusable async function or generator?

Question:

Answers: