Merge pull request #8 from zhangzhw8/feat_support_retry_times

feat: support MAX_DBCONN_RETRY_TIMES

Author: jdelic

README.rst

@@ -68,6 +68,17 @@ Both signals send a parameter ``dbwrapper`` which points to the current instance
 of ``django.db.backends.base.BaseDatabaseWrapper`` which allows the signal
 receiver to act on the database connection.
 
+Settings
+-------
+Here’s a list of settings available in django-dbconn-retry and their default values.
+You can change the value in your ``settings.py``.
+
+===========================  ==================================================
+Setting                       Description
+===========================  ==================================================
+``MAX_DBCONN_RETRY_TIMES``   Default: `1`
+                             The max times which django-dbconn-retry will try.
+===========================  ==================================================
 
 License
 =======

django_dbconn_retry/__init__.py

@@ -1,7 +1,6 @@
 # -* encoding: utf-8 *-
 import django
 
-
 from django_dbconn_retry.apps import pre_reconnect, post_reconnect, monkeypatch_django, DjangoIntegration
 
 

django_dbconn_retry/apps.py

@@ -1,6 +1,7 @@
 import logging
 
 from django.apps.config import AppConfig
+from django.conf import settings
 from django.db.backends.base import base as django_db_base
 from django.dispatch import Signal
 
@@ -34,6 +35,8 @@
 
 def monkeypatch_django() -> None:
     def ensure_connection_with_retries(self: django_db_base.BaseDatabaseWrapper) -> None:
+        self._max_dbconn_retry_times = getattr(settings, "MAX_DBCONN_RETRY_TIMES", 1)
+
         if self.connection is not None and hasattr(self.connection, 'closed') and self.connection.closed:
             _log.debug("failed connection detected")
             self.connection = None
@@ -45,15 +48,22 @@ def ensure_connection_with_retries(self: django_db_base.BaseDatabaseWrapper) ->
                     self.connect()
                 except Exception as e:
                     if isinstance(e, _operror_types):
-                        if hasattr(self, "_connection_retries") and self._connection_retries >= 1:
+                        if (
+                                hasattr(self, "_connection_retries") and
+                                self._connection_retries >= self._max_dbconn_retry_times
+                        ):
                             _log.error("Reconnecting to the database didn't help %s", str(e))
                             del self._in_connecting
                             post_reconnect.send(self.__class__, dbwrapper=self)
                             raise
                         else:
                             _log.info("Database connection failed. Refreshing...")
                             # mark the retry
-                            self._connection_retries = 1
+                            try:
+                                self._connection_retries += 1
+                            except AttributeError:
+                                self._connection_retries = 1
+
                             # ensure that we retry the connection. Sometimes .closed isn't set correctly.
                             self.connection = None
                             del self._in_connecting

django_dbconn_retry/tests/__init__.py

@@ -1,4 +1,5 @@
 # -* encoding: utf-8 *-
+import random
 import sys
 import logging
 
@@ -10,7 +11,7 @@
 
 from django.db.backends.base.base import BaseDatabaseWrapper
 from django.db import connection, OperationalError, transaction
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 
 logging.basicConfig(stream=sys.stderr)
@@ -24,6 +25,8 @@ class FullErrorTests(TestCase):
     database connections reliably fail. If I had been able to think of
     a better way, I'd have used it.
     """
+    max_dbconn_retry_times = random.randint(1, 100)
+
     def test_getting_root(self) -> None:
         self.client.get('/')
 
@@ -38,26 +41,31 @@ def tearDown(self) -> None:
         BaseDatabaseWrapper.connect = self.s_connect
         del BaseDatabaseWrapper.connection
 
-    def test_prehook(self) -> None:
-        cb = Mock(name='pre_reconnect_hook')
-        ddr.pre_reconnect.connect(cb)
+    def do_assert(self, cb):
         self.assertRaises(OperationalError, connection.ensure_connection)
         self.assertTrue(cb.called)
+        self.assertEqual(connection._connection_retries, self.max_dbconn_retry_times)
         del connection._connection_retries
 
+    @override_settings(MAX_DBCONN_RETRY_TIMES=max_dbconn_retry_times)
+    def test_prehook(self) -> None:
+        cb = Mock(name='pre_reconnect_hook')
+        ddr.pre_reconnect.connect(cb)
+        self.do_assert(cb)
+
+    @override_settings(MAX_DBCONN_RETRY_TIMES=max_dbconn_retry_times)
     def test_posthook(self) -> None:
         cb = Mock(name='post_reconnect_hook')
         ddr.post_reconnect.connect(cb)
-        self.assertRaises(OperationalError, connection.ensure_connection)
-        self.assertTrue(cb.called)
-        del connection._connection_retries
+        self.do_assert(cb)
 
 
 def fix_connection(sender: type, *, dbwrapper: BaseDatabaseWrapper, **kwargs: Any) -> None:
     dbwrapper.connect = dbwrapper.s_connect
 
 
 class ReconnectTests(TestCase):
+
     @classmethod
     def tearDownClass(cls) -> None:
         return
@@ -67,10 +75,7 @@ def test_ensure_closed(self) -> None:
         connection.close()
         self.assertFalse(connection.is_usable())  # should be true after setUp
 
-    def test_prehook(self) -> None:
-        cb = Mock(name='pre_reconnect_hook')
-        ddr.pre_reconnect.connect(fix_connection)
-        ddr.pre_reconnect.connect(cb)
+    def do_assert(self, cb):
         from django.db import connection
         connection.close()
         connection.s_connect = connection.connect
@@ -80,17 +85,16 @@ def test_prehook(self) -> None:
         ReconnectTests.cls_atomics['default'].__enter__()
         self.assertTrue(cb.called)
         self.assertTrue(connection.is_usable())
+        self.assertEqual(connection._connection_retries, 0)
+
+    def test_prehook(self) -> None:
+        cb = Mock(name='pre_reconnect_hook')
+        ddr.pre_reconnect.connect(fix_connection)
+        ddr.pre_reconnect.connect(cb)
+        self.do_assert(cb)
 
     def test_posthook(self) -> None:
         cb = Mock(name='post_reconnect_hook')
         ddr.pre_reconnect.connect(fix_connection)
         ddr.post_reconnect.connect(cb)
-        from django.db import connection
-        connection.close()
-        connection.s_connect = connection.connect
-        connection.connect = Mock(side_effect=OperationalError('reconnect testing'))
-        connection.ensure_connection()
-        ReconnectTests.cls_atomics['default'] = transaction.atomic(using='default')
-        ReconnectTests.cls_atomics['default'].__enter__()
-        self.assertTrue(cb.called)
-        self.assertTrue(connection.is_usable())
+        self.do_assert(cb)