Add retry logic for all management commands

extlinks/aggregates/tests.py

@@ -35,7 +35,19 @@
 from extlinks.organisations.models import Organisation
 
 
-class LinkAggregateCommandTest(TransactionTestCase):
+class BaseTransactionTest(TransactionTestCase):
+    @classmethod
+    def setUpClass(cls):
+        super(BaseTransactionTest, cls).setUpClass()
+        cls.tenacity_patcher = mock.patch('tenacity.nap.time')
+        cls.mock_tenacity = cls.tenacity_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        super(BaseTransactionTest, cls).tearDownClass()
+        cls.tenacity_patcher.stop()
+
+class LinkAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         # Creating one Collection
         self.organisation = OrganisationFactory(name="ACME Org")
@@ -233,7 +245,7 @@ def test_link_aggregate_with_argument_delete_org(self):
             call_command("fill_link_aggregates", collections=[new_collection.pk])
 
 
-class UserAggregateCommandTest(TransactionTestCase):
+class UserAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         # Creating one Collection
         self.organisation = OrganisationFactory(name="ACME Org")
@@ -481,7 +493,7 @@ def test_user_aggregate_with_argument_delete_org(self):
             call_command("fill_user_aggregates", collections=[new_collection.pk])
 
 
-class PageProjectAggregateCommandTest(TransactionTestCase):
+class PageProjectAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         # Creating one Collection
         self.organisation = OrganisationFactory(name="ACME Org")
@@ -739,7 +751,7 @@ def test_pageproject_aggregate_with_argument_delete_org(self):
             call_command("fill_pageproject_aggregates", collections=[new_collection.pk])
 
 
-class MonthlyLinkAggregateCommandTest(TransactionTestCase):
+class MonthlyLinkAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         self.organisation = OrganisationFactory(name="ACME Org")
         self.collection = CollectionFactory(name="ACME", organisation=self.organisation)
@@ -863,7 +875,7 @@ def test_specific_year_month(self):
             )
 
 
-class MonthlyUserAggregateCommandTest(TransactionTestCase):
+class MonthlyUserAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         self.organisation = OrganisationFactory(name="ACME Org")
         self.collection = CollectionFactory(name="ACME", organisation=self.organisation)
@@ -1051,7 +1063,7 @@ def test_specific_year_month(self):
             )
 
 
-class MonthlyPageProjectAggregateCommandTest(TransactionTestCase):
+class MonthlyPageProjectAggregateCommandTest(BaseTransactionTest):
     def setUp(self):
         self.organisation = OrganisationFactory(name="ACME Org")
         self.collection = CollectionFactory(name="ACME", organisation=self.organisation)
@@ -1259,7 +1271,7 @@ def test_specific_year_month(self):
             )
 
 
