Create RepeaterStubs

corehq/motech/repeaters/dbaccessors.py

@@ -185,6 +185,17 @@ def _iter_repeat_records_by_repeater(domain, repeater_id, chunk_size,
             yield doc['id']
 
 
+def iter_repeaters():
+    from .models import Repeater
+
+    for result in Repeater.get_db().view(
+        'repeaters/repeaters',
+        include_docs=True,
+        reduce=False,
+    ).all():
+        yield Repeater.wrap(result['doc'])
+
+
 def get_repeat_records_by_payload_id(domain, payload_id):
     repeat_records = get_sql_repeat_records_by_payload_id(domain, payload_id)
     if repeat_records:

corehq/motech/repeaters/migration_functions.py

@@ -0,0 +1,27 @@
+"""
+Functions used by RunPython operations in migrations.
+
+.. note::
+   Placing them in a separate module makes it easier to import them in
+   unittests.
+
+"""
+from django.db import transaction, IntegrityError
+
+from corehq.motech.repeaters.dbaccessors import iter_repeaters
+from corehq.motech.repeaters.models import RepeaterStub
+
+
+def create_repeaterstubs(apps, schema_editor):
+    for repeater in iter_repeaters():
+        with transaction.atomic():
+            try:
+                RepeaterStub.objects.create(
+                    domain=repeater.domain,
+                    repeater_id=repeater.get_id,
+                    is_paused=repeater.paused,
+                )
+            except IntegrityError:
+                # It's faster to violate uniqueness constraint and ask
+                # forgiveness than to use `.get()` or `.exists()` first.
+                pass

corehq/motech/repeaters/migrations/0004_create_repeaterstubs.py

@@ -0,0 +1,16 @@
+from django.db import migrations
+
+from ..migration_functions import create_repeaterstubs
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('repeaters', '0003_migrate_connectionsettings'),
+    ]
+
+    operations = [
+        migrations.RunPython(create_repeaterstubs,
+                             reverse_code=migrations.RunPython.noop,
+                             elidable=True),
+    ]

corehq/motech/repeaters/migrations/0005_repeater_uniq.py

@@ -0,0 +1,16 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('repeaters', '0004_create_repeaterstubs'),
+    ]
+
+    operations = [
+        migrations.AddConstraint(
+            model_name='repeaterstub',
+            constraint=models.UniqueConstraint(fields=('repeater_id',),
+                                               name='one_to_one_repeater'),
+        ),
+    ]

corehq/motech/repeaters/models.py

@@ -76,8 +76,8 @@
 from django.utils.translation import ugettext_lazy as _
 
 from couchdbkit.exceptions import ResourceConflict, ResourceNotFound
-from memoized import memoized
-from requests.exceptions import ConnectionError, Timeout, RequestException
+from memoized import memoized, memoized_property
+from requests.exceptions import ConnectionError, RequestException, Timeout
 
 from casexml.apps.case.xml import LEGAL_VERSIONS, V2
 from couchforms.const import DEVICE_LOG_XMLNS
@@ -120,6 +120,7 @@
 from corehq.util.couch import stale_ok
 from corehq.util.metrics import metrics_counter
 from corehq.util.quickcache import quickcache
+from corehq.util.urlsanitize.urlsanitize import PossibleSSRFAttempt
 
 from .const import (
     MAX_ATTEMPTS,
@@ -151,7 +152,6 @@
     ShortFormRepeaterJsonPayloadGenerator,
     UserPayloadGenerator,
 )
-from ...util.urlsanitize.urlsanitize import PossibleSSRFAttempt
 
 
 def log_repeater_timeout_in_datadog(domain):
@@ -176,6 +176,12 @@ def log_repeater_success_in_datadog(domain, status_code, repeater_type):
 
 class RepeaterStubManager(models.Manager):
 
