45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
# -*- 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'])
|