pgbot/pgbotlib/dbstuff.py

import random
import psycopg


class DBConn:
    def __init__(self, *args, **kwargs) -> None:
        self.connection = psycopg.connect(*args, **kwargs)
        self.cursor = self.connection.cursor()

    def update(self, query: str, values: tuple) -> None:
        failure = None
        try:
            self.cursor.execute('SAVEPOINT sp1')
            self.cursor.execute(query, values)
        except Exception as e:
            failure = e
            self.cursor.execute('ROLLBACK TO SAVEPOINT sp1')
        else:
            self.cursor.execute('RELEASE SAVEPOINT sp1')
        self.connection.commit()
        if failure:
            raise failure

    def query_raw(self, query: str, values: tuple) -> list:
        self.cursor.execute(query, values)
        return self.cursor.fetchall()

    def query_random(self, query: str, values: tuple) -> str:
        result = self.query_raw(query, values)
        if not result:
            return None
        return random.choice(result)[0]