diff options
Diffstat (limited to 'lib/sqlalchemy/testing/util.py')
-rw-r--r-- | lib/sqlalchemy/testing/util.py | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/lib/sqlalchemy/testing/util.py b/lib/sqlalchemy/testing/util.py index 7b3f721a6..8230f923a 100644 --- a/lib/sqlalchemy/testing/util.py +++ b/lib/sqlalchemy/testing/util.py @@ -147,6 +147,10 @@ def run_as_contextmanager(ctx, fn, *arg, **kw): simulating the behavior of 'with' to support older Python versions. + This is not necessary anymore as we have placed 2.6 + as minimum Python version, however some tests are still using + this structure. + """ obj = ctx.__enter__() @@ -194,6 +198,25 @@ def provide_metadata(fn, *args, **kw): self.metadata = prev_meta +def force_drop_names(*names): + """Force the given table names to be dropped after test complete, + isolating for foreign key cycles + + """ + from . import config + from sqlalchemy import inspect + + @decorator + def go(fn, *args, **kw): + + try: + return fn(*args, **kw) + finally: + drop_all_tables( + config.db, inspect(config.db), include_names=names) + return go + + class adict(dict): """Dict keys available as attributes. Shadows.""" @@ -207,3 +230,39 @@ class adict(dict): return tuple([self[key] for key in keys]) get_all = __call__ + + +def drop_all_tables(engine, inspector, schema=None, include_names=None): + from sqlalchemy import Column, Table, Integer, MetaData, \ + ForeignKeyConstraint + from sqlalchemy.schema import DropTable, DropConstraint + + if include_names is not None: + include_names = set(include_names) + + with engine.connect() as conn: + for tname, fkcs in reversed( + inspector.get_sorted_table_and_fkc_names(schema=schema)): + if tname: + if include_names is not None and tname not in include_names: + continue + conn.execute(DropTable( + Table(tname, MetaData()) + )) + elif fkcs: + if not engine.dialect.supports_alter: + continue + for tname, fkc in fkcs: + if include_names is not None and \ + tname not in include_names: + continue + tb = Table( + tname, MetaData(), + Column('x', Integer), + Column('y', Integer), + schema=schema + ) + conn.execute(DropConstraint( + ForeignKeyConstraint( + [tb.c.x], [tb.c.y], name=fkc) + )) |