Added "managestate" command

django_extensions/management/commands/managestate.py

@@ -0,0 +1,174 @@
+import json
+from operator import itemgetter
+from pathlib import Path
+
+from django import get_version
+from django.core.management import call_command
+from django.core.management.base import BaseCommand, CommandError
+from django.db import DEFAULT_DB_ALIAS, connections
+from django.db.migrations.loader import MigrationLoader
+from django.db.migrations.recorder import MigrationRecorder
+from django.utils import timezone
+from django.utils.version import get_version_tuple
+
+DEFAULT_FILENAME = 'managestate.json'
+DEFAULT_STATE = 'default'
+VERSION = get_version_tuple(get_version())
+
+
+class Command(BaseCommand):
+    help = 'Manage database state in the convenient way.'
+    conn = database = filename = verbosity = None
+    migrate_args = migrate_options = None
+    _applied_migrations = None
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            'action', choices=('dump', 'load'),
+            help='An action to do. '
+                 'Dump action saves applied migrations to a file. '
+                 'Load action applies migrations specified in a file.',
+        )
+        parser.add_argument(
+            'state', nargs='?', default=DEFAULT_STATE,
+            help=f'A name of a state. Usually a name of a git branch. Defaults to "{DEFAULT_STATE}"',
+        )
+        parser.add_argument(
+            '-d', '--database', default=DEFAULT_DB_ALIAS,
+            help=f'Nominates a database to synchronize. Defaults to the "{DEFAULT_DB_ALIAS}" database.',
+        )
+        parser.add_argument(
+            '-f', '--filename', default=DEFAULT_FILENAME,
+            help=f'A file to write to. Defaults to "{DEFAULT_FILENAME}"',
+        )
+
+        # migrate command arguments
+        parser.add_argument(
+            '--noinput', '--no-input', action='store_false', dest='interactive',
+            help='The argument for "migrate" command. '
+                 'Tells Django to NOT prompt the user for input of any kind.',
+        )
+        parser.add_argument(
+            '--fake', action='store_true',
+            help='The argument for "migrate" command. '
+                 'Mark migrations as run without actually running them.',
+        )
+        parser.add_argument(
+            '--fake-initial', action='store_true',
+            help='The argument for "migrate" command. '
+                 'Detect if tables already exist and fake-apply initial migrations if so. Make sure '
+                 'that the current database schema matches your initial migration before using this '
+                 'flag. Django will only check for an existing table name.',
+        )
+        parser.add_argument(
+            '--plan', action='store_true',
+            help='The argument for "migrate" command. '
+                 'Shows a list of the migration actions that will be performed.',
+        )
+        parser.add_argument(
+            '--run-syncdb', action='store_true',
+            help='The argument for "migrate" command. '
+                 'Creates tables for apps without migrations.',
+        )
+        parser.add_argument(
+            '--check', action='store_true', dest='check_unapplied',
+            help='The argument for "migrate" command. '
+                 'Exits with a non-zero status if unapplied migrations exist.',
+        )
+
+    def handle(self, action, database, filename, state, *args, **options):
+        self.migrate_args = args
+        self.migrate_options = options
+        self.verbosity = options['verbosity']
+        self.conn = connections[database]
+        self.database = database
+        self.filename = filename
+        getattr(self, action)(state)
+
+    def dump(self, state: str):
+        """Save applied migrations to a file."""
+        migrated_apps = self.get_migrated_apps()
+        migrated_apps.update(self.get_applied_migrations())
+        self.write({state: migrated_apps})
+        self.stdout.write(self.style.SUCCESS(
+            f'Migrations for state "{state}" have been successfully saved to {self.filename}.'
+        ))
+
+    def load(self, state: str):
+        """Apply migrations from a file."""
+        migrations = self.read().get(state)
+        if migrations is None:
+            raise CommandError(f'No such state saved: {state}')
+
+        if VERSION < (3, 0):
+            self.migrate_options.pop('check_unapplied', None)
+
+        kwargs = {
+            **self.migrate_options,
+            'database': self.database,
+            'verbosity': self.verbosity - 1 if self.verbosity > 1 else 0
+        }
+
+        for app, migration in migrations.items():
+            if self.is_applied(app, migration):
+                continue
+
+            if self.verbosity > 1:
+                self.stdout.write(self.style.WARNING(f'Applying migrations for "{app}"'))
+            args = (app, migration, *self.migrate_args)
+            call_command('migrate', *args, **kwargs)
+
+        self.stdout.write(self.style.SUCCESS(
+            f'Migrations for "{state}" have been successfully applied.'
+        ))
+
+    def get_migrated_apps(self) -> dict:
+        """Installed apps having migrations."""
+        apps = MigrationLoader(self.conn).migrated_apps
+        migrated_apps = dict.fromkeys(apps, 'zero')
+        if self.verbosity > 1:
+            self.stdout.write('Apps having migrations: ' + ', '.join(sorted(migrated_apps)))
+        return migrated_apps
+
+    def get_applied_migrations(self) -> dict:
+        """Installed apps with last applied migrations."""
+        if self._applied_migrations:
+            return self._applied_migrations
+
+        migrations = MigrationRecorder(self.conn).applied_migrations()
+        migrations = migrations if VERSION < (3, 0) else migrations.keys()
+        last_applied = sorted(migrations, key=itemgetter(1))
+
+        self._applied_migrations = dict(last_applied)
+        return self._applied_migrations
+
+    def is_applied(self, app: str, migration: str) -> bool:
+        """Check whether a migration for an app is applied or not."""
+        applied = self.get_applied_migrations().get(app)
+        if applied == migration:
+            if self.verbosity > 1:
+                self.stdout.write(self.style.WARNING(
+                    f'Migrations for "{app}" are already applied.'
+                ))
+            return True
+        return False
+
+    def read(self) -> dict:
+        """Get saved state from the file."""
+        path = Path(self.filename)
+        if not path.exists() or not path.is_file():
+            raise CommandError(f'No such file: {self.filename}')
+
+        with open(self.filename) as file:
+            return json.load(file)
+
+    def write(self, data: dict):
+        """Write new data to the file using existent one."""
+        try:
+            saved = self.read()
+        except CommandError:
+            saved = {}
+
+        saved.update(data, updated_at=str(timezone.now()))
+        with open(self.filename, 'w') as file:
+            json.dump(saved, file, indent=2, sort_keys=True)

