Upgrade and downgrade based on revision existence

Before there was no way to upgrade to a specific revision, because the
revision passed to the manager was passed to all the plugins, some of
which failed to process it.

Add a new method to every plugin `has_revision`, which returns
whether the plugin has the revision. Use it in upgrade and
downgrade.

Change-Id: I89b02d7ad479da6bff3c492c88edfee9c19abc22
Closes-Bug: 1474067

5 files changed, 161 insertions, 6 deletions

diff --git a/oslo_db/sqlalchemy/migration_cli/ext_alembic.py b/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
index 1dbf88f..825c638 100644
--- a/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
+++ b/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
@@ -15,6 +15,7 @@ import os
 import alembic
 from alembic import config as alembic_config
 import alembic.migration as alembic_migration
+from alembic import script as alembic_script
 
 from oslo_db.sqlalchemy.migration_cli import ext_base
 
@@ -89,3 +90,14 @@ class AlembicExtension(ext_base.MigrationExtensionBase):
         with self.engine.begin() as connection:
             self.config.attributes['connection'] = connection
             return alembic.command.stamp(self.config, revision=revision)
+
+    def has_revision(self, rev_id):
+        if rev_id in ['base', 'head']:
+            return True
+        script = alembic_script.ScriptDirectory(
+            self.config.get_main_option('alembic_repo_path'))
+        try:
+            script.get_revision(rev_id)
+            return True
+        except alembic.util.CommandError:
+            return False
diff --git a/oslo_db/sqlalchemy/migration_cli/ext_base.py b/oslo_db/sqlalchemy/migration_cli/ext_base.py
index 205b7b2..184003d 100644
--- a/oslo_db/sqlalchemy/migration_cli/ext_base.py
+++ b/oslo_db/sqlalchemy/migration_cli/ext_base.py
@@ -70,6 +70,15 @@ class MigrationExtensionBase(object):
         """
         raise NotImplementedError()
 
+    def has_revision(self, rev_id):
+        """Checks whether the repo contains a revision
+
+        :param rev_id: Revision to check
+        :returns: Whether the revision is in the repo
+        :rtype: bool
+        """
+        raise NotImplementedError()
+
     def __cmp__(self, other):
         """Used for definition of plugin order.
 
diff --git a/oslo_db/sqlalchemy/migration_cli/ext_migrate.py b/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
index eb72818..4002972 100644
--- a/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
+++ b/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
@@ -13,6 +13,8 @@
 import logging
 import os
 
+from migrate.versioning import version as migrate_version
+
 from oslo_db._i18n import _LE
 from oslo_db.sqlalchemy import migration
 from oslo_db.sqlalchemy.migration_cli import ext_base
@@ -65,3 +67,14 @@ class MigrateExtension(ext_base.MigrationExtensionBase):
     def version(self):
         return migration.db_version(
             self.engine, self.repository, init_version=self.init_version)
+
+    def has_revision(self, rev_id):
+        collection = migrate_version.Collection(self.repository)
+        try:
+            collection.version(rev_id)
+            return True
+        except (KeyError, ValueError):
+            # NOTE(breton): migrate raises KeyError if an int is passed but not
+            # found in the list of revisions and ValueError if non-int is
+            # passed. Both mean there is no requested revision.
+            return False
diff --git a/oslo_db/sqlalchemy/migration_cli/manager.py b/oslo_db/sqlalchemy/migration_cli/manager.py
index 9b55887..d335fa5 100644
--- a/oslo_db/sqlalchemy/migration_cli/manager.py
+++ b/oslo_db/sqlalchemy/migration_cli/manager.py
@@ -13,6 +13,8 @@
 import sqlalchemy
 from stevedore import enabled
 
+from oslo_db import exception
+
 
 MIGRATION_NAMESPACE = 'oslo.db.migration'
 
@@ -50,17 +52,40 @@ class MigrationManager(object):
 
     def upgrade(self, revision):
         """Upgrade database with all available backends."""
