Merge "Refactor database migration manager to use given engine"

5 files changed, 92 insertions, 31 deletions

diff --git a/oslo_db/sqlalchemy/migration_cli/ext_alembic.py b/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
index 243ae47..1dbf88f 100644
--- a/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
+++ b/oslo_db/sqlalchemy/migration_cli/ext_alembic.py
@@ -17,7 +17,6 @@ from alembic import config as alembic_config
 import alembic.migration as alembic_migration
 
 from oslo_db.sqlalchemy.migration_cli import ext_base
-from oslo_db.sqlalchemy import session as db_session
 
 
 class AlembicExtension(ext_base.MigrationExtensionBase):
@@ -28,31 +27,41 @@ class AlembicExtension(ext_base.MigrationExtensionBase):
     def enabled(self):
         return os.path.exists(self.alembic_ini_path)
 
-    def __init__(self, migration_config):
+    def __init__(self, engine, migration_config):
         """Extension to provide alembic features.
 
+        :param engine: SQLAlchemy engine instance for a given database
+        :type engine: sqlalchemy.engine.Engine
         :param migration_config: Stores specific configuration for migrations
         :type migration_config: dict
         """
         self.alembic_ini_path = migration_config.get('alembic_ini_path', '')
         self.config = alembic_config.Config(self.alembic_ini_path)
+        # TODO(viktors): Remove this, when we will use Alembic 0.7.5 or
+        #                higher, because the ``attributes`` dictionary was
+        #                added to Alembic in version 0.7.5.
+        if not hasattr(self.config, 'attributes'):
+            self.config.attributes = {}
         # option should be used if script is not in default directory
         repo_path = migration_config.get('alembic_repo_path')
         if repo_path:
             self.config.set_main_option('script_location', repo_path)
-        self.db_url = migration_config['db_url']
+        self.engine = engine
 
     def upgrade(self, version):
-        return alembic.command.upgrade(self.config, version or 'head')
+        with self.engine.begin() as connection:
+            self.config.attributes['connection'] = connection
+            return alembic.command.upgrade(self.config, version or 'head')
 
     def downgrade(self, version):
         if isinstance(version, int) or version is None or version.isdigit():
             version = 'base'
-        return alembic.command.downgrade(self.config, version)
+        with self.engine.begin() as connection:
+            self.config.attributes['connection'] = connection
+            return alembic.command.downgrade(self.config, version)
 
     def version(self):
-        engine = db_session.create_engine(self.db_url)
-        with engine.connect() as conn:
+        with self.engine.connect() as conn:
             context = alembic_migration.MigrationContext.configure(conn)
             return context.get_current_revision()
 
@@ -65,8 +74,10 @@ class AlembicExtension(ext_base.MigrationExtensionBase):
                              state
         :type autogenerate: bool
         """
-        return alembic.command.revision(self.config, message=message,
-                                        autogenerate=autogenerate)
+        with self.engine.begin() as connection:
+            self.config.attributes['connection'] = connection
+            return alembic.command.revision(self.config, message=message,
+                                            autogenerate=autogenerate)
 
     def stamp(self, revision):
         """Stamps database with provided revision.
