From ed276a1cb2b036650439d031f76a3eee717180ce Mon Sep 17 00:00:00 2001
From: Radko Krkos <krkos@cesnet.cz>
Date: Tue, 21 May 2019 12:20:00 +0200
Subject: [PATCH 2/3] Fix: Fixed enricher failure on DB restart.

* The enricher plugin's WhoIs module is reloading configuration from the DB
  periodically, what results in delayed failure after DB restart. This allows
  for reconnection in that case, effectively eliminating the issue.
* SQLAlchemy's Query class is extended to RetryingQuery and that one is used
  instead for the DB sessions.

(Redmine issue: #4569)
---
 lib/mentat/services/sqlstorage.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/mentat/services/sqlstorage.py b/lib/mentat/services/sqlstorage.py
index 4935787d..40320c6d 100644
--- a/lib/mentat/services/sqlstorage.py
+++ b/lib/mentat/services/sqlstorage.py
@@ -27,6 +27,7 @@ __credits__ = "Pavel Kácha <pavel.kacha@cesnet.cz>, Andrea Kropáčová <andrea
 
 import copy
 import sqlalchemy
+from sqlalchemy.orm import Query
 
 #
 # Custom libraries
@@ -38,6 +39,20 @@ from mentat.datatype.sqldb import MODEL
 _MANAGER = None
 
 
+class RetryingQuery(Query):
+    """
+    An override of SQLAlchemy's Query class, allowing for recovery from a lost DB
+    connection.
+    """
+    def _execute_and_instances(self, querycontext):
+        for _ in range(2):
+            try:
+                return super()._execute_and_instances(querycontext)
+            except sqlalchemy.exc.OperationalError:
+                self.session.close()
+                continue
+
+
 class StorageService:
     """
     Proxy object for working with persistent SQL storages. Maintains and provides
@@ -54,7 +69,7 @@ class StorageService:
         """
         self.dbengine     = sqlalchemy.engine_from_config(enginecfg, prefix = '')
         self.sessionmaker = sqlalchemy.orm.sessionmaker(bind = self.dbengine)
-        self.session      = self.sessionmaker()
+        self.session      = self.sessionmaker(query_cls = RetryingQuery)
 
     def __del__(self):
         self.close()
-- 
2.14.5