+        # a revision exists only in a single plugin. Until we reached it, we
+        # should upgrade to the plugins' heads.
+        # revision=None is a special case meaning latest revision.
+        rev_in_plugins = [p.has_revision(revision) for p in self._plugins]
+        if not any(rev_in_plugins) and revision is not None:
+            raise exception.DbMigrationError('Revision does not exist')
+
         results = []
-        for plugin in self._plugins:
-            results.append(plugin.upgrade(revision))
+        for plugin, has_revision in zip(self._plugins, rev_in_plugins):
+            if not has_revision or revision is None:
+                results.append(plugin.upgrade(None))
+            else:
+                results.append(plugin.upgrade(revision))
+                break
         return results
 
     def downgrade(self, revision):
         """Downgrade database with available backends."""
+        # a revision exists only in a single plugin. Until we reached it, we
+        # should upgrade to the plugins' first revision.
+        # revision=None is a special case meaning initial revision.
+        rev_in_plugins = [p.has_revision(revision) for p in self._plugins]
+        if not any(rev_in_plugins) and revision is not None:
+            raise exception.DbMigrationError('Revision does not exist')
+
         # downgrading should be performed in reversed order
         results = []
-        for plugin in reversed(self._plugins):
-            results.append(plugin.downgrade(revision))
+        for plugin, has_revision in zip(reversed(self._plugins),
+                                        reversed(rev_in_plugins)):
+            if not has_revision or revision is None:
+                results.append(plugin.downgrade(None))
+            else:
+                results.append(plugin.downgrade(revision))
+                break
         return results
 
     def version(self):
diff --git a/oslo_db/tests/sqlalchemy/test_migrate_cli.py b/oslo_db/tests/sqlalchemy/test_migrate_cli.py
index 209dfab..6bdb09a 100644
--- a/oslo_db/tests/sqlalchemy/test_migrate_cli.py
+++ b/oslo_db/tests/sqlalchemy/test_migrate_cli.py
@@ -10,10 +10,12 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import alembic
 import mock
 from oslotest import base as test_base
 import sqlalchemy
 
+from oslo_db import exception
 from oslo_db.sqlalchemy.migration_cli import ext_alembic
 from oslo_db.sqlalchemy.migration_cli import ext_migrate
 from oslo_db.sqlalchemy.migration_cli import manager
@@ -94,6 +96,29 @@ class TestAlembicExtension(test_base.BaseTestCase):
         version = self.alembic.version()
         self.assertIsNone(version)
 
+    def test_has_revision(self, command):
+        with mock.patch(('oslo_db.sqlalchemy.migration_cli.'
+                         'ext_alembic.alembic_script')) as mocked:
+            self.alembic.config.get_main_option = mock.Mock()
+            # since alembic_script is mocked and no exception is raised, call
+            # will result in success
+            self.assertIs(True, self.alembic.has_revision('test'))
+            mocked.ScriptDirectory().get_revision.assert_called_once_with(
+                'test')
+            self.assertIs(True, self.alembic.has_revision(None))
+            self.assertIs(True, self.alembic.has_revision('head'))
+
+    def test_has_revision_negative(self, command):
+        with mock.patch(('oslo_db.sqlalchemy.migration_cli.'
+                         'ext_alembic.alembic_script')) as mocked:
+            mocked.ScriptDirectory().get_revision.side_effect = (
+                alembic.util.CommandError)
+            self.alembic.config.get_main_option = mock.Mock()
+            # exception is raised, the call should be false
+            self.assertIs(False, self.alembic.has_revision('test'))
+            mocked.ScriptDirectory().get_revision.assert_called_once_with(
+                'test')
+
 
 @mock.patch(('oslo_db.sqlalchemy.migration_cli.'
              'ext_migrate.migration'))
@@ -159,6 +184,21 @@ class TestMigrateExtension(test_base.BaseTestCase):
             self.migration_config['init_version'],
             init_version=self.migration_config['init_version'])
 
