Merge 1c6df6f into b4839b8

README.rst

@@ -212,6 +212,17 @@ To write a repeat migration, make sure your migration has ``should_run`` defined
 The above migration will run **every time** ``migrate`` is called, except if it
 is marked as "deferred".  ``up`` is run if ``should_run`` returns True.
 
+To write a deferrable migration, add ``@deferrable`` to the up function::
+
+    from dbmigrator import deferrable
+
+
+    @deferrable
+    def up(cursor):
+        # this migration is automatically deferred
+
+The above migration will not run unless you use ``migrate --run-defers``.
+
 rollback
 --------
 

dbmigrator/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from .utils import logger
+from .utils import logger, deferrable
 
 
-__all__ = ('logger',)
+__all__ = ('logger', 'deferrable')

dbmigrator/commands/mark.py

@@ -19,17 +19,26 @@ def cli_command(cursor, migrations_directory='', migration_timestamp='',
     if completed is None:
         raise Exception('-t, -f or -d must be supplied.')
 
-    migrations = utils.get_migrations(
-        migrations_directory, import_modules=False)
-    for version, _ in migrations:
-        if version == migration_timestamp:
-            break
-    else:
+    migrations = {version: migration
+                  for version, _, migration in utils.get_migrations(
+                      migrations_directory, import_modules=True)}
+    migration = migrations.get(migration_timestamp)
+    if migration is None:
         migrated_versions = list(utils.get_schema_versions(cursor))
         if migration_timestamp not in migrated_versions:
-            logger.warning(
+            raise SystemExit(
                 'Migration {} not found'.format(migration_timestamp))
 
+    if not completed:
+        try:
+            if migration.up.dbmigrator_deferrable:
+                logger.error(
+                    'Deferrable migration {} cannot be marked as not deferred'
+                    .format(migration_timestamp))
+                return
+        except AttributeError:
+            # not a deferrable migration
+            pass
     utils.mark_migration(cursor, migration_timestamp, completed)
     if not completed:
         message = 'not been run'

dbmigrator/commands/migrate.py

@@ -15,10 +15,10 @@
 
 @utils.with_cursor
 def cli_command(cursor, migrations_directory='', version='',
-                db_connection_string='', **kwargs):
+                db_connection_string='', run_defers=False, **kwargs):
     pending_migrations = utils.get_pending_migrations(
         migrations_directory, cursor, import_modules=True,
-        up_to_version=version)
+        up_to_version=version, include_defers=run_defers)
 
     migrated = False
     for version, migration_name, migration in pending_migrations:
@@ -28,7 +28,8 @@ def cli_command(cursor, migrations_directory='', version='',
                              cursor,
                              version,
                              migration_name,
-                             migration)
+                             migration,
+                             run_defers)
 
     if not migrated:
         print('No pending migrations.  Database is up to date.')
