Reconnect to DB on lost connection

zuevmaxim · Jul 27, 2024 · 37db378 · 37db378
1 parent 09ab094
commit 37db378
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/storage/postgres.py b/storage/postgres.py
@@ -14,7 +14,7 @@ def __init__(self,
                  port: str,
                  database: str,
                  reconnect: bool = True,
-                 limit_reties: int = 0):
+                 limit_reties: int = 3):
         self.user = user
         self.password = password
         self.host = host
@@ -26,13 +26,13 @@ def __init__(self,
 
         self._connection = None
 
-        self.init()
+        self.connect()
 
     def connect(self, retry_counter=0):
         if not self._connection:
             try:
                 self._connection = psycopg2.connect(user=self.user, password=self.password, host=self.host,
-                                                    port=self.port, database=self.database, connect_timeout=3, )
+                                                    port=self.port, database=self.database, connect_timeout=3)
                 retry_counter = 0
                 self._connection.autocommit = True
                 return self._connection
@@ -41,7 +41,7 @@ def connect(self, retry_counter=0):
                     raise error
                 else:
                     retry_counter += 1
-                    logging.error("Got error {}. Reconnecting {}.".format(str(error).strip(), retry_counter))
+                    logging.error(f"Got error {str(error).strip()}. Reconnecting {retry_counter}.")
                     time.sleep(5)
                     self.connect(retry_counter)
             except (Exception, psycopg2.Error) as error:
@@ -62,7 +62,7 @@ def execute(self, query, vars=None, retry_counter=0):
                 raise error
             else:
                 retry_counter += 1
-                logging.error("Got error {}. Retrying {}.".format(str(error).strip(), retry_counter))
+                logging.error(f"Got error {str(error).strip()}. Retrying {retry_counter}.")
                 time.sleep(1)
                 self.reset()
                 return self.execute(query, vars, retry_counter)
@@ -79,9 +79,6 @@ def close(self):
             logging.info("PostgreSQL connection is closed.")
         self._connection = None
 
-    def init(self):
-        self.connect()
-
 
 db = DB(os.environ["POSTGRES_USER"],
         os.environ["POSTGRES_PASSWORD"],