PR #26 with added docs

arnold/__init__.py

@@ -1,15 +1,16 @@
+from __future__ import print_function
 import os
 import sys
 
 import argparse
 
+from peewee import sort_models_topologically
 from termcolor import colored
 
 from arnold.exceptions import DirectionNotFoundException
 from arnold.models import Migration
 from importlib import import_module
 
-
 class Terminator:
     IGNORED_FILES = ["__init__"]
 
@@ -44,7 +45,7 @@ def _retreive_filenames(self):
         for f in files:
             splits = f.rsplit(".", 1)
             if len(splits) <= 1 or splits[1] != "py" or \
-               splits[0] in self.IGNORED_FILES:
+                            splits[0] in self.IGNORED_FILES:
                 continue
             filenames.append(splits[0])
         return sorted(filenames, key=lambda fname: int(fname.split("_")[0]))
@@ -108,6 +109,11 @@ def get_latest_migration(self):
             self.model.migration.desc()
         ).first()
 
+    def get_applied_migrations(self):
+        return self.model.select().order_by(
+            self.model.migration.desc()
+        )
+
     def perform_migrations(self, direction):
         """
         Find the migration if it is passed in and call the up or down method as
@@ -118,47 +124,38 @@ def perform_migrations(self, direction):
 
         filenames = self._retreive_filenames()
 
-        if self.direction == "down":
-            filenames.reverse()
-
         if len(filenames) <= 0:
             return True
 
-        start = 0
-
-        latest_migration = self.get_latest_migration()
-
-        if latest_migration:
-            migration_index = filenames.index(latest_migration.migration)
+        applied_migrations = self.get_applied_migrations()
+        applied_migrations_names = set([x.migration for x in applied_migrations])
+        filenames_set = set(filenames)
+        if self.direction == 'up':
+            migrations_to_run = filenames_set.difference(applied_migrations_names)
+        else:
+            migrations_to_run = filenames_set.intersection(applied_migrations_names)
+        migrations_to_run = sorted(list(migrations_to_run), key=lambda fname: int(fname.split("_")[0]))
 
-            if migration_index == len(filenames) - 1 and \
-                self.direction == 'up':
-                print("Nothing to go {0}.".format(
-                    colored(self.direction, "magenta"))
-                )
-                return False
+        if self.direction == "down":
+            migrations_to_run.reverse()
 
-            if self.direction == 'up':
-                start = migration_index + 1
-            else:
-                start = migration_index
-        if not latest_migration and self.direction == 'down':
+        if len(migrations_to_run) == 0:
             print("Nothing to go {0}.".format(
                 colored(self.direction, "magenta"))
             )
             return False
 
         if self.count == 0:
-            end = len(filenames)
+            end = len(migrations_to_run)
         else:
-            end = start + self.count
+            end = self.count
 
-        migrations_to_complete = filenames[start:end]
+        migrations_to_complete = migrations_to_run[:end]
 
         if self.count > len(migrations_to_complete):
             print(
                 "Count {0} greater than available migrations. Going {1} {2} times."
-                .format(
+                    .format(
                     colored(self.count, "green"),
                     colored(self.direction, "magenta"),
                     colored(len(migrations_to_complete), "red"),
@@ -196,6 +193,71 @@ def init(args):
     return True
 
 
+def create_tables(args):
+    def my_import(name):
+        components = name.split('.')
+        mod = __import__('.'.join(components[:-1]))
+        for comp in components[1:]:
+            mod = getattr(mod, comp)
+        return mod
+
+    model_classes = []
+    errors = False
+    for model in args.models:
+        print("Importing model {0}... ".format(colored(model, "magenta")), end='')
+        try:
+            model_class = my_import(model)
+            model_classes.append(model_class)
+            print(colored('OK', 'green'))
+        except ImportError:
+            print(colored('import error', 'red'))
+            errors = True
+    if errors:
+        sys.exit(1)
+
+    terminator = Terminator(args)
+
+    next_migration_num = '0001'
+    if terminator._retreive_filenames():
+        next_migration_num = terminator._retreive_filenames()[-1].split('_')[0]
+        next_migration_num = "%04d" % (int(next_migration_num) + 1)
+
+    migration_file_name = '{0}/{1}/{2}_auto_create_tables_{3}.py'.format(
+        terminator.folder, 'migrations', next_migration_num,
+        "_".join([m.__name__.lower() for m in model_classes]) if len(model_classes) < 4 else len(model_classes)
+    )
+
+    qc = terminator.database.compiler()
+
+    print("Writing down migration file", colored(migration_file_name, 'blue'))
+    with open(migration_file_name, 'w') as migration_file:
+
+        print("from {0} import {1} as model_class".format(model_classes[0].__module__, model_classes[0].__name__),
+              file=migration_file)
+        print("database = model_class._meta.database", file=migration_file)
+
+        print("\n\ndef up():", file=migration_file)
+        for m in sort_models_topologically(model_classes):
+            print("\n    # Create model", m.__module__ + '.' + m.__name__, file=migration_file)
+            print("    database.execute_sql('%s')\n" % qc.create_table(m)[0], file=migration_file)
+
+            for field in m._fields_to_index():
+                print("    database.execute_sql('%s')" % qc.create_index(m, [field], field.unique)[0],
+                      file=migration_file)
+
+            if m._meta.indexes:
+                for fields, unique in m._meta.indexes:
+                    fobjs = [m._meta.fields[f] for f in fields]
+                    print("    database.execute_sql('%s')" % qc.create_index(m, fobjs, unique)[0],
+                          file=migration_file)
+
+        print("\n\ndef down():", file=migration_file)
+
+        for m in reversed(sort_models_topologically(model_classes)):
+            print("\n    # Drop model", m.__module__ + '.' + m.__name__, file=migration_file)
+            print("    database.execute_sql('%s')\n" % qc.drop_table(m, cascade=True)[0], file=migration_file)
+
+
 def parse_args(args):
     sys.path.insert(0, os.getcwd())
     parser = argparse.ArgumentParser(description='Migrations. Down. Up.')
@@ -222,7 +284,7 @@ def parse_args(args):
     up_cmd = subparsers.add_parser('up', help='Migrate up.')
     up_cmd.set_defaults(func=up)
     up_cmd.add_argument(
-        'count', type=int, help='How many migrations to go up.'
+        'count', type=int, default=0, nargs='?', help='How many migrations to go up.'
     )
     up_cmd.add_argument(
         '--fake', type=bool, default=False, help='Fake the migration.'
@@ -234,8 +296,15 @@ def parse_args(args):
         'count', type=int, help='How many migrations to go down.'
     )
 
+    add_create_tables = subparsers.add_parser('add_create_tables',
+                                              help='Add migration with freezed SQL for create tables.')
+    add_create_tables.set_defaults(func=create_tables)
+    add_create_tables.add_argument('models', metavar='model', type=str, nargs='+',
+                                   help='models given with full import path')
+
     return parser.parse_args(args)
 
+
 def main():
     sys.argv.pop(0)
     args = parse_args(sys.argv)

docs/creating_migrations.rst

@@ -23,3 +23,17 @@ The basic template for a migration is as follows: ::
       pass
 
 A more complete example of a migration file can be found `here <https://github.com/cam-stitt/arnold/blob/master/tests/arnold_config/migrations/001_initial.py>`_.
+
+You can automatically create a migration file for migrations that create models. You want to do that to freeze state of models.
+See following example, file 001_initial.py: ::
+
+    db.create_table(SomeModel)
+
+File 002_some_change.py: ::
+
+    Add some_column to SomeModel
+
+In this scenario if you run this migration against clean database, migration 002 will fail because migration 001 already created some_column.
+To freeze models properly you can use `add_create_tables` method: ::
+
+    $ arnold add_create_tables fully.qualified.path.SomeModel1 fully.qualified.path.SomeModel2 ...