Psycopg2 auto reconnect inside a class

后端 未结 1 609
既然无缘
既然无缘 2021-01-26 06:10

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 M         


        
相关标签:
1条回答
  • 2021-01-26 06:40

    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!

    0 讨论(0)
提交回复
热议问题