nxy/bot/storage.py

# -*- coding: utf-8 -*-
import os
import contextlib

import psycopg2
from psycopg2.extras import DictCursor

from . import Plugin, Bot

class DBConn:
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self._reconnect()

    def _reconnect(self):
        self._con = psycopg2.connect(*self.args, **self.kwargs)
        return self._con

    def commit(self, *args, **kwargs):
        try:
            return self._con.commit(*args, **kwargs)
        except psycopg2.InterfaceError:
            return self._reconnect().commit(*args, **kwargs)

    def rollback(self, *args, **kwargs):
        try:
            return self._con.rollback(*args, **kwargs)
        except psycopg2.InterfaceError:
            return self._reconnect().rollback(*args, **kwargs)

    @contextlib.contextmanager
    def cursor(self, *args, **kwargs):
        try:
            yield self._con.cursor(cursor_factory=DictCursor, *args, **kwargs)
        except psycopg2.InterfaceError:
            yield self._reconnect().cursor(cursor_factory=DictCursor, *args, **kwargs)


class Storage(Plugin):
    def __init__(self, bot: Bot):
        super().__init__(bot)
        self.bot.sql = self
        self.con = DBConn(os.environ['DATABASE_URI'])