Psycopg2在类内自动重新连接

Question

我有课要连接我的数据库。

import psycopg2, psycopg2.extensions
from parseini import config
import pandas as pd, pandas.io.sql as sqlio


class MyDatabase:
    def __init__(self, name='mydb.ini'):
        self.params = config(filename=name)
        self.my_connection = psycopg2.connect(**self.params)
        self.my_cursor = self.my_connection.cursor()

    def fetch_all_as_df(self, sql_statement):
        return sqlio.read_sql_query(sql_statement, self.my_connection)

    def df_to_sql(self, df):
        table = 'sometable'
        return sqlio.to_sql(df, table, self.my_connection)

    def __del__(self):
        self.my_cursor.close()
        self.my_connection.close()

在我的情况下，我如何重新连接到数据库并处理phocopg2.OperationalError？

Answer 1

您可以制作一个装饰器，在引发psycopg2.InterfaceError或psycopg2.OperationalError时尝试重新连接。

这只是一个例子，它可以工作，可能需要调整：

import time
from functools import wraps
import psycopg2, psycopg2.extensions


def retry(fn):
    @wraps(fn)
    def wrapper(*args, **kw):
        cls = args[0]
        for x in range(cls._reconnectTries):
            print(x, cls._reconnectTries)
            try:
                return fn(*args, **kw)
            except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
                print ("\nDatabase Connection [InterfaceError or OperationalError]")
                print ("Idle for %s seconds" % (cls._reconnectIdle))
                time.sleep(cls._reconnectIdle)
                cls._connect()
    return wrapper


class MyDatabase:
    _reconnectTries = 5
    _reconnectIdle = 2  # wait seconds before retying

    def __init__(self, name='mydb.ini'):
        self.my_connection = None
        self.my_cursor = None
        self.params = config(filename=name)
        self._connect()

    def _connect(self):
        self.my_connection = psycopg2.connect(**self.params)
        self.my_cursor = self.my_connection.cursor()

    @retry
    def fetch_all_as_df(self, sql_statement):
        return sqlio.read_sql_query(sql_statement, self.my_connection)

    @retry
    def dummy(self):
        self.my_cursor.execute('select 1+2 as result')
        return self.my_cursor.fetchone()

    @retry
    def df_to_sql(self, df):
        table = 'sometable'
        return sqlio.to_sql(df, table, self.my_connection)

    def __del__(self):
        # Maybe there is a connection but no cursor, whatever close silently!
        for c in (self.my_cursor, self.my_connection):
            try:
                c.close()
            except:
                pass


db = MyDatabase()
time.sleep(30)  # some time to shutdown the database
print(db.dummy())

输出：

Database Connection [InterfaceError or OperationalError]
Idle for 2 seconds

Database Connection [InterfaceError or OperationalError]
Idle for 2 seconds

Database Connection [InterfaceError or OperationalError]
Idle for 2 seconds

Database Connection [InterfaceError or OperationalError]
Idle for 2 seconds
(3,)

注意：_connect本身没有修饰，所以这段代码假设初始连接总是工作！