retry a query after a OperationalError or InternalError

[dfroger-idnow]

Hi,

The problem

I want to implement a "retry" mechanism for a web service (FastAPI) I'm working on.

Basically, a MariaDB Connection and a Cursor are provided (via Dependency Injection) to the Python coroutine function that handles an HTTP request and, is case of OperationalError (maybe a network issues closes the underlying connection) or InternalError , I want to acquire a new Connection/Cursor to retry the query.

The aiomysql library provides context manager for both Connection and Cursor. Putting all the coroutine function that handles an HTTP request in their context was my first idea, but it's not suitable to re-acquiring the Connection/Cursor for retry:

# That's peudo-code

# HTTP server level
POOL = await create_pool(...)

# HTTP request level
async with POOL.acquire() as conn:
    async with conn.cursor() as cur:
        try:
            cur.execute("QUERY-1")
            conn.commit()
        except (OperationalError, InternalError):
            # Here, we would like to release (conn, cur) and
            # acquire new ones, but we are stuck in the context
            # managers.

        # The next queries must be performed with the new acquired
        # (cur, conn), so inside it's managed context. try/except
        # logic must be present too.
        cur.execute("QUERY-2")
        conn.commit()

Instead, I propose to put the coroutine function for the HTTP request inside the context of a DatabaseRetryContext context manager, which exposes a Connection/Cursor, and a method to re-acquire them. The pseudo-code now becomes:

# That's peudo-code

# HTTP server level
POOL = await create_pool(...)

# HTTP request level
async with DatabaseRetryContext(pool) as ctx:
    try:
        ctx.cur.execute("QUERY-1")
        ctx.conn.commit()
    except (OperationalError, InternalError):
        ctx.reconnect()
        # try again
        ctx.cur.execute("QUERY-1")
        ctx.conn.commit()

    # A second query is now performed with the new acquired
    # (cur, conn). try/except logic must be present too.
    ctx.cur.execute("QUERY-2")
    ctx.conn.commit()

The solution implementation

Here is the proposed implementation of the DatabaseRetryContext:

from contextlib import AsyncExitStack
import logging
from types import TracebackType
from typing import Self, Type

from aiomysql import Connection, Pool


class DatabaseContext:
    """
    A MariaDB connection and cursor dedicated to a given HTTP API request
    """

    conn: Connection
    cur: DictCursor

    def __init__(self, conn: Connection, cur: DictCursor) -> None:
        self.conn = conn
        self.cur = cur


class DatabaseRetryContext(DatabaseContext):
    """
    A MariaDB connection and cursor dedicated to a given HTTP API request

    In case of OperationalError/InternalError during an SQL query, the `.reconnect()`
    method can be used to release the potentially broken (conn, cur) to the
    connection pool (which will properly manage it) and acquire a new (conn, cur),
    to retry the query.
    """

    pool: Pool
    _stack: AsyncExitStack

    def __init__(self, pool: Pool) -> None:
        self.pool = pool

    async def __aenter__(self) -> Self:
        """
        Acquire a (conn, cur) from the connection pool.
        """
        await self._acquire()
        return self

    async def __aexit__(
        self,
        exc_type: Type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        """
        Release the (conn, cur) to the connection pool.

        It's either the original (conn, cur), or new ones from .reconnect()
        """
        await self._stack.__aexit__(exc_type, exc_value, traceback)

    async def _acquire(self) -> None:
        """
        Acquire a connection and cursor from the connection pool

        Here we call __aenter__ and __aexit__ directly instead of a "with" block.
        It's done with AsyncExitStack to make sure the cleanup occurs in case of error, see
        this good explanation: https://stackoverflow.com/a/39172487

        """
        async with AsyncExitStack() as stack:
            self.conn = await stack.enter_async_context(self.pool.acquire())
            self.cur = await stack.enter_async_context(self.conn.cursor(DictCursor))
            self._stack = stack.pop_all()

    async def _release(self) -> None:
        """
        Release the (conn, cur) to the connection pool.

        The connection pool will take care of discarding broken connections (if closed by the
        SQL server, if underlying TCP connection is closed, if exception is raised...): See
        aiomysql/aiomysql/pool.py:Pool._fill_free_pool.
        """
        await self._stack.__aexit__(None, None, None)

    async def reconnect(self) -> None:
        """
        Release the (conn, cur) and acquire new ones

        This is triggered by the retry() function in case of OperationalError/InternalError.
        """
        await self._release()
        await self._acquire()

How to use it

Functions that perform database queries can takes a DatabaseRetryContext as first argument, and be decorated by a retry function. For example:

@retry
async def select_num_things(ctx: DatabaseRetryContext) -> int:
    await ctx.cur.execute("SELECT count(*) AS num FROM things")
    row = await ctx.cur.fetchone() 
    return row["num"]

If a function do not use connection pool or retry, a simple DatabaseContext can be used instead.

And a retry function could be for example:

import asyncio
from collections.abc import Awaitable
from typing import Callable, Concatenate, ParamSpec, TypeVar

from pymysql import OperationalError, InternalError
from fastapi import HTTPException


RETRIES_ERROR = (OperationalError, InternalError)

P = ParamSpec("P")
R = TypeVar("R")


def retry(
    f: Callable[Concatenate[DatabaseRetryContext, P], Awaitable[R]]
) -> Callable[Concatenate[DatabaseRetryContext, P], Awaitable[R]]:
    """
    Decorator to retry a function that performs database queries.

    The retries are limited to OperationError/InternalError, that is errors
    that does not comes from our code, but from network errors, database server
    memory, ...

    The decorated function must takes a "DatabaseRetryContext" as first argument: in case of failure
    due to DatabaseRetryContext.connection, the connection/cursor are re-acquired.

    A maximum of 2 retries are performed, after waiting 0.2 and 0.5 seconds.
    """

    async def f_with_retries(
        ctx: DatabaseRetryContext, *args: P.args, **kwargs: P.kwargs
    ) -> R:
        # First try
        try:
            return await f(ctx, *args, **kwargs)
        except RETRIES_ERROR:
            logger.exception(
                f"{f.__name__} failed (1/3), trying again after 0.2 seconds"
            )

        # Second try
        await asyncio.sleep(0.2)
        await ctx.reconnect()
        try:
            result = await f(ctx, *args, **kwargs)
            logger.info(f"{f.__name__} success (2/3)")
            return result
        except RETRIES_ERROR:
            logger.exception(
                f"{f.__name__} failed (2/3), trying again after 0.5 seconds"
            )

        # Third and last try
        await asyncio.sleep(0.5)
        await ctx.reconnect()
        try:
            result = await f(ctx, *args, **kwargs)
            logger.info(f"{f.__name__} success (3/3)")
            return result
        except RETRIES_ERROR:
            logger.exception(f"{f.__name__} failed (3/3), giving up")
            raise HTTPException(
                500,
                detail={
                    "type": "database_errror",
                    "msg": "A database error occured",
                },
            )

    return f_with_retries

Questions

What do you thing about the problem I'm trying to solve: retry with a new Connection/Cursor when OperationalError/InternalError occurs?

What do you thing about the proposed solution, is there are better/simpler approach?

What do you thing about a Pull Request to include DatabaseRetryContext in aiomysql: is the problem it solves common enough, and the solution generic enough? I would expect the retry function (maybe simplified version) to be in the documentation as possible example of how to use DatabaseRetryContext.

Thanks for readying!