docs/managestate.rst

@@ -0,0 +1,70 @@
+managestate
+==========
+
+:synopsis: Saves current applied migrations to a file or applies migrations from this file.
+
+The `managestate` command fetches last applied migrations from a specified database
+and saves them to a specified file. After that, you may easily apply saved migrations.
+The advantage of this approach is that you may have at hand several database states
+and quickly switch between them.
+
+Why?
+----
+
+While you develop several features or fix some bugs at the same time you often meet
+the situation when you need to apply or unapply database migrations before you check out
+to another feature/bug branch. You always need to view current migrations by `showmigrations`,
+then apply or unapply it manually using `migrate` and there is no problem if you work with
+one Django app. But when there is more than one, it starts to annoy. To forget about the problem
+and quickly switch between branches use the `managestate` command.
+
+How?
+----
+
+To dump current migrations use::
+
+    $ ./manage.py managestate dump
+
+A state will be saved to `managestate.json` just about the following::
+
+    {
+      "default": {
+        "admin": "0003_logentry_add_action_flag_choices",
+        "auth": "0012_alter_user_first_name_max_length",
+        "contenttypes": "0002_remove_content_type_name",
+        "sessions": "0001_initial",
+        "sites": "0002_alter_domain_unique",
+        "myapp": "zero"
+      },
+      "updated_at": "2021-06-27 10:42:50.364070"
+    }
+
+As you see, migrations have been saved as the state called "default".
+You can specify it using the positional argument::
+
+    $ ./manage.py managestate dump my_feature_branch
+
+Then migrations will be added to `managestate.json` under the key "my_feature_branch".
+To change the filename use `-f` or `--filename` flag.
+
+When you load a state from a file, you may also use all arguments defined for the `migrate` command.
+
+Examples
+----
+
+Save an initial database state of the branch "master/main" before developing features::
+
+    $ ./manage.py managestate dump master
+
+Check out to your branch, develop your feature, and dump its state when you are going to get reviewed::
+
+    $ ./manage.py managestate dump super-feature
+
+Check out to the "master" branch back and rollback a database state with just one command::
+
+    $ ./manage.py managestate load master
+
+If you need to add some improvements to your feature, just use::
+
+    $ ./manage.py managestate load super-feature
+

