Make body of migration run only once

corehq/apps/domain_migration_flags/api.py

@@ -1,6 +1,7 @@
 import logging
 from collections import defaultdict
 from datetime import datetime
+from functools import wraps
 
 from django.db.models import Q
 
@@ -126,3 +127,25 @@ def all_domains_with_migrations_in_progress():
 def reset_caches(domain, slug):
     any_migrations_in_progress(domain, strict=True)
     get_migration_status(domain, slug, strict=True)
+
+
+def once_off_migration(slug):
+    """Ensure that the body of a migration is run only once
+
+    You can still run it again if there is an exception
+    """
+    def outer(migration_fn):
+        @wraps(migration_fn)
+        def inner(*args, **kwargs):
+            if get_migration_complete(ALL_DOMAINS, slug):
+                return
+            set_migration_started(ALL_DOMAINS, slug)
+            try:
+                res = migration_fn(*args, **kwargs)
+            except Exception:
+                set_migration_not_started(ALL_DOMAINS, slug)
+                raise
+            set_migration_complete(ALL_DOMAINS, slug)
+            return res
+        return inner
+    return outer

corehq/apps/domain_migration_flags/tests/test_domain_migration_progress.py

@@ -1,13 +1,16 @@
 import uuid
 from datetime import datetime
+from unittest.mock import Mock
 
 from django.test import TestCase
 
 from corehq.apps.domain_migration_flags.api import (
+    ALL_DOMAINS,
     any_migrations_in_progress,
     get_migration_complete,
     get_migration_status,
     migration_in_progress,
+    once_off_migration,
     set_migration_complete,
     set_migration_not_started,
     set_migration_started,
@@ -20,10 +23,8 @@
 
 class DomainMigrationProgressTest(TestCase):
 
-    @classmethod
-    def setUpClass(cls):
-        super(DomainMigrationProgressTest, cls).setUpClass()
-        cls.slug = uuid.uuid4().hex
+    def setUp(self):
+        self.slug = uuid.uuid4().hex
 
     def get_progress(self, domain):
         return DomainMigrationProgress.objects.get(domain=domain, migration_slug=self.slug)
@@ -109,3 +110,42 @@ def test_any_migrations_in_progress(self):
         self.assertTrue(any_migrations_in_progress('purple'))
         set_migration_complete('purple', 'other_slug')
         self.assertFalse(any_migrations_in_progress('purple'))
+
+    def test_once_off_decorator(self):
+        actual_migration = Mock()
+
+        @once_off_migration(self.slug)
+        def basic_migration():
+            self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.IN_PROGRESS)
+            actual_migration()
+
+        self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.NOT_STARTED)
+        actual_migration.assert_not_called()
+
+        basic_migration()
+        self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.COMPLETE)
+        actual_migration.assert_called_once()
+
+        basic_migration()
+        actual_migration.assert_called_once()
+
+    def test_once_off_decorator_failure(self):
+        actual_migration = Mock()
+
+        @once_off_migration(self.slug)
+        def failing_migration():
+            self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.IN_PROGRESS)
+            actual_migration()
+            raise ValueError('this migration failed')
+
+        self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.NOT_STARTED)
+        self.assertEqual(actual_migration.call_count, 0)
+
+        with self.assertRaises(ValueError):
+            failing_migration()
+        self.assertEqual(get_migration_status(ALL_DOMAINS, self.slug), MigrationStatus.NOT_STARTED)
+        self.assertEqual(actual_migration.call_count, 1)
+
+        with self.assertRaises(ValueError):
+            failing_migration()
+        self.assertEqual(actual_migration.call_count, 2)

corehq/apps/users/management/commands/rm_couch_user_data.py

@@ -1,9 +1,9 @@
 # One-off migration, June 2024
-
 from itertools import chain
 
 from django.core.management.base import BaseCommand
 
+from corehq.apps.domain_migration_flags.api import once_off_migration
 from corehq.apps.users.models import CommCareUser
 from corehq.dbaccessors.couchapps.all_docs import (
     get_doc_count_by_type,
@@ -16,13 +16,18 @@
 class Command(BaseCommand):
     help = "Populate SQL user data from couch"
 
-    def handle(self, **options):
-        db = CommCareUser.get_db()
-        count = (get_doc_count_by_type(db, 'WebUser')
-                 + get_doc_count_by_type(db, 'CommCareUser'))
-        all_ids = chain(iter_all_doc_ids(db, 'WebUser'),
-                        iter_all_doc_ids(db, 'CommCareUser'))
-        iter_update(db, _update_user, with_progress_bar(all_ids, count), verbose=True)
+    def handle(self, *args, **options):
+        do_migration()
+
+
+@once_off_migration("rm_couch_user_data")
+def do_migration():
+    db = CommCareUser.get_db()
+    count = (get_doc_count_by_type(db, 'WebUser')
+             + get_doc_count_by_type(db, 'CommCareUser'))
+    all_ids = chain(iter_all_doc_ids(db, 'WebUser'),
+                    iter_all_doc_ids(db, 'CommCareUser'))
+    iter_update(db, _update_user, with_progress_bar(all_ids, count), verbose=True)
 
 
 def _update_user(user_doc):