feat(ci): add command to rollback migrations (#4768)

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"

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"

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 "<datetime from step 1>"
+```
+
+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