mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-06 05:14:38 +08:00
MAINT: Clean up usage of engine versus connection
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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 _
|
||||
|
||||
Reference in New Issue
Block a user