migrate PrimaryHero model to use addon instead of promoted_addon in a migration safe way (#23333)

* migrate PrimaryHero model to use addon instead of promoted_addon in a migration safe way

- Changed the `addon` field in the `PrimaryHero` model to be nullable to be able to continue querying before the migration is complete
- Introduced a new task `sync_primary_hero_addon` to synchronize the `addon` field with the `promoted_addon` where applicable.

* Fix from review

Author: KevinMind

src/olympia/hero/migrations/0020_primaryhero_addon.py

@@ -36,12 +36,4 @@ class Migration(migrations.Migration):
             name='addon',
             field=models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to='addons.addon', null=True),
         ),
-        # Migrate promoted_addon to addon
-        migrations.RunPython(primary_hero_addon_up, primary_hero_addon_down),
-        # Make the addon field non-nullable after populating the existing data
-        migrations.AlterField(
-            model_name='primaryhero',
-            name='addon',
-            field=models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to='addons.addon', null=False),
-        ),
     ]

src/olympia/hero/models.py

@@ -179,7 +179,7 @@ class PrimaryHero(ModelBase):
     promoted_addon = models.OneToOneField(
         PromotedAddon, on_delete=models.CASCADE, null=True
     )
-    addon = models.OneToOneField(Addon, on_delete=models.CASCADE, null=False)
+    addon = models.OneToOneField(Addon, on_delete=models.CASCADE, null=True)
 
     is_external = models.BooleanField(default=False)
 

src/olympia/hero/tasks.py

@@ -0,0 +1,25 @@
+import olympia.core.logger
+from olympia.amo.celery import task
+from olympia.amo.decorators import use_primary_db
+from olympia.hero.models import PrimaryHero
+
+
+log = olympia.core.logger.getLogger('z.hero')
+
+
+@task
+@use_primary_db
+def sync_primary_hero_addon():
+    invalid_heroes = []
+    for hero in PrimaryHero.objects.filter(addon__isnull=True):
+        promoted_addon = hero.promoted_addon
+
+        if promoted_addon is None:
+            invalid_heroes.append(hero)
+        else:
+            log.info(f'Syncing hero {hero} with promoted addon {hero.promoted_addon}')
+            hero.addon = promoted_addon.addon
+            hero.save()
+
+    if len(invalid_heroes) > 0:
+        log.error(f'Invalid PrimaryHero records: {invalid_heroes}')

src/olympia/hero/tests/test_tasks.py

@@ -0,0 +1,44 @@
+from olympia.amo.tests import TestCase, addon_factory
+from olympia.hero.models import PrimaryHero
+from olympia.hero.tasks import sync_primary_hero_addon
+from olympia.promoted.models import PromotedAddon
+
+
+class TestSyncPrimaryHeroAddon(TestCase):
+    def test_no_heros(self):
+        assert PrimaryHero.objects.count() == 0
+        sync_primary_hero_addon.apply()
+        assert PrimaryHero.objects.count() == 0
+
+    def test_with_legacy_promoted_addon(self):
+        addon = addon_factory()
+        # Create legacy promoted_addon
+        promoted_addon = PromotedAddon.objects.create(addon=addon)
+        # Also make the addon promoted with the new promoted addon model
+        self.make_addon_promoted(addon, 1)
+        hero = PrimaryHero.objects.create(promoted_addon=promoted_addon)
+        sync_primary_hero_addon.apply()
+        hero.reload()
+        assert hero.addon == addon
+
+    def test_with_new_promoted_addon(self):
+        addon = addon_factory()
+        self.make_addon_promoted(addon, 1)
+        hero = PrimaryHero.objects.create(addon=addon)
+        sync_primary_hero_addon.apply()
+        hero.reload()
+        assert hero.addon == addon
+
+    def test_with_no_addon_or_legacy_promoted_addon_should_raise(self):
+        """
+        During the transition to using `addon` instead of `promoted_addon`,
+        It is possible that there are some heroes that have neither an addon
+        or a promoted addon. This should not happen, but we should handle the edge case
+        until we have removed the legacy promoted_addon field.
+        """
+        PrimaryHero.objects.create()
+
+        with self.assertLogs(logger='z.hero', level='ERROR') as cm:
+            sync_primary_hero_addon.apply()
+
+        assert 'Invalid PrimaryHero records' in cm.output[0]

src/olympia/lib/settings_base.py

@@ -743,6 +743,7 @@ def get_db_config(environ_var, atomic_requests=True):
     'olympia.files.tasks.backfill_file_manifest': {'queue': 'adhoc'},
     'olympia.files.tasks.extract_host_permissions': {'queue': 'adhoc'},
     'olympia.lib.crypto.tasks.bump_and_resign_addons': {'queue': 'adhoc'},
+    'olympia.hero.tasks.sync_primary_hero_addon': {'queue': 'adhoc'},
     # Misc AMO tasks.
     'olympia.blocklist.tasks.monitor_remote_settings': {'queue': 'amo'},
     'olympia.abuse.tasks.appeal_to_cinder': {'queue': 'amo'},