Merge "Add index modifying methods"

oslo/db/sqlalchemy/utils.py

@@ -654,3 +654,73 @@ def get_db_connection_info(conn_pieces):
         password = auth_pieces[1].strip()
 
     return (user, password, database, host)
+
+
+def index_exists(migrate_engine, table_name, index_name):
+    """Check if given index exists.
+
+    :param migrate_engine: sqlalchemy engine
+    :param table_name:     name of the table
+    :param index_name:     name of the index
+    """
+    inspector = reflection.Inspector.from_engine(migrate_engine)
+    indexes = inspector.get_indexes(table_name)
+    index_names = [index['name'] for index in indexes]
+    return index_name in index_names
+
+
+def add_index(migrate_engine, table_name, index_name, idx_columns):
+    """Create an index for given columns.
+
+    :param migrate_engine: sqlalchemy engine
+    :param table_name:     name of the table
+    :param index_name:     name of the index
+    :param idx_columns:    tuple with names of columns that will be indexed
+    """
+    table = get_table(migrate_engine, table_name)
+    if not index_exists(migrate_engine, table_name, index_name):
+        index = Index(
+            index_name, *[getattr(table.c, col) for col in idx_columns]
+        )
+        index.create()
+    else:
+        raise ValueError("Index '%s' already exists!" % index_name)
+
+
+def drop_index(migrate_engine, table_name, index_name):
+    """Drop index with given name.
+
+    :param migrate_engine: sqlalchemy engine
+    :param table_name:     name of the table
+    :param index_name:     name of the index
+    """
+    table = get_table(migrate_engine, table_name)
+    for index in table.indexes:
+        if index.name == index_name:
+            index.drop()
+            break
+    else:
+        raise ValueError("Index '%s' not found!" % index_name)
+
+
+def change_index_columns(migrate_engine, table_name, index_name, new_columns):
+    """Change set of columns that are indexed by given index.
+
+    :param migrate_engine: sqlalchemy engine
+    :param table_name:     name of the table
+    :param index_name:     name of the index
+    :param new_columns:    tuple with names of columns that will be indexed
+    """
+    drop_index(migrate_engine, table_name, index_name)
+    add_index(migrate_engine, table_name, index_name, new_columns)
+
+
+def column_exists(engine, table_name, column):
+    """Check if table has given column.
+
+    :param engine:         sqlalchemy engine
+    :param table_name:     name of the table
+    :param column:         name of the colmn
+    """
+    t = get_table(engine, table_name)
+    return column in t.c

tests/sqlalchemy/test_utils.py

@@ -37,6 +37,7 @@ from oslo.db import exception
 from oslo.db.openstack.common.fixture import moxstubout
 from oslo.db.sqlalchemy import models
 from oslo.db.sqlalchemy import session
+from oslo.db.sqlalchemy import test_base as db_test_base
 from oslo.db.sqlalchemy import test_migrations
 from oslo.db.sqlalchemy import utils
 from tests import utils as test_utils
@@ -843,3 +844,73 @@ class TestModelQuery(test_base.BaseTestCase):
             self.session.query, MyModel, self.user_context.read_deleted)
         _project_filter.assert_called_with(
             self.session.query, MyModel, self.user_context, False)
+
+
+class TestUtils(db_test_base.DbTestCase):
+    def setUp(self):
+        super(TestUtils, self).setUp()
+        meta = MetaData(bind=self.engine)
+        self.test_table = Table(
+            'test_table',
+            meta,
+            Column('a', Integer),
+            Column('b', Integer)
+        )
+        self.test_table.create()
+        self.addCleanup(meta.drop_all)
+
+    def test_index_exists(self):
+        self.assertFalse(utils.index_exists(self.engine, 'test_table',
+                                            'new_index'))
+        Index('new_index', self.test_table.c.a).create(self.engine)
+        self.assertTrue(utils.index_exists(self.engine, 'test_table',
+                                           'new_index'))
+
+    def test_add_index(self):
+        self.assertFalse(utils.index_exists(self.engine, 'test_table',
+                                            'new_index'))
+        utils.add_index(self.engine, 'test_table', 'new_index', ('a',))
+        self.assertTrue(utils.index_exists(self.engine, 'test_table',
+                                           'new_index'))
+
+    def test_add_existing_index(self):
+        Index('new_index', self.test_table.c.a).create(self.engine)
+        self.assertRaises(ValueError, utils.add_index, self.engine,
+                          'test_table', 'new_index', ('a',))
+
+    def test_drop_index(self):
+        Index('new_index', self.test_table.c.a).create(self.engine)
+        utils.drop_index(self.engine, 'test_table', 'new_index')
+        self.assertFalse(utils.index_exists(self.engine, 'test_table',
+                         'new_index'))
+
+    def test_drop_unexisting_index(self):
+        self.assertRaises(ValueError, utils.drop_index, self.engine,
+                          'test_table', 'new_index')
+
+    @mock.patch('oslo.db.sqlalchemy.utils.drop_index')
+    @mock.patch('oslo.db.sqlalchemy.utils.add_index')
+    def test_change_index_columns(self, add_index, drop_index):
+        utils.change_index_columns(self.engine, 'test_table', 'a_index',
+                                  ('a',))
+        utils.drop_index.assert_called_once_with(self.engine, 'test_table',
+                                                 'a_index')
+        utils.add_index.assert_called_once_with(self.engine, 'test_table',
+                                                'a_index', ('a',))
+
+    def test_column_exists(self):
+        for col in ['a', 'b']:
+            self.assertTrue(utils.column_exists(self.engine, 'test_table',
+                                                col))
+        self.assertFalse(utils.column_exists(self.engine, 'test_table',
+                                             'fake_column'))
+
+
+class TestUtilsMysqlOpportunistically(
+        TestUtils, db_test_base.MySQLOpportunisticTestCase):
+    pass
+
+
+class TestUtilsPostgresqlOpportunistically(
+        TestUtils, db_test_base.PostgreSQLOpportunisticTestCase):
+    pass