@@ -75,4 +86,6 @@ class AlembicExtension(ext_base.MigrationExtensionBase):
                          database with most recent revision
         :type revision: string
         """
-        return alembic.command.stamp(self.config, revision=revision)
+        with self.engine.begin() as connection:
+            self.config.attributes['connection'] = connection
+            return alembic.command.stamp(self.config, revision=revision)
diff --git a/oslo_db/sqlalchemy/migration_cli/ext_migrate.py b/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
index e31ee3d..eb72818 100644
--- a/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
+++ b/oslo_db/sqlalchemy/migration_cli/ext_migrate.py
@@ -16,7 +16,6 @@ import os
 from oslo_db._i18n import _LE
 from oslo_db.sqlalchemy import migration
 from oslo_db.sqlalchemy.migration_cli import ext_base
-from oslo_db.sqlalchemy import session as db_session
 
 
 LOG = logging.getLogger(__name__)
@@ -31,11 +30,10 @@ class MigrateExtension(ext_base.MigrationExtensionBase):
 
     order = 1
 
-    def __init__(self, migration_config):
+    def __init__(self, engine, migration_config):
+        self.engine = engine
         self.repository = migration_config.get('migration_repo_path', '')
         self.init_version = migration_config.get('init_version', 0)
-        self.db_url = migration_config['db_url']
-        self.engine = db_session.create_engine(self.db_url)
 
     @property
     def enabled(self):
diff --git a/oslo_db/sqlalchemy/migration_cli/manager.py b/oslo_db/sqlalchemy/migration_cli/manager.py
index bda3c2a..9b55887 100644
--- a/oslo_db/sqlalchemy/migration_cli/manager.py
+++ b/oslo_db/sqlalchemy/migration_cli/manager.py
@@ -10,6 +10,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import sqlalchemy
 from stevedore import enabled
 
 
@@ -23,11 +24,21 @@ def check_plugin_enabled(ext):
 
 class MigrationManager(object):
 
-    def __init__(self, migration_config):
+    def __init__(self, migration_config, engine=None):
+        if engine is None:
+            if migration_config.get('db_url'):
+                engine = sqlalchemy.create_engine(
+                    migration_config['db_url'],
+                    poolclass=sqlalchemy.pool.NullPool,
+                )
+            else:
+                raise ValueError('Either database url or engine'
+                                 ' must be provided.')
+
         self._manager = enabled.EnabledExtensionManager(
             MIGRATION_NAMESPACE,
             check_plugin_enabled,
-            invoke_kwds={'migration_config': migration_config},
+            invoke_args=(engine, migration_config),
             invoke_on_load=True
         )
         if not self._plugins:
diff --git a/oslo_db/tests/old_import_api/sqlalchemy/test_migrate_cli.py b/oslo_db/tests/old_import_api/sqlalchemy/test_migrate_cli.py
index 135d44e..660a6af 100644
--- a/oslo_db/tests/old_import_api/sqlalchemy/test_migrate_cli.py
+++ b/oslo_db/tests/old_import_api/sqlalchemy/test_migrate_cli.py
@@ -12,6 +12,7 @@
 
 import mock
 from oslotest import base as test_base
+import sqlalchemy
 
 from oslo.db.sqlalchemy.migration_cli import ext_alembic
 from oslo.db.sqlalchemy.migration_cli import ext_migrate
@@ -35,7 +36,9 @@ class TestAlembicExtension(test_base.BaseTestCase):
     def setUp(self):
         self.migration_config = {'alembic_ini_path': '.',
                                  'db_url': 'sqlite://'}
-        self.alembic = ext_alembic.AlembicExtension(self.migration_config)
+        self.engine = sqlalchemy.create_engine(self.migration_config['db_url'])
+        self.alembic = ext_alembic.AlembicExtension(
+            self.engine, self.migration_config)
         super(TestAlembicExtension, self).setUp()
 
     def test_check_enabled_true(self, command):
@@ -52,7 +55,8 @@ class TestAlembicExtension(test_base.BaseTestCase):
         Verifies enabled returns False on empty alembic_ini_path variable
         """
         self.migration_config['alembic_ini_path'] = ''
-        alembic = ext_alembic.AlembicExtension(self.migration_config)
+        alembic = ext_alembic.AlembicExtension(
+            self.engine, self.migration_config)
         self.assertFalse(alembic.enabled)
 
     def test_upgrade_none(self, command):
@@ -91,14 +95,16 @@ class TestAlembicExtension(test_base.BaseTestCase):
         self.assertIsNone(version)
 
 
