From b374aab2d9524c2fca9126980321fd7034e4144a Mon Sep 17 00:00:00 2001
From: Jonathan Zernik <yzernik@gmail.com>
Date: Sat, 16 Oct 2021 02:04:28 -0700
Subject: [PATCH] Add init with retries method to squeak db (#1616)

* Add init with retries method to squeak db

* Use db init_with_retries method in squeak node class

* Add unit test for db init with retries method

* Raise exception with message on init with retry fail
---
 squeaknode/db/squeak_db.py     | 23 ++++++++++++++++++++++
 squeaknode/node/squeak_node.py |  9 +--------
 tests/db/test_squeak_db.py     | 35 ++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/squeaknode/db/squeak_db.py b/squeaknode/db/squeak_db.py
index fb169343..fe761b10 100644
--- a/squeaknode/db/squeak_db.py
+++ b/squeaknode/db/squeak_db.py
@@ -53,6 +53,8 @@ from squeaknode.db.models import Models
 
 MAX_INT = 999999999999
 MAX_HASH = b'\xff' * 32
+INIT_NUM_RETRIES = 10
+INIT_RETRY_INTERVAL_S = 1
 
 
 logger = logging.getLogger(__name__)
@@ -74,6 +76,27 @@ class SqueakDb:
         logger.debug("SqlAlchemy version: {}".format(sqlalchemy.__version__))
         run_migrations(self.engine)
 
+    def init_with_retries(
+            self,
+            num_retries=INIT_NUM_RETRIES,
+            retry_interval_s=INIT_RETRY_INTERVAL_S,
+    ):
+        """ Try repeatedly to init the database.
+
+        Raises exception if db init fails more than `num_retries` times.
+        """
+        n = 0
+        while True:
+            try:
+                self.init()
+                return
+            except Exception:
+                logger.exception("Failed to initialize database.")
+                n += 1
+                if n >= num_retries:
+                    raise Exception("Failed to initialize database.")
+                time.sleep(retry_interval_s)
+
     @property
     def squeaks(self):
         return self.models.squeaks
diff --git a/squeaknode/node/squeak_node.py b/squeaknode/node/squeak_node.py
index aad9e0e1..0380ce61 100644
--- a/squeaknode/node/squeak_node.py
+++ b/squeaknode/node/squeak_node.py
@@ -20,7 +20,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 import logging
-import time
 
 from squeak.params import SelectParams
 
@@ -112,13 +111,7 @@ class SqueakNode:
             connection_string))
         engine = get_engine(connection_string)
         self.squeak_db = SqueakDb(engine)
-        for _ in range(10):
-            try:
-                self.squeak_db.init()
-                break
-            except Exception:
-                logger.exception("Failed to initialize database.")
-                time.sleep(10)
+        self.squeak_db.init_with_retries()
 
     def initialize_lightning_client(self):
         # load the lightning client
diff --git a/tests/db/test_squeak_db.py b/tests/db/test_squeak_db.py
index 5672185a..d607ea0a 100644
--- a/tests/db/test_squeak_db.py
+++ b/tests/db/test_squeak_db.py
@@ -19,6 +19,7 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
+import mock
 import pytest
 from sqlalchemy import create_engine
 
@@ -161,6 +162,40 @@ def inserted_peer_id(squeak_db, peer):
     yield squeak_db.insert_peer(peer)
 
 
+def test_init_with_retries(squeak_db):
+    with mock.patch.object(squeak_db, 'init', autospec=True) as mock_init, \
+            mock.patch('squeaknode.db.squeak_db.time.sleep', autospec=True) as mock_sleep:
+        ret = squeak_db.init_with_retries(num_retries=5, retry_interval_s=100)
+
+        assert ret is None
+        mock_init.assert_called_once_with()
+        assert mock_sleep.call_count == 0
+
+
+def test_init_with_retries_fail_once(squeak_db):
+    with mock.patch.object(squeak_db, 'init', autospec=True) as mock_init, \
+            mock.patch('squeaknode.db.squeak_db.time.sleep', autospec=True) as mock_sleep:
+        mock_init.side_effect = [Exception('some db error'), None]
+        ret = squeak_db.init_with_retries(num_retries=5, retry_interval_s=100)
+
+        assert ret is None
+        mock_init.call_count == 2
+        assert mock_sleep.call_count == 1
+
+
+def test_init_with_retries_fail_many_times(squeak_db):
+    with mock.patch.object(squeak_db, 'init', autospec=True) as mock_init, \
+            mock.patch('squeaknode.db.squeak_db.time.sleep', autospec=True) as mock_sleep:
+        mock_init.side_effect = [Exception('some db error')] * 5
+
+        with pytest.raises(Exception) as excinfo:
+            squeak_db.init_with_retries(num_retries=5, retry_interval_s=100)
+        assert "Failed to initialize database." in str(excinfo.value)
+
+        mock_init.call_count == 5
+        assert mock_sleep.call_count == 4
+
+
 def test_get_squeak(squeak_db, squeak, inserted_squeak_hash):
     retrieved_squeak = squeak_db.get_squeak(inserted_squeak_hash)