-class ArchiveLinkAggregatesCommandTest(TransactionTestCase):
+class ArchiveLinkAggregatesCommandTest(BaseTransactionTest):
     def setUp(self):
         self.organisation = OrganisationFactory(name="JSTOR")
         self.collection = CollectionFactory(
@@ -1518,7 +1530,7 @@ def test_archive_link_aggregates_no_dates(self, mock_swift_connection):
         self.assertEqual(LinkAggregate.objects.count(), 1)
 
 
-class ArchiveUserAggregatesCommandTest(TransactionTestCase):
+class ArchiveUserAggregatesCommandTest(BaseTransactionTest):
     def setUp(self):
         self.user = UserFactory(username="jonsnow")
         self.organisation = OrganisationFactory(name="JSTOR")
@@ -1782,7 +1794,7 @@ def test_archive_user_aggregates_no_dates(self, mock_swift_connection):
         self.assertEqual(UserAggregate.objects.count(), 1)
 
 
-class ArchivePageProjectAggregatesCommandTest(TransactionTestCase):
+class ArchivePageProjectAggregatesCommandTest(BaseTransactionTest):
     def setUp(self):
         self.page = "TestPage"
         self.project = "en.wikipedia.org"
@@ -2051,7 +2063,7 @@ def test_archive_pageproject_aggregates_no_dates(self, mock_swift_connection):
         self.assertEqual(PageProjectAggregate.objects.count(), 1)
 
 
-class UploadAllArchivedAggregatesCommandTest(TransactionTestCase):
+class UploadAllArchivedAggregatesCommandTest(BaseTransactionTest):
     @mock.patch.dict(
         os.environ,
         {

extlinks/common/management/commands/__init__.py

@@ -4,21 +4,47 @@
 import logging
 from os import remove
 from os.path import basename
+import sys
+from tenacity import Retrying, stop_after_attempt, wait_exponential
+
+logger = logging.getLogger("django")
 
 
 class BaseCommand(DjangoBaseCommand):
     """
-    Django BaseCommand wrapper that adds file locks to management commands
+    Django BaseCommand wrapper that adds
+        - file locks
+        - up to 5 retries with exponential backoff
     """
 
+    def __init__(self, *args, **options):
+        super().__init__(*args, **options)
+        self.filename = basename(inspect.getfile(self.__class__))
+        self.e = None
+
+    def retry_log(self, retry_state):
+        logger.warning(f"{self.filename} attempt {retry_state.attempt_number} failed")
+
     def handle(self, *args, **options):
-        lockname = basename(inspect.getfile(self.__class__))
         # Use a lockfile to prevent overruns.
-        lockfile = "/tmp/{}.lock".format(lockname)
+        logger.info(f"Executing {self.filename}")
+        lockfile = f"/tmp/{self.filename}.lock"
         lock = FileLock(lockfile)
         lock.acquire()
         try:
-            self._handle(*args, **options)
-        finally:
-            lock.release()
-            remove(lockfile)
+            for attempt in Retrying(
+                after=self.retry_log,
+                reraise=True,
+                stop=stop_after_attempt(5),
+                wait=wait_exponential(multiplier=1, min=60, max=300),
+            ):
+                with attempt:
+                    self._handle(*args, **options)
+        except Exception as e:
+            logger.warning(f"Retries exhausted for {self.filename}")
+            logger.error(e)
+            self.e = e
+        lock.release()
+        remove(lockfile)
+        if self.e is not None:
+            raise self.e

extlinks/links/management/commands/linksearchtotal_collect.py

@@ -32,16 +32,12 @@ def _handle(self, *args, **options):
         total_links_dictionary = {}
         for i, language in enumerate(wiki_list_data):
             logger.info(f"connecting to db {language}")
-            try:
-                db = MySQLdb.connect(
-                    host=f"{language}wiki.analytics.db.svc.wikimedia.cloud",
-                    user=os.environ["REPLICA_DB_USER"],
-                    passwd=os.environ["REPLICA_DB_PASSWORD"],
-                    db=f"{language}wiki_p",
-                )
-            except MySQLdb.OperationalError as e:
-                logger.error(str(e))
-                sys.exit(1)
+            db = MySQLdb.connect(
+                host=f"{language}wiki.analytics.db.svc.wikimedia.cloud",
+                user=os.environ["REPLICA_DB_USER"],
+                passwd=os.environ["REPLICA_DB_PASSWORD"],
+                db=f"{language}wiki_p",
+            )
 
             cur = db.cursor()
 

extlinks/links/tests.py

@@ -23,8 +23,19 @@
 from .helpers import link_is_tracked, reverse_host
 from .models import URLPattern, LinkEvent
 
-
-class LinksHelpersTest(TestCase):
+class BaseTest(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super(BaseTest, cls).setUpClass()
+        cls.tenacity_patcher = mock.patch('tenacity.nap.time')
+        cls.mock_tenacity = cls.tenacity_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        super(BaseTest, cls).tearDownClass()
+        cls.tenacity_patcher.stop()
+
+class LinksHelpersTest(BaseTest):
     def test_reverse_host(self):
         """
         Test the reverse_host function in helpers.py.
@@ -58,7 +69,7 @@ def test_get_organisation(self):
         self.assertEqual(new_link.get_organisation, organisation2)
 
 
-class LinksTrackedTest(TestCase):
+class LinksTrackedTest(BaseTest):
     def setUp(self):
         _ = URLPatternFactory(url="test.com")
 
@@ -123,7 +134,7 @@ def test_link_is_tracked_false_other_proxy(self):
         )
 
 
-class URLPatternModelTest(TestCase):
+class URLPatternModelTest(BaseTest):
     def test_get_proxied_url_1(self):
         """
         Test that URLPattern.get_proxied_url transforms a URL correctly
@@ -140,7 +151,7 @@ def test_get_proxied_url_2(self):
         self.assertEqual(test_urlpattern.get_proxied_url, "platform-almanhal-com")
 
 
-class LinkEventsCollectCommandTest(TestCase):
+class LinkEventsCollectCommandTest(BaseTest):
     def setUp(self):
         self.organisation1 = OrganisationFactory(name="JSTOR")
         self.collection1 = CollectionFactory(
@@ -1250,7 +1261,7 @@ def test_fix_proxy_linkevents_on_user_list_command(self):
         self.assertEqual(PageProjectAggregate.objects.count(), 2)
 
 
-class UploadAllArchivedTestCase(TestCase):
+class UploadAllArchivedTestCase(BaseTest):
     @mock.patch.dict(
         os.environ,
         {

extlinks/settings/logging.py

@@ -25,7 +25,7 @@
         },
         "handlers": {
             "nodebug_console": {
-                "level": "WARNING",
+                "level": "INFO",
                 "filters": ["require_debug_false"],
                 "class": "logging.StreamHandler",
                 "formatter": "django.server",

requirements/django.txt

@@ -8,6 +8,7 @@ pymemcache==4.0.0
 python-dotenv==1.1.0
 sentry-sdk==2.10.0
 sseclient==0.0.27
+tenacity==9.1.2
 time_machine==2.14.2
 python-swiftclient>=4.6.0,<5.0.0
 keystoneauth1>=5.9.2,<6.0.0