-@mock.patch(('oslo.db.sqlalchemy.migration_cli.'
+@mock.patch(('oslo_db.sqlalchemy.migration_cli.'
              'ext_migrate.migration'))
 class TestMigrateExtension(test_base.BaseTestCase):
 
     def setUp(self):
         self.migration_config = {'migration_repo_path': '.',
                                  'db_url': 'sqlite://'}
-        self.migrate = ext_migrate.MigrateExtension(self.migration_config)
+        self.engine = sqlalchemy.create_engine(self.migration_config['db_url'])
+        self.migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         super(TestMigrateExtension, self).setUp()
 
     def test_check_enabled_true(self, migration):
@@ -106,7 +112,8 @@ class TestMigrateExtension(test_base.BaseTestCase):
 
     def test_check_enabled_false(self, migration):
         self.migration_config['migration_repo_path'] = ''
-        migrate = ext_migrate.MigrateExtension(self.migration_config)
+        migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         self.assertFalse(migrate.enabled)
 
     def test_upgrade_head(self, migration):
@@ -143,7 +150,8 @@ class TestMigrateExtension(test_base.BaseTestCase):
 
     def test_change_init_version(self, migration):
         self.migration_config['init_version'] = 101
-        migrate = ext_migrate.MigrateExtension(self.migration_config)
+        migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         migrate.downgrade(None)
         migration.db_sync.assert_called_once_with(
             migrate.engine,
@@ -158,9 +166,11 @@ class TestMigrationManager(test_base.BaseTestCase):
         self.migration_config = {'alembic_ini_path': '.',
                                  'migrate_repo_path': '.',
                                  'db_url': 'sqlite://'}
+        engine = sqlalchemy.create_engine(self.migration_config['db_url'])
         self.migration_manager = manager.MigrationManager(
-            self.migration_config)
+            self.migration_config, engine)
         self.ext = mock.Mock()
+        self.ext.obj.version = mock.Mock(return_value=0)
         self.migration_manager._manager.extensions = [self.ext]
         super(TestMigrationManager, self).setUp()
 
@@ -180,6 +190,10 @@ class TestMigrationManager(test_base.BaseTestCase):
         self.migration_manager.version()
         self.ext.obj.version.assert_called_once_with()
 
+    def test_version_return_value(self):
+        version = self.migration_manager.version()
+        self.assertEqual(0, version)
+
     def test_revision_message_autogenerate(self):
         self.migration_manager.revision('test', True)
         self.ext.obj.revision.assert_called_once_with('test', True)
@@ -192,6 +206,13 @@ class TestMigrationManager(test_base.BaseTestCase):
         self.migration_manager.stamp('stamp')
         self.ext.obj.stamp.assert_called_once_with('stamp')
 
+    def test_wrong_config(self):
+        err = self.assertRaises(ValueError,
+                                manager.MigrationManager,
+                                {'wrong_key': 'sqlite://'})
+        self.assertEqual('Either database url or engine must be provided.',
+                         err.args[0])
+
 
 class TestMigrationRightOrder(test_base.BaseTestCase):
 
@@ -199,8 +220,9 @@ class TestMigrationRightOrder(test_base.BaseTestCase):
         self.migration_config = {'alembic_ini_path': '.',
                                  'migrate_repo_path': '.',
                                  'db_url': 'sqlite://'}
+        engine = sqlalchemy.create_engine(self.migration_config['db_url'])
         self.migration_manager = manager.MigrationManager(
-            self.migration_config)
+            self.migration_config, engine)
         self.first_ext = MockWithCmp()
         self.first_ext.obj.order = 1
         self.first_ext.obj.upgrade.return_value = 100
diff --git a/oslo_db/tests/sqlalchemy/test_migrate_cli.py b/oslo_db/tests/sqlalchemy/test_migrate_cli.py
index 6ad693c..209dfab 100644
--- a/oslo_db/tests/sqlalchemy/test_migrate_cli.py
+++ b/oslo_db/tests/sqlalchemy/test_migrate_cli.py
@@ -12,6 +12,7 @@
 
 import mock
 from oslotest import base as test_base
+import sqlalchemy
 
 from oslo_db.sqlalchemy.migration_cli import ext_alembic
 from oslo_db.sqlalchemy.migration_cli import ext_migrate
@@ -35,7 +36,9 @@ class TestAlembicExtension(test_base.BaseTestCase):
     def setUp(self):
         self.migration_config = {'alembic_ini_path': '.',
                                  'db_url': 'sqlite://'}
-        self.alembic = ext_alembic.AlembicExtension(self.migration_config)
+        self.engine = sqlalchemy.create_engine(self.migration_config['db_url'])
+        self.alembic = ext_alembic.AlembicExtension(
+            self.engine, self.migration_config)
         super(TestAlembicExtension, self).setUp()
 
     def test_check_enabled_true(self, command):
@@ -52,7 +55,8 @@ class TestAlembicExtension(test_base.BaseTestCase):
         Verifies enabled returns False on empty alembic_ini_path variable
         """
         self.migration_config['alembic_ini_path'] = ''
-        alembic = ext_alembic.AlembicExtension(self.migration_config)
+        alembic = ext_alembic.AlembicExtension(
+            self.engine, self.migration_config)
         self.assertFalse(alembic.enabled)
 
     def test_upgrade_none(self, command):
@@ -98,7 +102,9 @@ class TestMigrateExtension(test_base.BaseTestCase):
     def setUp(self):
         self.migration_config = {'migration_repo_path': '.',
                                  'db_url': 'sqlite://'}
-        self.migrate = ext_migrate.MigrateExtension(self.migration_config)
+        self.engine = sqlalchemy.create_engine(self.migration_config['db_url'])
+        self.migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         super(TestMigrateExtension, self).setUp()
 
     def test_check_enabled_true(self, migration):
@@ -106,7 +112,8 @@ class TestMigrateExtension(test_base.BaseTestCase):
 
     def test_check_enabled_false(self, migration):
         self.migration_config['migration_repo_path'] = ''
-        migrate = ext_migrate.MigrateExtension(self.migration_config)
+        migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         self.assertFalse(migrate.enabled)
 
     def test_upgrade_head(self, migration):
@@ -143,7 +150,8 @@ class TestMigrateExtension(test_base.BaseTestCase):
 
     def test_change_init_version(self, migration):
         self.migration_config['init_version'] = 101
-        migrate = ext_migrate.MigrateExtension(self.migration_config)
+        migrate = ext_migrate.MigrateExtension(
+            self.engine, self.migration_config)
         migrate.downgrade(None)
         migration.db_sync.assert_called_once_with(
             migrate.engine,
@@ -158,8 +166,9 @@ class TestMigrationManager(test_base.BaseTestCase):
         self.migration_config = {'alembic_ini_path': '.',
                                  'migrate_repo_path': '.',
                                  'db_url': 'sqlite://'}
+        engine = sqlalchemy.create_engine(self.migration_config['db_url'])
         self.migration_manager = manager.MigrationManager(
-            self.migration_config)
+            self.migration_config, engine)
         self.ext = mock.Mock()
         self.ext.obj.version = mock.Mock(return_value=0)
         self.migration_manager._manager.extensions = [self.ext]
@@ -197,6 +206,13 @@ class TestMigrationManager(test_base.BaseTestCase):
         self.migration_manager.stamp('stamp')
         self.ext.obj.stamp.assert_called_once_with('stamp')
 
+    def test_wrong_config(self):
+        err = self.assertRaises(ValueError,
+                                manager.MigrationManager,
+                                {'wrong_key': 'sqlite://'})
+        self.assertEqual('Either database url or engine must be provided.',
+                         err.args[0])
+
 
 class TestMigrationRightOrder(test_base.BaseTestCase):
 
@@ -204,8 +220,9 @@ class TestMigrationRightOrder(test_base.BaseTestCase):
         self.migration_config = {'alembic_ini_path': '.',
                                  'migrate_repo_path': '.',
                                  'db_url': 'sqlite://'}
+        engine = sqlalchemy.create_engine(self.migration_config['db_url'])
         self.migration_manager = manager.MigrationManager(
-            self.migration_config)
+            self.migration_config, engine)
         self.first_ext = MockWithCmp()
         self.first_ext.obj.order = 1
         self.first_ext.obj.upgrade.return_value = 100