Merge "Implement backend-specific drop_all_objects for provisioning."

5 files changed, 169 insertions, 26 deletions

diff --git a/oslo_db/sqlalchemy/compat/utils.py b/oslo_db/sqlalchemy/compat/utils.py
index fa6c3e7..e817718 100644
--- a/oslo_db/sqlalchemy/compat/utils.py
+++ b/oslo_db/sqlalchemy/compat/utils.py
@@ -24,3 +24,21 @@ sqla_097 = _SQLA_VERSION >= (0, 9, 7)
 sqla_094 = _SQLA_VERSION >= (0, 9, 4)
 sqla_090 = _SQLA_VERSION >= (0, 9, 0)
 sqla_08 = _SQLA_VERSION >= (0, 8)
+
+
+def get_postgresql_enums(conn):
+    """Return a list of ENUM type names on a Postgresql backend.
+
+    For SQLAlchemy 0.9 and lower, makes use of the semi-private
+    _load_enums() method of the Postgresql dialect.  In SQLAlchemy
+    1.0 this feature is supported using get_enums().
+
+    This function may only be called when the given connection
+    is against the Postgresql backend.  It will fail for other
+    kinds of backends.
+
+    """
+    if sqla_100:
+        return [e['name'] for e in sqlalchemy.inspect(conn).get_enums()]
+    else:
+        return conn.dialect._load_enums(conn).keys()
diff --git a/oslo_db/sqlalchemy/provision.py b/oslo_db/sqlalchemy/provision.py
index cce7025..ffb0cca 100644
--- a/oslo_db/sqlalchemy/provision.py
+++ b/oslo_db/sqlalchemy/provision.py
@@ -27,9 +27,11 @@ import six
 from six import moves
 import sqlalchemy
 from sqlalchemy.engine import url as sa_url
+from sqlalchemy import schema
 
 from oslo_db._i18n import _LI
 from oslo_db import exception
+from oslo_db.sqlalchemy.compat import utils as compat_utils
 from oslo_db.sqlalchemy import session
 from oslo_db.sqlalchemy import utils
 
@@ -56,6 +58,9 @@ class ProvisionedDatabase(object):
         self.backend.create_named_database(self.db_token)
         self.engine = self.backend.provisioned_engine(self.db_token)
 
+    def drop_all_objects(self):
+        self.backend.drop_all_objects(self.engine)
+
     def dispose(self):
         self.engine.dispose()
         self.backend.drop_named_database(self.db_token)
@@ -179,6 +184,15 @@ class Backend(object):
             self.engine, ident,
             conditional=conditional)
 
+    def drop_all_objects(self, engine):
+        """Drop all database objects.
+
+        Drops all database objects remaining on the default schema of the
+        given engine.
+
+        """
+        self.impl.drop_all_objects(engine)
+
     def database_exists(self, ident):
         """Return True if a database of the given name exists."""
 
@@ -246,6 +260,8 @@ class BackendImpl(object):
 
     default_engine_kwargs = {}
 
+    supports_drop_fk = True
+
     @classmethod
     def all_impls(cls):
         """Return an iterator of all possible BackendImpl objects.
@@ -294,6 +310,49 @@ class BackendImpl(object):
     def drop_named_database(self, engine, ident, conditional=False):
         """Drop a database with the given name."""
 
+    def drop_all_objects(self, engine):
+        """Drop all database objects.
+
+        Drops all database objects remaining on the default schema of the
+        given engine.
+
+        Per-db implementations will also need to drop items specific to those
+        systems, such as sequences, custom types (e.g. pg ENUM), etc.
+
+        """
+
+        with engine.begin() as conn:
+            inspector = sqlalchemy.inspect(engine)
+            metadata = schema.MetaData()
+            tbs = []
+            all_fks = []
+
+            for table_name in inspector.get_table_names():
+                fks = []
+                for fk in inspector.get_foreign_keys(table_name):
+                    # note that SQLite reflection does not have names
+                    # for foreign keys until SQLAlchemy 1.0
+                    if not fk['name']:
+                        continue
+                    fks.append(
+                        schema.ForeignKeyConstraint((), (), name=fk['name'])
+                        )
+                table = schema.Table(table_name, metadata, *fks)
+                tbs.append(table)
+                all_fks.extend(fks)
+
+            if self.supports_drop_fk:
+                for fkc in all_fks:
+                    conn.execute(schema.DropConstraint(fkc))
+
+            for table in tbs:
+                conn.execute(schema.DropTable(table))
+
+            self.drop_additional_objects(conn)
+
+    def drop_additional_objects(self, conn):
+        pass
+
     def provisioned_engine(self, base_url, ident):
         """Return a provisioned engine.
 
@@ -344,6 +403,9 @@ class MySQLBackendImpl(BackendImpl):
 
 @BackendImpl.impl.dispatch_for("sqlite")
 class SQLiteBackendImpl(BackendImpl):
+
+    supports_drop_fk = False
+
     def create_opportunistic_driver_url(self):
         return "sqlite://"
 
@@ -394,6 +456,12 @@ class PostgresqlBackendImpl(BackendImpl):
             else:
                 conn.execute("DROP DATABASE %s" % ident)
 
