Fix data_seed migration data being skipped (#22860)

* Skip initial migration if seeding on initialize.py

* Remove redundant code

* Replace flush db with reset_db to ensure --clean resets migrations

* Add django check for db charset

* Remove uneeded import_licenses command

* Fix tests

src/olympia/amo/management/commands/data_seed.py

@@ -1,6 +1,3 @@
-import os
-import shutil
-
 from django.conf import settings
 from django.core.management import call_command
 
@@ -13,56 +10,25 @@ class Command(BaseDataCommand):
         'generated add-ons, and data from AMO production.'
     )
 
-    def _clean_storage(self, root: str, dir_dict: dict[str, str | dict]) -> None:
-        for key, value in dir_dict.items():
-            curr_path = os.path.join(root, key)
-            if isinstance(value, dict):
-                self._clean_storage(curr_path, value)
-            else:
-                shutil.rmtree(curr_path, ignore_errors=True)
-                os.makedirs(curr_path, exist_ok=True)
-
-    def clean_storage(self):
-        self.logger.info('Cleaning storage...')
-        self._clean_storage(
-            settings.STORAGE_ROOT,
-            {
-                'files': '',
-                'shared_storage': {
-                    'tmp': {
-                        'addons': '',
-                        'data': '',
-                        'file_viewer': '',
-                        'guarded-addons': '',
-                        'icon': '',
-                        'log': '',
-                        'persona_header': '',
-                        'preview': '',
-                        'test': '',
-                        'uploads': '',
-                    },
-                    'uploads': '',
-                },
-            },
-        )
-
     def handle(self, *args, **options):
         num_addons = 10
         num_themes = 5
 
+        # Delete any existing data_seed backup
         self.clean_dir(self.data_backup_init)
 
         self.logger.info('Resetting database...')
-        call_command('flush', '--noinput')
+        call_command('reset_db', '--no-utf8', '--noinput')
+        # Delete any local storage files
+        # This should happen after we reset the database to ensure any records
+        # relying on storage are deleted.
         self.clean_storage()
-        # reindex --wipe will force the ES mapping to be re-installed.
-        call_command('reindex', '--wipe', '--force', '--noinput')
+        # Migrate the database
         call_command('migrate', '--noinput')
 
         self.logger.info('Loading initial data...')
         call_command('loaddata', 'initial.json')
         call_command('import_prod_versions')
-        call_command('import_licenses')
         call_command(
             'createsuperuser',
             '--no-input',

src/olympia/amo/management/commands/initialize.py

@@ -33,6 +33,7 @@ class Command(BaseDataCommand):
         """
         Create the database.
         """
+        logging.info(f'options: {options}')
         # We need to support skipping loading/seeding when desired.
         # Like in CI environments where you don't want to load data every time.
         if settings.DATA_BACKUP_SKIP:
@@ -41,10 +42,12 @@ class Command(BaseDataCommand):
             )
             return
 
-        clean = options.get('clean')
-        load = options.get('load')
-        logging.info(f'options: {options}')
+        # If DB empty or we are explicitly cleaning, then bail with data_seed.
+        if options.get('clean') or not self.local_admin_exists():
+            call_command('data_seed')
+            return
 
+        load = options.get('load')
         # We always migrate the DB.
         logging.info('Migrating...')
         call_command('migrate', '--noinput')
@@ -52,9 +55,6 @@ class Command(BaseDataCommand):
         # If we specify a specifi backup, simply load that.
         if load:
             call_command('data_load', '--name', load)
-        # If DB empty or we are explicitly cleaning, then reseed.
-        elif clean or not self.local_admin_exists():
-            call_command('data_seed')
         # We should reindex even if no data is loaded/modified
         # because we might have a fresh instance of elasticsearch
         else:

src/olympia/amo/tests/test_commands.py

@@ -152,15 +152,13 @@ def test_generate_jsi18n_files():
 
 class BaseTestDataCommand(TestCase):
     class Commands:
-        flush = mock.call('flush', '--noinput')
+        reset_db = mock.call('reset_db', '--no-utf8', '--noinput')
         migrate = mock.call('migrate', '--noinput')
         data_seed = mock.call('data_seed')
 
-        flush = mock.call('flush', '--noinput')
         reindex = mock.call('reindex', '--wipe', '--force', '--noinput')
         load_initial_data = mock.call('loaddata', 'initial.json')
         import_prod_versions = mock.call('import_prod_versions')