tests/management/commands/test_managestate.py

@@ -0,0 +1,107 @@
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+from django.core.management import CommandError, call_command
+from django.db import connection
+from django.db.migrations.recorder import MigrationRecorder
+
+from django_extensions.management.commands.managestate import DEFAULT_FILENAME, DEFAULT_STATE, Command
+
+pytestmark = [pytest.mark.django_db]
+
+COMMAND = 'managestate'
+DEFAULT_FILE = Path(DEFAULT_FILENAME)
+
+
+@pytest.fixture
+def make_dump(request):
+    request.addfinalizer(DEFAULT_FILE.unlink)
+    return call_command(COMMAND, 'dump', '-v', 0)
+
+
+@pytest.fixture
+def cmd_data():
+    cmd = Command()
+    cmd.verbosity = 0
+    cmd.conn = connection
+    data = cmd.get_migrated_apps()
+    data.update(cmd.get_applied_migrations())
+    return data
+
+
+class TestManageStateExceptions:
+    """Tests for managestate management command exceptions."""
+
+    def test_bad_action(self):
+        with pytest.raises(CommandError):
+            call_command(COMMAND, 'bad_action')
+
+    def test_no_such_file(self):
+        with pytest.raises(CommandError):
+            call_command(COMMAND, 'load', '-f', 'non_existent_file')
+
+    def test_not_a_file(self):
+        with pytest.raises(CommandError), TemporaryDirectory() as tmp_dir:
+            call_command(COMMAND, 'load', '-f', tmp_dir)
+
+    @pytest.mark.usefixtures('make_dump')
+    def test_no_such_state(self):
+        with pytest.raises(CommandError):
+            call_command(COMMAND, 'load', 'non_existent_state')
+
+
+class TestManageState:
+    """Tests managestate management command"""
+
+    def test_dump(self, request, capsys):
+        request.addfinalizer(DEFAULT_FILE.unlink)
+        call_command(COMMAND, 'dump', '-v', 0)
+        stdout, _ = capsys.readouterr()
+        assert DEFAULT_FILE.exists()
+        assert 'successfully saved' in stdout
+
+    @pytest.mark.parametrize('filename', [DEFAULT_FILENAME, 'custom_f1l3n4m3.json'])
+    def test_dump_files(self, request, filename):
+        path = Path(filename)
+        request.addfinalizer(path.unlink)
+        call_command(COMMAND, 'dump', '-f', filename)
+        assert path.exists()
+
+    @pytest.mark.parametrize('state', [DEFAULT_STATE, 'custom_state'])
+    def test_dump_states(self, request, state):
+        request.addfinalizer(DEFAULT_FILE.unlink)
+        call_command(COMMAND, 'dump', state)
+        with open(DEFAULT_FILE) as file:
+            data = json.load(file)
+        assert isinstance(data, dict)
+        assert data.get(state) is not None
+
+    @pytest.mark.usefixtures('make_dump')
+    def test_load(self, capsys):
+        call_command(COMMAND, 'load', '-v', 0)
+        stdout, _ = capsys.readouterr()
+        assert 'successfully applied' in stdout
+
+    @pytest.mark.parametrize('filename', [DEFAULT_FILENAME, 'custom_f1l3n4m3.json'])
+    def test_load_files(self, request, capsys, filename):
+        request.addfinalizer(Path(filename).unlink)
+        call_command(COMMAND, 'dump', '-f', filename)
+        call_command(COMMAND, 'load', '-f', filename)
+        stdout, _ = capsys.readouterr()
+        assert 'successfully applied' in stdout
+
+    @pytest.mark.parametrize('state', [DEFAULT_STATE, 'custom_state'])
+    def test_load_states(self, request, capsys, state):
+        request.addfinalizer(DEFAULT_FILE.unlink)
+        call_command(COMMAND, 'dump', state)
+        call_command(COMMAND, 'load', state)
+        stdout, _ = capsys.readouterr()
+        assert 'successfully applied' in stdout
+
+    def test_migration_is_last_applied(self, cmd_data):
+        migrations = MigrationRecorder(connection).applied_migrations()
+        for app, migration in cmd_data.items():
+            last_migration = sorted(filter(lambda x: x[0] == app, migrations))[-1][1]
+            assert migration == last_migration