+    def drop_additional_objects(self, conn):
+        enums = compat_utils.get_postgresql_enums(conn)
+
+        for e in enums:
+            conn.execute("DROP TYPE %s" % e)
+
     def database_exists(self, engine, ident):
         return bool(
             engine.scalar(
diff --git a/oslo_db/sqlalchemy/test_base.py b/oslo_db/sqlalchemy/test_base.py
index aaff621..601a2c8 100644
--- a/oslo_db/sqlalchemy/test_base.py
+++ b/oslo_db/sqlalchemy/test_base.py
@@ -61,6 +61,7 @@ class DbFixture(fixtures.Fixture):
             msg = '%s backend is not available.' % self.DRIVER
             return self.test.skip(msg)
         else:
+            self.test.provision = self.provision
             self.test.engine = self.provision.engine
             self.addCleanup(setattr, self.test, 'engine', None)
             self.test.sessionmaker = session.get_maker(self.test.engine)
diff --git a/oslo_db/sqlalchemy/test_migrations.py b/oslo_db/sqlalchemy/test_migrations.py
index 9b65421..7627d21 100644
--- a/oslo_db/sqlalchemy/test_migrations.py
+++ b/oslo_db/sqlalchemy/test_migrations.py
@@ -25,9 +25,7 @@ import alembic.migration
 import pkg_resources as pkg
 import six
 import sqlalchemy
-from sqlalchemy.engine import reflection
 import sqlalchemy.exc
-from sqlalchemy import schema
 import sqlalchemy.sql.expression as expr
 import sqlalchemy.types as types
 
@@ -486,30 +484,7 @@ class ModelsMigrationsSync(object):
             return insp_def != "'%s'::character varying" % meta_def.arg
 
     def _cleanup(self):
-        engine = self.get_engine()
-        with engine.begin() as conn:
-            inspector = reflection.Inspector.from_engine(engine)
-            metadata = schema.MetaData()
-            tbs = []
-            all_fks = []
-
-            for table_name in inspector.get_table_names():
-                fks = []
-                for fk in inspector.get_foreign_keys(table_name):
-                    if not fk['name']:
-                        continue
-                    fks.append(
-                        schema.ForeignKeyConstraint((), (), name=fk['name'])
-                        )
-                table = schema.Table(table_name, metadata, *fks)
-                tbs.append(table)
-                all_fks.extend(fks)
-
-            for fkc in all_fks:
-                conn.execute(schema.DropConstraint(fkc))
-
-            for table in tbs:
-                conn.execute(schema.DropTable(table))
+        self.provision.drop_all_objects()
 
     FKInfo = collections.namedtuple('fk_info', ['constrained_columns',
                                                 'referred_table',
diff --git a/oslo_db/tests/sqlalchemy/test_provision.py b/oslo_db/tests/sqlalchemy/test_provision.py
new file mode 100644
index 0000000..7c57de3
--- /dev/null
+++ b/oslo_db/tests/sqlalchemy/test_provision.py
@@ -0,0 +1,81 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+
+from sqlalchemy import inspect
+from sqlalchemy import schema
+from sqlalchemy import types
+
+from oslo_db.sqlalchemy import test_base
+
+
+class DropAllObjectsTest(test_base.DbTestCase):
+
+    def setUp(self):
+        super(DropAllObjectsTest, self).setUp()
+
+        self.metadata = metadata = schema.MetaData()
+        schema.Table(
+            'a', metadata,
+            schema.Column('id', types.Integer, primary_key=True),
+            mysql_engine='InnoDB'
+        )
+        schema.Table(
+            'b', metadata,
+            schema.Column('id', types.Integer, primary_key=True),
+            schema.Column('a_id', types.Integer, schema.ForeignKey('a.id')),
+            mysql_engine='InnoDB'
+        )
+        schema.Table(
+            'c', metadata,
+            schema.Column('id', types.Integer, primary_key=True),
+            schema.Column('b_id', types.Integer, schema.ForeignKey('b.id')),
+            schema.Column(
+                'd_id', types.Integer,
+                schema.ForeignKey('d.id', use_alter=True, name='c_d_fk')),
+            mysql_engine='InnoDB'
+        )
+        schema.Table(
+            'd', metadata,
+            schema.Column('id', types.Integer, primary_key=True),
+            schema.Column('c_id', types.Integer, schema.ForeignKey('c.id')),
+            mysql_engine='InnoDB'
+        )
+
+        metadata.create_all(self.engine, checkfirst=False)
+        # will drop nothing if the test worked
+        self.addCleanup(metadata.drop_all, self.engine, checkfirst=True)
+
+    def test_drop_all(self):
+        insp = inspect(self.engine)
+        self.assertEqual(
+            set(['a', 'b', 'c', 'd']),
+            set(insp.get_table_names())
+        )
+
+        self.provision.drop_all_objects()
+
+        insp = inspect(self.engine)
+        self.assertEqual(
+            [],
+            insp.get_table_names()
+        )
+
+
+class MySQLRetainSchemaTest(
+        DropAllObjectsTest, test_base.MySQLOpportunisticTestCase):
+    pass
+
+
+class PostgresqlRetainSchemaTest(
+        DropAllObjectsTest, test_base.PostgreSQLOpportunisticTestCase):
+    pass