diff --git a/api/core/management/commands/rollbackmigrationsappliedafter.py b/api/core/management/commands/rollbackmigrationsappliedafter.py new file mode 100644 index 000000000000..f0473d91dd27 --- /dev/null +++ b/api/core/management/commands/rollbackmigrationsappliedafter.py @@ -0,0 +1,58 @@ +from argparse import ArgumentParser +from datetime import datetime + +from django.core.management import BaseCommand, CommandError, call_command +from django.db.migrations.recorder import MigrationRecorder + + +class Command(BaseCommand): + """ + Rollback all migrations applied on or after a given datetime. + + Usage: python manage.py rollbackmigrationsappliedafter "2024-10-24 08:23:45" + """ + + def add_arguments(self, parser: ArgumentParser): + parser.add_argument( + "dt", + type=str, + help="Rollback all migrations applied on or after this datetime (provided in ISO format)", + ) + + def handle(self, *args, dt: str, **kwargs) -> None: + try: + _dt = datetime.fromisoformat(dt) + except ValueError: + raise CommandError("Date must be in ISO format") + + applied_migrations = MigrationRecorder.Migration.objects.filter(applied__gte=_dt).order_by("applied") + if not applied_migrations.exists(): + self.stdout.write(self.style.NOTICE("No migrations to rollback.")) + + # Since we've ordered by the date applied, we know that the first entry in the qs for each app + # is the earliest migration after the supplied date. + earliest_migration_by_app = {} + for migration in applied_migrations: + if migration.app in earliest_migration_by_app: + continue + earliest_migration_by_app[migration.app] = migration.name + + for app, migration_name in earliest_migration_by_app.items(): + call_command( + "migrate", app, _get_previous_migration_number(migration_name) + ) + + +def _get_previous_migration_number(migration_name: str) -> str: + """ + Returns the previous migration number (0 padded number to 4 characters), or zero + if the provided migration name is the first for a given app (usually 0001_initial). + + Examples: + _get_previous_migration_number("0001_initial") -> "zero" + _get_previous_migration_number("0009_migration_9") -> "0008" + _get_previous_migration_number("0103_migration_103") -> "0102" + """ + + migration_number = int(migration_name.split("_", maxsplit=1)[0]) + return f"{migration_number - 1:04}" if migration_number > 1 else "zero" diff --git a/api/tests/unit/core/test_unit_core_management.py b/api/tests/unit/core/test_unit_core_management.py new file mode 100644 index 000000000000..8f4c023e5ab0 --- /dev/null +++ b/api/tests/unit/core/test_unit_core_management.py @@ -0,0 +1,92 @@ +from unittest.mock import call + +import pytest +from _pytest.capture import CaptureFixture +from django.core.management import CommandError, call_command +from django.db.migrations.recorder import MigrationRecorder +from pytest_mock import MockerFixture + + +class MockQuerySet(list): + def exists(self) -> bool: + return self.__len__() > 0 + + +def test_rollbackmigrationsappliedafter(mocker: MockerFixture) -> None: + # Given + dt_string = "2024-10-24 08:23:45" + + migration_1 = mocker.MagicMock(app="foo", spec=MigrationRecorder.Migration) + migration_1.name = "0001_initial" + + migration_2 = mocker.MagicMock(app="bar", spec=MigrationRecorder.Migration) + migration_2.name = "0002_some_migration_description" + + migration_3 = mocker.MagicMock(app="bar", spec=MigrationRecorder.Migration) + migration_3.name = "0003_some_other_migration_description" + + migrations = MockQuerySet([migration_1, migration_2, migration_3]) + + mocked_migration_recorder = mocker.patch( + "core.management.commands.rollbackmigrationsappliedafter.MigrationRecorder" + ) + mocked_migration_recorder.Migration.objects.filter.return_value.order_by.return_value = ( + migrations + ) + + mocked_call_command = mocker.patch( + "core.management.commands.rollbackmigrationsappliedafter.call_command" + ) + + # When + call_command("rollbackmigrationsappliedafter", dt_string) + + # Then + assert mocked_call_command.mock_calls == [ + call("migrate", "foo", "zero"), + call("migrate", "bar", "0001"), + ] + + +def test_rollbackmigrationsappliedafter_invalid_date(mocker: MockerFixture) -> None: + # Given + dt_string = "foo" + + mocked_call_command = mocker.patch( + "core.management.commands.rollbackmigrationsappliedafter.call_command" + ) + + # When + with pytest.raises(CommandError) as e: + call_command("rollbackmigrationsappliedafter", dt_string) + + # Then + assert mocked_call_command.mock_calls == [] + assert e.value.args == ("Date must be in ISO format",) + + +def test_rollbackmigrationsappliedafter_no_migrations( + mocker: MockerFixture, capsys: CaptureFixture +) -> None: + # Given + dt_string = "2024-10-01" + + mocked_migration_recorder = mocker.patch( + "core.management.commands.rollbackmigrationsappliedafter.MigrationRecorder" + ) + mocked_migration_recorder.Migration.objects.filter.return_value.order_by.return_value = MockQuerySet( + [] + ) + + mocked_call_command = mocker.patch( + "core.management.commands.rollbackmigrationsappliedafter.call_command" + ) + + # When + call_command("rollbackmigrationsappliedafter", dt_string) + + # Then + assert mocked_call_command.mock_calls == [] + + captured = capsys.readouterr() + assert captured.out == "No migrations to rollback.\n" diff --git a/docs/docs/deployment/hosting/locally-api.md b/docs/docs/deployment/hosting/locally-api.md index d2b051883ab9..6f080e5cc406 100644 --- a/docs/docs/deployment/hosting/locally-api.md +++ b/docs/docs/deployment/hosting/locally-api.md @@ -477,7 +477,20 @@ ORDER BY applied DESC ::: -2. Replace the datetime in the query below with a datetime after the deployment of the version you want to roll back to, +2. Run the following command inside a Flagsmith API container running the _current_ version of Flagsmith + +```bash +python manage.py rollbackmigrationsafter "" +``` + +3. Roll back the Flagsmith API to the desired version. + +### Steps pre v2.151.0 + +If you are rolling back from a version earlier than v2.151.0, you will need to replace step 2 above with the following 2 +steps. + +1. Replace the datetime in the query below with a datetime after the deployment of the version you want to roll back to, and before any subsequent deployments. Execute the subsequent query against the Flagsmith database. ```sql {14} showLineNumbers @@ -511,8 +524,7 @@ Example output: python manage.py migrate token_blacklist zero ``` -3. Run the generated commands inside a Flagsmith API container running the _current_ version of Flagsmith -4. Roll back the Flagsmith API to the desired version. +2. Run the generated commands inside a Flagsmith API container running the _current_ version of Flagsmith ## Information for Developers working on the project