Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How do I promote PostgreSQL warnings to exceptions in psycopg2?

From the PostgreSQL docs on BEGIN:

Issuing BEGIN when already inside a transaction block will provoke a warning message. The state of the transaction is not affected.

How can I make psycopg2 raise an exception on any such warning?

like image 319
planetp Avatar asked Oct 30 '22 00:10

planetp


1 Answers

I am very far from being psycopg2 or Postgres expert, and, I am sure there is a better solution to increase the warning level, but here is something that worked for me - a custom cursor which looks into connection notices and, if there is something there - it throws an exception. The implementation itself is for education purposes mostly - I am sure it needs to be adjusted to work in your use case:

import psycopg2

# this "cursor" class needs to be used as a base for custom cursor classes
from psycopg2.extensions import cursor  

class ErrorThrowingCursor(cursor):
    def __init__(self, conn, *args, **kwargs):
        self.conn = conn
        super(ErrorThrowingCursor, self).__init__(*args, **kwargs)

    def execute(self, query, vars=None):
        result = super(ErrorThrowingCursor, self).execute(query, vars)

        for notice in self.conn.notices:
            level, message = notice.split(": ")
            if level == "WARNING":
                raise psycopg2.Warning(message.strip())

        return result

Usage sample:

conn = psycopg2.connect(user="user", password="secret")
cursor = conn.cursor(conn, cursor_factory=ErrorThrowingCursor)

This would throw an exception (of a psycopg2.Warning type) if a warning was issued after a query execution. Sample:

psycopg2.Warning: there is already a transaction in progress
like image 119
alecxe Avatar answered Nov 09 '22 22:11

alecxe