Aggregating psycopg2 Errors under a Context Manager

Question

I have various functions which attempt to insert into a PostgreSQL database. In an attempt to remove redundancy for common errors that I was occurring I have created a context manager.

Solution

from contextlib import contextmanager

import psycopg2


class DomainSpecificException(Exception):
    ...


@contextmanager
def raise_on_insert(table_name: str) -> None:
    try:
        yield
    except psycopg2.errors.InFailedSqlTransaction:
        raise DomainSpecificException(
            description=(
                f'Failed to insert into: {table_name} because InFailedSqlTransaction, this is usually a sign that the upstream connection has spoiled.'
            )
        )
    except psycopg2.errors.NotNullViolation:
        raise DomainSpecificException(
            description=(
                f'Failed to insert into: {table_name} because NotNullViolation for column.'
            )
        )
    except psycopg2.errors.UniqueViolation:
        raise DomainSpecificException(
            description=(
                f'Failed to insert into: {table_name} because UniqueViolation for column'
            ),
        )

    except psycopg2.errors.ForeignKeyViolation:
        raise DomainSpecificException(
            description=(
                f'Failed to insert into: {table_name} because ForeignKeyViolation for column'
            )
        )

Usage

with raise_on_insert(table_name='table_1'):
  ...  # Insert into db

Edit

I have improved the code as follows:

from contextlib import contextmanager

from psycopg2.errors import InFailedSqlTransaction, NotNullViolation, UniqueViolation, ForeignKeyViolation


INSERT_EXCEPTIONS: tuple[Exception, ...] = (
    InFailedSqlTransaction,
    NotNullViolation,
    UniqueViolation,
    ForeignKeyViolation
)


class DomainSpecificException(Exception):
    pass


@contextmanager
def raise_on_insert(table_name: str) -> None:
    try:
        yield
    except INSERT_EXCEPTIONS as err:
        raise DomainSpecificException(description=f'Failed to insert into: {table_name} because {type(err).__name__}')

Answer 1

We have two implementations here, one better than the other. Cool.

Kudos on annotating and passing mypy lint. And defining app-specific exceptions is very helpful to callers who want coarse- or fine-grained try / catches.

---

In the 1st implementation, my chief critique is DRY, it just seems a bit chatty. No biggie.

I do appreciate that the InFailedSqlTransaction case is trying hard to be helpful, to offer a diagnostic error that helps some poor maintainer figure out the root cause and the fix.

I still have some minor reservations, but they show up in both implementations.

---

The 2nd implementation is very nice. Short, sweet, clear code.

Here is one concern. I don't know the right answer here, I will just voice the concern. Maybe four exceptions is the right number? But maybe other bad things can go wrong? I'm just wondering if psycopg2 maybe offers some umbrella exception that you would also like to throw in there.

Taking a step back, it's not clear from these source documents exactly what your error handling objectives are. The code I'm reading mostly seems focused on exposing details to a handler that's not very far up the call stack which can maybe recover / ignore / retry when a low-level bad thing happens, so the high-level business goal still ends up getting accomplished. We're missing some review context here. I would have loved the opportunity to read calling code that has a try, or unit tests which provoke each of the four exceptions.

Summarize this as: "do we need to exhaustively cover all failure modes?"

Here's another concern. The description parameter is helpful to humans. I'm sad that we didn't wrap the original err exception. That is, we've discarded traceback details like the function and source line number that triggered the DB error.

Maybe a single exception type makes sense, because all your callers have similar try / except behavior? But consider creating four child exception classes, so callers can catch just "the one thing" they know how to recover from, if they want to.

Final concern is very minor. We offer a maintainer a little less diagnostic advice in the InFailedSqlTransaction case. Maybe remedy that by tacking on a # comment in that list of four.

---

This code achieves its design objectives.

I would be willing to delegate or accept maintenance tasks on this codebase.

Answer 2

Based on comments.

from contextlib import contextmanager

from psycopg2.errors import DatabaseError


class DomainSpecificException(Exception):
    pass


@contextmanager
def raise_on_insert(table_name: str) -> None:
    try:
        yield
    except DatabaseError as err:
        raise DomainSpecificException(description=f'Failed to insert into: {table_name} because {type(err).__name__}') from err