Andrew Godwin avatar Andrew Godwin committed 95eb414

--add-index option to startmigration

Comments (0)

Files changed (1)

south/management/commands/startmigration.py

             help='Generate a Create Table migration for the specified model.  Add multiple models to this migration with subsequent --model parameters.'),
         make_option('--add-field', action='append', dest='added_field_list', type='string',
             help='Generate an Add Column migration for the specified modelname.fieldname - you can use this multiple times to add more than one column.'),
+        make_option('--add-index', action='append', dest='added_index_list', type='string',
+            help='Generate an Add Index migration for the specified modelname.fieldname - you can use this multiple times to add more than one column.'),
         make_option('--initial', action='store_true', dest='initial', default=False,
             help='Generate the initial schema for the app.'),
         make_option('--auto', action='store_true', dest='auto', default=False,
     help = "Creates a new template migration for the given app"
     usage_str = "Usage: ./manage.py startmigration appname migrationname [--initial] [--auto] [--model ModelName] [--add-field ModelName.field_name] [--freeze] [--stdout]"
     
-    def handle(self, app=None, name="", added_model_list=None, added_field_list=None, initial=False, freeze_list=None, auto=False, stdout=False, **options):
+    def handle(self, app=None, name="", added_model_list=None, added_field_list=None, initial=False, freeze_list=None, auto=False, stdout=False, added_index_list=None, **options):
         
         # Any supposed lists that are None become empty lists
         added_model_list = added_model_list or []
         added_field_list = added_field_list or []
-        
+        added_index_list = added_index_list or []
+
         # --stdout means name = -
         if stdout:
             name = "-"
         changed_fields = [] # (mkey, fname, old_def, new_def)
         added_uniques = set() # (mkey, field_names)
         deleted_uniques = set() # (mkey, field_names)
+
+        added_indexes = set()
+        deleted_indexes = []
+        
         
         # --initial means 'add all models in this app'.
         if initial:
             else:
                 added_fields.add(("%s.%s" % (app, modelname), fieldname))
         
+        # same thing as above, but for indexes
+        for fielddef in added_index_list:
+            try:
+                modelname, fieldname = fielddef.split(".", 1)
+            except ValueError:
+                print "The field specification '%s' is not in modelname.fieldname format." % fielddef
+            else:
+                added_indexes.add(("%s.%s" % (app, modelname), fieldname))
+        
         # Add anything frozen (I almost called the dict Iceland...)
         if freeze_list:
             for item in freeze_list:
                 model._meta.app_label,
                 model._meta.object_name,
             )
-        
+
+        ### Added indexes. going here, since it might add to added_uniques ###
+        for mkey, field_name in added_indexes:
+            # Get the model
+            model = model_unkey(mkey)
+            # Get the field
+            try:
+                field = model._meta.get_field(field_name)
+            except FieldDoesNotExist:
+                print "Model '%s' doesn't have a field '%s'" % (mkey, field_name)
+                return
+
+            if field.unique:
+                ut = (mkey, (field.name,))
+                added_uniques.add(ut)
+
+            elif field.db_index:
+                # Create migrations
+                forwards += CREATE_INDEX_SNIPPET % (
+                    model._meta.object_name,
+                    field.name,
+                    model._meta.db_table,
+                    field.name,
+                )
+
+                backwards += DELETE_INDEX_SNIPPET % (
+                    model._meta.object_name,
+                    field.name,
+                    model._meta.db_table,
+                    field.column,
+                )
+                print " + Added index for '%s.%s'" % (mkey, field_name)
+
+            else:
+                print "Field '%s.%s' does not have db_index or unique set to True" % (mkey, field_name)
+                return
         
         ### Changed fields ###
         for mkey, field_name, old_triple, new_triple, last_orm in changed_fields:
         # Renaming column for field '%(field_name)s'.
         db.rename_column(%(old_column)r, %(new_column)r)
         '''
-FIELD_NEEDS_DEF_SNIPPET = "<< PUT FIELD DEFINITION HERE >>"
+FIELD_NEEDS_DEF_SNIPPET = "<< PUT FIELD DEFINITION HERE >>"
+
+CREATE_INDEX_SNIPPET = '''
+        # Adding index on '%s.%s'
+        db.create_index(%r, [%r])
+        '''
+DELETE_INDEX_SNIPPET = '''
+        # Deleting index on '%s.%s'
+        db.delete_index(%r, [%r])
+        '''
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.