Ticket #27844: optimizemigration.py

File optimizemigration.py, 3.1 KB (added by Raphael Gaschignard, 4 years ago)
Line 
1from django.core.management import BaseCommand
2from django.core.management import CommandError
3from django.db import DEFAULT_DB_ALIAS
4from django.db import connections
5from django.db.migrations.exceptions import AmbiguityError
6from django.db.migrations.loader import MigrationLoader
7from django.db.migrations.optimizer import MigrationOptimizer
8from django.db.migrations.writer import MigrationWriter
9
10
11class Command(BaseCommand):
12    def add_arguments(self, parser):
13        parser.add_argument('app_label',
14                            help='App label of the application to optimize a migration for.')
15        parser.add_argument('migration_name',
16                            help='Migration to optimize')
17
18    def handle(self, *args, **options):
19        verbosity = options.get('verbosity')
20
21        app_label = options['app_label']
22        migration_name = options['migration_name']
23
24        loader = MigrationLoader(connections[DEFAULT_DB_ALIAS])
25
26        if app_label not in loader.migrated_apps:
27            raise CommandError(
28                "App '%s' does not have migrations (so squashmigrations on "
29                "it makes no sense)" % app_label
30            )
31
32        migration = self.find_migration(loader, app_label, migration_name)
33
34        optimizer = MigrationOptimizer()
35
36        new_operations = optimizer.optimize(migration.operations, migration.app_label)
37
38        if len(new_operations) == len(migration.operations):
39            if verbosity > 0:
40                self.stdout.write("  No optimizations possible.")
41            return
42        else:
43            if verbosity > 0:
44                self.stdout.write(
45                    "  Optimized from %s operations to %s operations." %
46                    (len(migration.operations), len(new_operations))
47                )
48
49        # set the new migration optimizations
50        migration.operations = new_operations
51
52        # write the migration back out
53        writer = MigrationWriter(migration)
54        with open(writer.path, "wb") as fh:
55            fh.write(writer.as_string())
56
57        if verbosity > 0:
58            self.stdout.write(self.style.MIGRATE_HEADING("Optimized migration %s" % writer.path))
59            if writer.needs_manual_porting:
60                self.stdout.write(self.style.MIGRATE_HEADING("Manual porting required"))
61                self.stdout.write("  Your migrations contained functions that must be manually copied over,")
62                self.stdout.write("  as we could not safely copy their implementation.")
63                self.stdout.write("  See the comment at the top of the squashed migration for details.")
64
65    def find_migration(self, loader, app_label, name):
66        try:
67            return loader.get_migration_by_prefix(app_label, name)
68        except AmbiguityError:
69            raise CommandError(
70                "More than one migration matches '%s' in app '%s'. Please be "
71                "more specific." % (name, app_label)
72            )
73        except KeyError:
74            raise CommandError(
75                "Cannot find a migration matching '%s' from app '%s'." %
76                (name, app_label)
77            )
Back to Top