+    def test_has_revision(self, command):
+        with mock.patch(('oslo_db.sqlalchemy.migration_cli.'
+                         'ext_migrate.migrate_version')) as mocked:
+            self.migrate.has_revision('test')
+            mocked.Collection().version.assert_called_once_with('test')
+            # tip of the branch should always be True
+            self.assertIs(True, self.migrate.has_revision(None))
+
+    def test_has_revision_negative(self, command):
+        with mock.patch(('oslo_db.sqlalchemy.migration_cli.'
+                         'ext_migrate.migrate_version')) as mocked:
+            mocked.Collection().version.side_effect = ValueError
+            self.assertIs(False, self.migrate.has_revision('test'))
+            mocked.Collection().version.assert_called_once_with('test')
+
 
 class TestMigrationManager(test_base.BaseTestCase):
 
@@ -214,7 +254,7 @@ class TestMigrationManager(test_base.BaseTestCase):
                          err.args[0])
 
 
-class TestMigrationRightOrder(test_base.BaseTestCase):
+class TestMigrationMultipleExtensions(test_base.BaseTestCase):
 
     def setUp(self):
         self.migration_config = {'alembic_ini_path': '.',
@@ -233,7 +273,7 @@ class TestMigrationRightOrder(test_base.BaseTestCase):
         self.second_ext.obj.downgrade.return_value = 100
         self.migration_manager._manager.extensions = [self.first_ext,
                                                       self.second_ext]
-        super(TestMigrationRightOrder, self).setUp()
+        super(TestMigrationMultipleExtensions, self).setUp()
 
     def test_upgrade_right_order(self):
         results = self.migration_manager.upgrade(None)
@@ -242,3 +282,59 @@ class TestMigrationRightOrder(test_base.BaseTestCase):
     def test_downgrade_right_order(self):
         results = self.migration_manager.downgrade(None)
         self.assertEqual(results, [100, 0])
+
+    def test_upgrade_does_not_go_too_far(self):
+        self.first_ext.obj.has_revision.return_value = True
+        self.second_ext.obj.has_revision.return_value = False
+        self.second_ext.obj.upgrade.side_effect = AssertionError(
+            'this method should not have been called')
+
+        results = self.migration_manager.upgrade(100)
+        self.assertEqual([100], results)
+
+    def test_downgrade_does_not_go_too_far(self):
+        self.second_ext.obj.has_revision.return_value = True
+        self.first_ext.obj.has_revision.return_value = False
+        self.first_ext.obj.downgrade.side_effect = AssertionError(
+            'this method should not have been called')
+
+        results = self.migration_manager.downgrade(100)
+        self.assertEqual([100], results)
+
+    def test_upgrade_checks_rev_existence(self):
+        self.first_ext.obj.has_revision.return_value = False
+        self.second_ext.obj.has_revision.return_value = False
+
+        # upgrade to a specific non-existent revision should fail
+        self.assertRaises(exception.DbMigrationError,
+                          self.migration_manager.upgrade, 100)
+
+        # upgrade to the "head" should succeed
+        self.assertEqual([100, 200], self.migration_manager.upgrade(None))
+
+        # let's assume the second ext has the revision, upgrade should succeed
+        self.second_ext.obj.has_revision.return_value = True
+        self.assertEqual([100, 200], self.migration_manager.upgrade(200))
+
+        # upgrade to the "head" should still succeed
+        self.assertEqual([100, 200], self.migration_manager.upgrade(None))
+
+    def test_downgrade_checks_rev_existence(self):
+        self.first_ext.obj.has_revision.return_value = False
+        self.second_ext.obj.has_revision.return_value = False
+
+        # upgrade to a specific non-existent revision should fail
+        self.assertRaises(exception.DbMigrationError,
+                          self.migration_manager.downgrade, 100)
+
+        # downgrade to the "base" should succeed
+        self.assertEqual([100, 0], self.migration_manager.downgrade(None))
+
+        # let's assume the second ext has the revision, downgrade should
+        # succeed
+        self.first_ext.obj.has_revision.return_value = True
+        self.assertEqual([100, 0], self.migration_manager.downgrade(200))
+
+        # downgrade to the "base" should still succeed
+        self.assertEqual([100, 0], self.migration_manager.downgrade(None))
+        self.assertEqual([100, 0], self.migration_manager.downgrade('base'))