MAINT: Clean up usage of engine versus connection

This commit is contained in:
Richard Frank
2016-08-23 10:07:58 -04:00
parent a788191b8e
commit 67dbeab820
2 changed files with 31 additions and 31 deletions
+3 -3
View File
@@ -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)
+28 -28
View File
@@ -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 _