+    def get_or_none(self, *args, **kwargs):
+        try:
+            return self.get(*args, **kwargs)
+        except self.model.DoesNotExist:
+            return None
+
     def all_ready(self):
         """
         Return all RepeaterStubs ready to be forwarded.
@@ -212,6 +218,10 @@ class Meta:
             models.Index(fields=['domain']),
             models.Index(fields=['repeater_id']),
         ]
+        constraints = [
+            models.UniqueConstraint(fields=['repeater_id'],
+                                    name='one_to_one_repeater')
+        ]
 
     @property
     @memoized
@@ -299,6 +309,28 @@ def _get_connection_settings(self):
     def name(self):
         return self.connection_settings.name
 
+    def save(self, **kwargs):
+        is_new = not self._id
+        super().save(**kwargs)
+        if is_new:
+            RepeaterStub.objects.create(
+                domain=self.domain,
+                repeater_id=self._id,
+                is_paused=self.paused,
+            )
+
+    def delete(self):
+        if self.repeater_stub.id:
+            self.repeater_stub.delete()
+        super().delete()
+
+    @memoized_property
+    def repeater_stub(self):
+        return RepeaterStub.objects.get(
+            domain=self.domain,
+            repeater_id=self._id,
+        )
+
     @classmethod
     def available_for_domain(cls, domain):
         """Returns whether this repeater can be used by a particular domain
@@ -448,14 +480,24 @@ def retire(self):
             self['base_doc'] += DELETED_SUFFIX
         self.paused = False
         self.save()
+        # Deleting RepeaterStub will cascade-delete SQLRepeatRecords
+        # and SQLRepeatRecordAttempts too. (RequestLog will still keep a
+        # record of all send attempts.)
+        self.repeater_stub.delete()
+        # NOTE: Undeleting a Repeater needs to include creating a
+        #       RepeaterStub for it.
 
     def pause(self):
         self.paused = True
         self.save()
+        self.repeater_stub.is_paused = True
+        self.repeater_stub.save()
 
     def resume(self):
         self.paused = False
         self.save()
+        self.repeater_stub.is_paused = False
+        self.repeater_stub.save()
 
     def get_url(self, repeat_record):
         # to be overridden

corehq/motech/repeaters/tests/test_models.py

@@ -3,6 +3,7 @@
 from uuid import uuid4
 
 from django.conf import settings
+from django.db import IntegrityError, transaction
 from django.test import SimpleTestCase, TestCase
 from django.utils import timezone
 
@@ -21,6 +22,7 @@
     RECORD_PENDING_STATE,
     RECORD_SUCCESS_STATE,
 )
+from ..migration_functions import create_repeaterstubs
 from ..models import (
     FormRepeater,
     RepeaterStub,
@@ -49,13 +51,9 @@ def setUp(self):
             url='https://www.example.com/api/',
         )
         self.repeater.save()
-        self.repeater_stub = RepeaterStub.objects.create(
-            domain=DOMAIN,
-            repeater_id=self.repeater.get_id,
-        )
+        self.repeater_stub = self.repeater.repeater_stub
 
     def tearDown(self):
-        self.repeater_stub.delete()
         self.repeater.delete()
         super().tearDown()
 
@@ -402,3 +400,122 @@ def test_yes(self):
         with make_repeat_record(self.repeater_stub, RECORD_PENDING_STATE):
             is_migrated = are_repeat_records_migrated(DOMAIN)
         self.assertTrue(is_migrated)