-        import_licenses = mock.call('import_licenses')
         createsuperuser = mock.call(
             'createsuperuser',
             '--no-input',
@@ -278,16 +276,15 @@ class TestInitializeDataCommand(BaseTestDataCommand):
     def test_handle_with_clean_and_load_arguments(self):
         """
         Test running the 'initialize' command with both '--clean' and '--load'
-        arguments. Expected: Command should prioritize '--load' and perform
-        migration, loading.
+        arguments. Expected: Command should prioritize '--clean' and perform
+        migration and seeding.
         """
         name = 'test'
         call_command('initialize', clean=True, load=name)
         self._assert_commands_called_in_order(
             self.mocks['mock_call_command'],
             [
-                self.mock_commands.migrate,
-                self.mock_commands.data_load(name),
+                self.mock_commands.data_seed,
             ],
         )
 
@@ -301,7 +298,6 @@ class TestInitializeDataCommand(BaseTestDataCommand):
         self._assert_commands_called_in_order(
             self.mocks['mock_call_command'],
             [
-                self.mock_commands.migrate,
                 self.mock_commands.data_seed,
             ],
         )
@@ -332,7 +328,6 @@ class TestInitializeDataCommand(BaseTestDataCommand):
         self._assert_commands_called_in_order(
             self.mocks['mock_call_command'],
             [
-                self.mock_commands.migrate,
                 self.mock_commands.data_seed,
             ],
         )
@@ -343,6 +338,7 @@ class TestInitializeDataCommand(BaseTestDataCommand):
         Expected: The command exits with an error and does not proceed to seeding
         or loading data.
         """
+        self.with_local_admin()
         self.mocks['mock_call_command'].side_effect = Exception('test')
         with pytest.raises(Exception) as context:
             call_command('initialize')
@@ -599,6 +595,10 @@ class TestSeedDataCommand(BaseTestDataCommand):
                 'mock_clean_dir',
                 'olympia.amo.management.commands.data_seed.BaseDataCommand.clean_dir',
             ),
+            (
+                'mock_clean_storage',
+                'olympia.amo.management.commands.data_seed.BaseDataCommand.clean_storage',
+            ),
         )
 
         self.mocks = {}
@@ -611,15 +611,18 @@ class TestSeedDataCommand(BaseTestDataCommand):
     def test_default(self):
         call_command('data_seed')
 
+        self.mocks['mock_clean_dir'].assert_called_once_with(
+            self.base_data_command.data_backup_init
+        )
+        self.mocks['mock_clean_storage'].assert_called_once()
+
         self._assert_commands_called_in_order(
             self.mocks['mock_call_command'],
             [
-                self.mock_commands.flush,
-                self.mock_commands.reindex,
+                self.mock_commands.reset_db,
                 self.mock_commands.migrate,
                 self.mock_commands.load_initial_data,
                 self.mock_commands.import_prod_versions,
-                self.mock_commands.import_licenses,
                 self.mock_commands.createsuperuser,
                 self.mock_commands.load_zadmin_users,
                 self.mock_commands.generate_addons('firefox', 10),

src/olympia/applications/management/commands/import_licenses.py

@@ -1,15 +0,0 @@
-from django.core.management.base import BaseCommand
-
-from olympia.constants.licenses import ALL_LICENSES
-from olympia.versions.models import License
-
-
-class Command(BaseCommand):
-    help = """Import a the licenses."""
-
-    def handle(self, *args, **options):
-        for license in ALL_LICENSES:
-            try:
-                License.objects.get_or_create(builtin=license.builtin)
-            except Exception:
-                continue

src/olympia/core/apps.py

@@ -9,6 +9,7 @@ from django.conf import settings
 from django.core.checks import Error, Tags, register
 from django.core.management import call_command
 from django.core.management.base import CommandError
+from django.db import connection
 from django.utils.translation import gettext_lazy as _
 
 from olympia.core.utils import get_version_json
@@ -98,6 +99,26 @@ def static_check(app_configs, **kwargs):
     return errors
 
 
+@register(CustomTags.custom_setup)
+def db_charset_check(app_configs, **kwargs):
+    errors = []
+
+    with connection.cursor() as cursor:
+        cursor.execute("SHOW VARIABLES LIKE 'character_set_database';")
+        result = cursor.fetchone()
+        if result[1] != settings.DB_CHARSET:
+            errors.append(
+                Error(
+                    'Database charset invalid. '
+                    f'Expected {settings.DB_CHARSET}, '
+                    f'recieved {result[1]}',
+                    id='setup.E005',
+                )
+            )
+
+    return errors
+
+
 class CoreConfig(AppConfig):
     name = 'olympia.core'
     verbose_name = _('Core')

src/olympia/core/tests/test_apps.py

@@ -2,10 +2,22 @@ from unittest import mock
 
 from django.core.management import call_command
 from django.core.management.base import SystemCheckError
-from django.test import SimpleTestCase
+from django.test import TestCase
 
 
-class SystemCheckIntegrationTest(SimpleTestCase):
+class SystemCheckIntegrationTest(TestCase):
+    @mock.patch('olympia.core.apps.connection.cursor')
+    def test_db_charset_check(self, mock_cursor):
+        mock_cursor.return_value.__enter__.return_value.fetchone.return_value = (
+            'character_set_database',
+            'utf8mb3',
+        )
+        with self.assertRaisesMessage(
+            SystemCheckError,
+            'Database charset invalid. Expected utf8mb4, recieved utf8mb3',
+        ):
+            call_command('check')
+
     def test_uwsgi_check(self):
         call_command('check')
 

src/olympia/lib/settings_base.py

@@ -130,6 +130,9 @@ CORS_ALLOW_HEADERS = list(default_headers) + [
     'x-country-code',
 ]
 
+DB_ENGINE = 'olympia.core.db.mysql'
+DB_CHARSET = 'utf8mb4'
+
 
 def get_db_config(environ_var, atomic_requests=True):
     values = env.db(var=environ_var, default='mysql://root:@127.0.0.1/olympia')
@@ -142,9 +145,9 @@ def get_db_config(environ_var, atomic_requests=True):
             'ATOMIC_REQUESTS': atomic_requests,
             # Pool our database connections up for 300 seconds
             'CONN_MAX_AGE': 300,
-            'ENGINE': 'olympia.core.db.mysql',
+            'ENGINE': DB_ENGINE,
             'OPTIONS': {
-                'charset': 'utf8mb4',
+                'charset': DB_CHARSET,
                 'sql_mode': 'STRICT_ALL_TABLES',
                 'isolation_level': 'read committed',
             },
@@ -167,6 +170,8 @@ REPLICA_DATABASES = []
 LOCAL_ADMIN_EMAIL = 'local_admin@mozilla.com'
 LOCAL_ADMIN_USERNAME = 'local_admin'
 
+DJANGO_EXTENSIONS_RESET_DB_MYSQL_ENGINES = [DB_ENGINE]
+
 # Local time zone for this installation. Choices can be found here:
 # http://en.wikipedia.org/wiki/List_of_tz_zones_by_name
 # although not all choices may be available on all operating systems.