nxy/bot/storage.py

45 lines
1.2 KiB
Python
Raw Normal View History

# -*- coding: utf-8 -*-
import os
import contextlib
import psycopg2
from psycopg2.extras import DictCursor
2017-08-22 15:43:48 +00:00
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):
2017-08-22 15:43:48 +00:00
def __init__(self, bot: Bot):
super().__init__(bot)
2017-08-22 15:43:48 +00:00
self.bot.sql = self
self.con = DBConn(os.environ['DATABASE_URI'])