Psycopg2 auto reconnect inside a class

问题

I've got class to connect to my Database.

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()

How could I reconnect to database and handle psycopg2.OperationalError in my case?

回答1:

You could make a decorator that tries to reconnect when psycopg2.InterfaceError or psycopg2.OperationalError are raised.
That's just an example how it could work and probably needs adjustments:

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())

Output:

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,)

Note: _connect itself is not decorated, so this code assumes an initial connect always works!

来源：https://stackoverflow.com/questions/62094660/psycopg2-auto-reconnect-inside-a-class

标签

python-3.x

psycopg2