Refactor database migration manager to use given engine

Removing global engine object introduces some changes
to migration api.
Basically we need to move engine creation to outer scope.
In order to address these issues we need to refactor alembic api.

Co-Authored-By: Dmitry Shulyak <dshulyak@mirantis.com>

Adopted from Ia2f9586c4e9c31339dd7dcd9f45a58e8a8409921

Change-Id: Ic7dcf1d2f38388bd1e5a517ce60190bff499686a

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