@@ -37,4 +38,7 @@ def cli_command(cursor, migrations_directory='', version='',
 def cli_loader(parser):
     parser.add_argument('--version',
                         help='Migrate database up to this version')
+    parser.add_argument('--run-defers',
+                        action='store_true',
+                        help='Also run the deferred migrations')
     return cli_command

dbmigrator/tests/data/md/20170810124056_empty.py

@@ -1,11 +1,14 @@
 # -*- coding: utf-8 -*-
 
+from dbmigrator import deferrable
+
 
 # Uncomment should_run if this is a repeat migration
 # def should_run(cursor):
 #     # TODO return True if migration should run
 
 
+@deferrable
 def up(cursor):
     # TODO migration code
     pass

dbmigrator/tests/test_cli.py

@@ -217,18 +217,19 @@ def test_mark_as_true_already_true(self):
         self.assertIn('20160228202637   add_table         True', stdout)
         self.assertIn('20160228212456   cool_stuff        True', stdout)
 
-    @mock.patch('dbmigrator.logger.warning')
-    def test_migration_not_found(self, warning):
+    def test_migration_not_found(self):
         testing.install_test_packages()
         cmd = ['--db-connection-string', testing.db_connection_string]
 
         self.target(cmd + ['--context', 'package-a', 'init', '--version', '0'])
 
-        self.target(cmd + ['mark', '-t', '012345'])
-        warning.assert_called_with('Migration 012345 not found')
+        with self.assertRaises(SystemExit) as cm:
+            self.target(cmd + ['mark', '-t', '012345'])
+        self.assertEqual('Migration 012345 not found', cm.exception.message)
 
-        self.target(cmd + ['mark', '-f', '012345'])
-        warning.assert_called_with('Migration 012345 not found')
+        with self.assertRaises(SystemExit) as cm:
+            self.target(cmd + ['mark', '-f', '012345'])
+        self.assertEqual('Migration 012345 not found', cm.exception.message)
 
     def test_mark_as_false(self):
         testing.install_test_packages()
@@ -325,6 +326,28 @@ def test_mark_as_deferred(self):
         self.assertIn('20160228202637   add_table         deferred', stdout)
         self.assertIn('20160228212456   cool_stuff        False', stdout)
 
+    @mock.patch('dbmigrator.logger.error')
+    def test_deferrable(self, error):
+        md = os.path.join(testing.test_data_path, 'md')
+        cmd = ['--db-connection-string', testing.db_connection_string,
+               '--migrations-directory', md]
+
+        self.target(cmd + ['init', '--version', '0'])
+
+        # it is not possible to mark a deferrable migration as not deferred
+        self.target(cmd + ['mark', '-f', '20170810124056'])
+
+        error.assert_called_with('Deferrable migration 20170810124056 '
+                                 'cannot be marked as not deferred')
+
+        # mark a deferrable migration as completed
+        with testing.captured_output() as (out, err):
+            self.target(cmd + ['mark', '-t', '20170810124056'])
+
+        stdout = out.getvalue()
+        self.assertEqual('Migration 20170810124056 marked as completed\n',
+                         stdout)
+
 
 class GenerateTestCase(BaseTestCase):
     @mock.patch('dbmigrator.utils.timestamp')
@@ -479,6 +502,45 @@ def cleanup():
         stdout = out.getvalue()
         self.assertIn('No pending migrations', stdout)
 
+    def test_deferrable(self):
+        md = os.path.join(testing.test_data_path, 'md')
+        cmd = ['--db-connection-string', testing.db_connection_string,
+               '--migrations-directory', md]
+
+        def cleanup():
+            if os.path.exists('insert_data.txt'):
+                os.remove('insert_data.txt')
+            with psycopg2.connect(testing.db_connection_string) as db_conn:
+                with db_conn.cursor() as cursor:
+                    cursor.execute('DROP TABLE IF EXISTS a_table')
+
+        self.addCleanup(cleanup)
+
+        self.target(cmd + ['init', '--version', '0'])
+        with testing.captured_output() as (out, err):
+            self.target(cmd + ['migrate'])
+
+        stdout = out.getvalue()
+        self.assertIn('+CREATE TABLE a_table', stdout)
+        self.assertIn('Marking migration 20170810124056 empty as deferred',
+                      stdout)
+
+        # Run the repeat migration by creating this file
+        with open('insert_data.txt', 'w') as f:
+            f.write('三好')
+
+        with testing.captured_output() as (out, err):
+            self.target(cmd + ['migrate'])
+
+        stdout = out.getvalue()
+        self.assertIn('Running migration 20170810093943', stdout)
+
+        with testing.captured_output() as (out, err):
+            self.target(cmd + ['migrate', '--run-defers'])
+
+        stdout = out.getvalue()
+        self.assertIn('Running migration 20170810124056', stdout)
+
 
 class RollbackTestCase(BaseTestCase):
     @mock.patch('dbmigrator.utils.timestamp')
@@ -568,7 +630,7 @@ def cleanup():
 
         self.target(cmd + ['init', '--version', '0'])
         with testing.captured_output() as (out, err):
-            self.target(cmd + ['migrate'])
+            self.target(cmd + ['migrate', '--run-defers'])
 
         # Run the repeat migration by creating this file
         with open('insert_data.txt', 'w') as f:

dbmigrator/utils.py

@@ -36,6 +36,14 @@
 logger.addHandler(handler)
 
 
+def deferrable(func):
+    """A decorator to mark the migration as deferred by default.  The state of
+    the migration can be changed later by using `mark`.
+    """
+    func.dbmigrator_deferrable = True
+    return func
+
+
 # psycopg2 / libpq doesn't respond to SIGINT (ctrl-c):
 # https://github.com/psycopg/psycopg2/issues/333
 # To get around this problem, using code from:
@@ -164,7 +172,7 @@ def get_schema_versions(cursor, versions_only=True, raise_error=True,
 
 
 def get_pending_migrations(migration_directories, cursor, import_modules=False,
-                           up_to_version=None):
+                           up_to_version=None, include_defers=False):
     migrated_versions = {i[0]: i[1] or 'deferred'
                          for i in get_schema_versions(
                              cursor, versions_only=False)}
@@ -182,14 +190,17 @@ def get_pending_migrations(migration_directories, cursor, import_modules=False,
         version, migration_name, mod = migration
         if not import_modules:
             migration = migration[:-1]
-        if migrated_versions.get(version) != 'deferred':
+        if migrated_versions.get(version) != 'deferred' or \
+                include_defers:
             try:
                 mod.should_run
                 # repeat migrations are always included
                 yield migration
             except AttributeError:
                 # not a repeat migration
-                if version not in migrated_versions:
+                if version not in migrated_versions or \
+                        migrated_versions[version] == 'deferred' and \
+                        include_defers:
                     yield migration
 
 
@@ -205,7 +216,8 @@ def compare_schema(db_connection_string, callback, *args, **kwargs):
                              n=10))).encode('utf-8'))
 
 
-def run_migration(cursor, version, migration_name, migration):
+def run_migration(cursor, version, migration_name, migration,
+                  run_defers=False):
     try:
         if not migration.should_run(cursor):
             print('Skipping migration {} {}: should_run is false'
@@ -214,6 +226,18 @@ def run_migration(cursor, version, migration_name, migration):
     except AttributeError:
         # not a repeat migration
         pass
+
+    if not run_defers:
+        try:
+            if migration.up.dbmigrator_deferrable:
+                print('Marking migration {} {} as deferred'
+                      .format(version, migration_name))
+                mark_migration(cursor, version, 'deferred')
+                return
+        except AttributeError:
+            # not a deferred migration
+            pass
+
     print('Running migration {} {}'.format(version, migration_name))
     migration.up(cursor)
     mark_migration(cursor, version, True)