diff --git a/tests/test_assets.py b/tests/test_assets.py index 3668a0cc..50842a06 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -1452,7 +1452,7 @@ class TestAssetDBVersioning(ZiplineTestCase): # first downgrade to v3 downgrade(self.engine, 3) metadata = sa.MetaData(conn) - metadata.reflect(bind=self.engine) + metadata.reflect() check_version_info(conn, metadata.tables['version_info'], 3) self.assertFalse('exchange_full' in metadata.tables) @@ -1461,9 +1461,9 @@ class TestAssetDBVersioning(ZiplineTestCase): # Verify that the db version is now 0 metadata = sa.MetaData(conn) - metadata.reflect(bind=self.engine) + metadata.reflect() version_table = metadata.tables['version_info'] - check_version_info(self.engine, version_table, 0) + check_version_info(conn, version_table, 0) # Check some of the v1-to-v0 downgrades self.assertTrue('futures_contracts' in metadata.tables) diff --git a/zipline/assets/asset_db_migrations.py b/zipline/assets/asset_db_migrations.py index 72d24b78..550e096c 100644 --- a/zipline/assets/asset_db_migrations.py +++ b/zipline/assets/asset_db_migrations.py @@ -21,39 +21,39 @@ def downgrade(engine, desired_version): """ # Check the version of the db at the engine - conn = engine.connect() - metadata = sa.MetaData(conn) - metadata.reflect(bind=engine) - version_info_table = metadata.tables['version_info'] - starting_version = sa.select((version_info_table.c.version,)).scalar() + with engine.connect() as conn: + metadata = sa.MetaData(conn) + metadata.reflect() + version_info_table = metadata.tables['version_info'] + starting_version = sa.select((version_info_table.c.version,)).scalar() - # Check for accidental upgrade - if starting_version < desired_version: - raise AssetDBImpossibleDowngrade(db_version=starting_version, - desired_version=desired_version) + # Check for accidental upgrade + if starting_version < desired_version: + raise AssetDBImpossibleDowngrade(db_version=starting_version, + desired_version=desired_version) - # Check if the desired version is already the db version - if starting_version == desired_version: - # No downgrade needed - return + # Check if the desired version is already the db version + if starting_version == desired_version: + # No downgrade needed + return - # Create alembic context - ctx = MigrationContext.configure(conn) - op = Operations(ctx) + # Create alembic context + ctx = MigrationContext.configure(conn) + op = Operations(ctx) - # Integer keys of downgrades to run - # E.g.: [5, 4, 3, 2] would downgrade v6 to v2 - downgrade_keys = range(desired_version, starting_version)[::-1] + # Integer keys of downgrades to run + # E.g.: [5, 4, 3, 2] would downgrade v6 to v2 + downgrade_keys = range(desired_version, starting_version)[::-1] - # Disable foreign keys until all downgrades are complete - _pragma_foreign_keys(conn, False) + # Disable foreign keys until all downgrades are complete + _pragma_foreign_keys(conn, False) - # Execute the downgrades in order - for downgrade_key in downgrade_keys: - _downgrade_methods[downgrade_key](op, engine, version_info_table) + # Execute the downgrades in order + for downgrade_key in downgrade_keys: + _downgrade_methods[downgrade_key](op, conn, version_info_table) - # Re-enable foreign keys - _pragma_foreign_keys(conn, True) + # Re-enable foreign keys + _pragma_foreign_keys(conn, True) def _pragma_foreign_keys(connection, on): @@ -96,10 +96,10 @@ def downgrades(src): @do(op.setitem(_downgrade_methods, destination)) @wraps(f) - def wrapper(op, engine, version_info_table): + def wrapper(op, conn, version_info_table): version_info_table.delete().execute() # clear the version f(op) - write_version_info(engine, version_info_table, destination) + write_version_info(conn, version_info_table, destination) return wrapper return _