+
+
+class RepeaterStubOneToOneRepeaterTests(RepeaterFixtureMixin, TestCase):
+
+    def test_one_in_domain(self):
+        self.assertEqual(RepeaterStub.objects.filter(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+
+    def test_one_globally(self):
+        self.assertEqual(RepeaterStub.objects.filter(
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+
+    def test_only_one_in_domain(self):
+        with transaction.atomic():
+            # (Run in its own transaction so as not to mess with tearDown)
+            with self.assertRaises(IntegrityError):
+                RepeaterStub.objects.create(
+                    domain=DOMAIN,
+                    repeater_id=self.repeater.get_id,
+                )
+
+    def test_only_one_globally(self):
+        with transaction.atomic():
+            with self.assertRaises(IntegrityError):
+                RepeaterStub.objects.create(
+                    domain='another-domain',
+                    repeater_id=self.repeater.get_id,
+                )
+
+
+class SaveDeleteRepeaterTests(TestCase):
+
+    def setUp(self):
+        self.repeater = FormRepeater(
+            domain=DOMAIN,
+            url='https://www.example.com/api/',
+        )
+
+    def tearDown(self):
+        try:
+            self.repeater.delete()
+        except TypeError:
+            # test_delete() causes tearDown() to throw TypeError because
+            # self.repeater is already deleted.
+            pass
+
+    def test_save_new(self):
+        with self.assertNumQueries(1):
+            self.repeater.save()
+        self.assertEqual(RepeaterStub.objects.filter(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+
+    def test_save_old(self):
+        self.repeater.save()
+        self.repeater.paused = True
+        with self.assertNumQueries(0):
+            self.repeater.save()
+        self.assertEqual(RepeaterStub.objects.filter(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+
+    def test_delete(self):
+        self.repeater.save()
+        with self.assertNumQueries(3):
+            # 1. SELECT ... FROM "repeaters_repeaterstub" ...
+            # 2. SELECT ... FROM "repeaters_repeatrecord" ...
+            # 3. DELETE FROM "repeaters_repeaterstub" ...
+            self.repeater.delete()
+        self.assertEqual(RepeaterStub.objects.filter(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        ).count(), 0)
+
+
+class TestMigrationCantDuplicate(RepeaterFixtureMixin, TestCase):
+
+    def test_migration_cant_duplicate(self):
+        self.assertEqual(RepeaterStub.objects.filter(
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+        create_repeaterstubs(None, None)
+        self.assertEqual(RepeaterStub.objects.filter(
+            repeater_id=self.repeater.get_id,
+        ).count(), 1)
+
+
+class PauseResumeRetireRepeaterTests(RepeaterFixtureMixin, TestCase):
+
+    def _get_repeater_stub(self):
+        return RepeaterStub.objects.get(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        )
+
+    def test_pause(self):
+        self.assertFalse(self._get_repeater_stub().is_paused)
+        self.repeater.pause()
+        self.assertTrue(self._get_repeater_stub().is_paused)
+
+    def test_resume(self):
+        repeater_stub = self._get_repeater_stub()
+        repeater_stub.is_paused = True
+        repeater_stub.save()
+
+        self.repeater.resume()
+        self.assertFalse(self._get_repeater_stub().is_paused)
+
+    def test_retire(self):
+        self.repeater.retire()
+        self.assertFalse(RepeaterStub.objects.filter(
+            domain=DOMAIN,
+            repeater_id=self.repeater.get_id,
+        ).exists())

corehq/motech/repeaters/tests/test_tasks.py

@@ -83,14 +83,14 @@ def setUpClass(cls):
             name='Test API',
             url="http://localhost/api/"
         )
-        cls.repeater = FormRepeater(
-            domain=DOMAIN,
-            connection_settings_id=cls.connection_settings.id,
-        )
-        cls.repeater.save()
 
     def setUp(self):
-        self.repeater_stub = RepeaterStub.objects.create(
+        self.repeater = FormRepeater(
+            domain=DOMAIN,
+            connection_settings_id=self.connection_settings.id,
+        )
+        self.repeater.save()
+        self.repeater_stub = RepeaterStub.objects.get(
             domain=DOMAIN,
             repeater_id=self.repeater.get_id,
         )
@@ -104,11 +104,10 @@ def setUp(self):
             just_now += timedelta(seconds=1)
 
     def tearDown(self):
-        self.repeater_stub.delete()
+        self.repeater.delete()
 
     @classmethod
     def tearDownClass(cls):
-        cls.repeater.delete()
         cls.connection_settings.delete()
         cls.domain.delete()
         super().tearDownClass()