mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-09 11:50:09 +08:00
312 lines
9.9 KiB
Python
312 lines
9.9 KiB
Python
from functools import wraps
|
|
|
|
from alembic.migration import MigrationContext
|
|
from alembic.operations import Operations
|
|
import sqlalchemy as sa
|
|
from toolz.curried import do, operator as op
|
|
|
|
from zipline.assets.asset_writer import write_version_info
|
|
from zipline.errors import AssetDBImpossibleDowngrade
|
|
from zipline.utils.preprocess import preprocess
|
|
from zipline.utils.sqlite_utils import coerce_string_to_eng
|
|
|
|
|
|
@preprocess(engine=coerce_string_to_eng)
|
|
def downgrade(engine, desired_version):
|
|
"""Downgrades the assets db at the given engine to the desired version.
|
|
|
|
Parameters
|
|
----------
|
|
engine : Engine
|
|
An SQLAlchemy engine to the assets database.
|
|
desired_version : int
|
|
The desired resulting version for the assets database.
|
|
"""
|
|
|
|
# Check the version of the db at the engine
|
|
with engine.begin() 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 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)
|
|
|
|
# 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)
|
|
|
|
# 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)
|
|
|
|
|
|
def _pragma_foreign_keys(connection, on):
|
|
"""Sets the PRAGMA foreign_keys state of the SQLite database. Disabling
|
|
the pragma allows for batch modification of tables with foreign keys.
|
|
|
|
Parameters
|
|
----------
|
|
connection : Connection
|
|
A SQLAlchemy connection to the db
|
|
on : bool
|
|
If true, PRAGMA foreign_keys will be set to ON. Otherwise, the PRAGMA
|
|
foreign_keys will be set to OFF.
|
|
"""
|
|
connection.execute("PRAGMA foreign_keys=%s" % ("ON" if on else "OFF"))
|
|
|
|
|
|
# This dict contains references to downgrade methods that can be applied to an
|
|
# assets db. The resulting db's version is the key.
|
|
# e.g. The method at key '0' is the downgrade method from v1 to v0
|
|
_downgrade_methods = {}
|
|
|
|
|
|
def downgrades(src):
|
|
"""Decorator for marking that a method is a downgrade to a version to the
|
|
previous version.
|
|
|
|
Parameters
|
|
----------
|
|
src : int
|
|
The version this downgrades from.
|
|
|
|
Returns
|
|
-------
|
|
decorator : callable[(callable) -> callable]
|
|
The decorator to apply.
|
|
"""
|
|
def _(f):
|
|
destination = src - 1
|
|
|
|
@do(op.setitem(_downgrade_methods, destination))
|
|
@wraps(f)
|
|
def wrapper(op, conn, version_info_table):
|
|
conn.execute(version_info_table.delete()) # clear the version
|
|
f(op)
|
|
write_version_info(conn, version_info_table, destination)
|
|
|
|
return wrapper
|
|
return _
|
|
|
|
|
|
@downgrades(1)
|
|
def _downgrade_v1(op):
|
|
"""
|
|
Downgrade assets db by removing the 'tick_size' column and renaming the
|
|
'multiplier' column.
|
|
"""
|
|
# Drop indices before batch
|
|
# This is to prevent index collision when creating the temp table
|
|
op.drop_index('ix_futures_contracts_root_symbol')
|
|
op.drop_index('ix_futures_contracts_symbol')
|
|
|
|
# Execute batch op to allow column modification in SQLite
|
|
with op.batch_alter_table('futures_contracts') as batch_op:
|
|
|
|
# Rename 'multiplier'
|
|
batch_op.alter_column(column_name='multiplier',
|
|
new_column_name='contract_multiplier')
|
|
|
|
# Delete 'tick_size'
|
|
batch_op.drop_column('tick_size')
|
|
|
|
# Recreate indices after batch
|
|
op.create_index('ix_futures_contracts_root_symbol',
|
|
table_name='futures_contracts',
|
|
columns=['root_symbol'])
|
|
op.create_index('ix_futures_contracts_symbol',
|
|
table_name='futures_contracts',
|
|
columns=['symbol'],
|
|
unique=True)
|
|
|
|
|
|
@downgrades(2)
|
|
def _downgrade_v2(op):
|
|
"""
|
|
Downgrade assets db by removing the 'auto_close_date' column.
|
|
"""
|
|
# Drop indices before batch
|
|
# This is to prevent index collision when creating the temp table
|
|
op.drop_index('ix_equities_fuzzy_symbol')
|
|
op.drop_index('ix_equities_company_symbol')
|
|
|
|
# Execute batch op to allow column modification in SQLite
|
|
with op.batch_alter_table('equities') as batch_op:
|
|
batch_op.drop_column('auto_close_date')
|
|
|
|
# Recreate indices after batch
|
|
op.create_index('ix_equities_fuzzy_symbol',
|
|
table_name='equities',
|
|
columns=['fuzzy_symbol'])
|
|
op.create_index('ix_equities_company_symbol',
|
|
table_name='equities',
|
|
columns=['company_symbol'])
|
|
|
|
|
|
@downgrades(3)
|
|
def _downgrade_v3(op):
|
|
"""
|
|
Downgrade assets db by adding a not null constraint on
|
|
``equities.first_traded``
|
|
"""
|
|
op.create_table(
|
|
'_new_equities',
|
|
sa.Column(
|
|
'sid',
|
|
sa.Integer,
|
|
unique=True,
|
|
nullable=False,
|
|
primary_key=True,
|
|
),
|
|
sa.Column('symbol', sa.Text),
|
|
sa.Column('company_symbol', sa.Text),
|
|
sa.Column('share_class_symbol', sa.Text),
|
|
sa.Column('fuzzy_symbol', sa.Text),
|
|
sa.Column('asset_name', sa.Text),
|
|
sa.Column('start_date', sa.Integer, default=0, nullable=False),
|
|
sa.Column('end_date', sa.Integer, nullable=False),
|
|
sa.Column('first_traded', sa.Integer, nullable=False),
|
|
sa.Column('auto_close_date', sa.Integer),
|
|
sa.Column('exchange', sa.Text),
|
|
)
|
|
op.execute(
|
|
"""
|
|
insert into _new_equities
|
|
select * from equities
|
|
where equities.first_traded is not null
|
|
""",
|
|
)
|
|
op.drop_table('equities')
|
|
op.rename_table('_new_equities', 'equities')
|
|
# we need to make sure the indices have the proper names after the rename
|
|
op.create_index(
|
|
'ix_equities_company_symbol',
|
|
'equities',
|
|
['company_symbol'],
|
|
)
|
|
op.create_index(
|
|
'ix_equities_fuzzy_symbol',
|
|
'equities',
|
|
['fuzzy_symbol'],
|
|
)
|
|
|
|
|
|
@downgrades(4)
|
|
def _downgrade_v4(op):
|
|
"""
|
|
Downgrades assets db by copying the `exchange_full` column to `exchange`,
|
|
then dropping the `exchange_full` column.
|
|
"""
|
|
op.drop_index('ix_equities_fuzzy_symbol')
|
|
op.drop_index('ix_equities_company_symbol')
|
|
|
|
op.execute("UPDATE equities SET exchange = exchange_full")
|
|
|
|
with op.batch_alter_table('equities') as batch_op:
|
|
batch_op.drop_column('exchange_full')
|
|
|
|
op.create_index('ix_equities_fuzzy_symbol',
|
|
table_name='equities',
|
|
columns=['fuzzy_symbol'])
|
|
op.create_index('ix_equities_company_symbol',
|
|
table_name='equities',
|
|
columns=['company_symbol'])
|
|
|
|
|
|
@downgrades(5)
|
|
def _downgrade_v5(op):
|
|
op.create_table(
|
|
'_new_equities',
|
|
sa.Column(
|
|
'sid',
|
|
sa.Integer,
|
|
unique=True,
|
|
nullable=False,
|
|
primary_key=True,
|
|
),
|
|
sa.Column('symbol', sa.Text),
|
|
sa.Column('company_symbol', sa.Text),
|
|
sa.Column('share_class_symbol', sa.Text),
|
|
sa.Column('fuzzy_symbol', sa.Text),
|
|
sa.Column('asset_name', sa.Text),
|
|
sa.Column('start_date', sa.Integer, default=0, nullable=False),
|
|
sa.Column('end_date', sa.Integer, nullable=False),
|
|
sa.Column('first_traded', sa.Integer),
|
|
sa.Column('auto_close_date', sa.Integer),
|
|
sa.Column('exchange', sa.Text),
|
|
sa.Column('exchange_full', sa.Text)
|
|
)
|
|
|
|
op.execute(
|
|
"""
|
|
insert into _new_equities
|
|
select
|
|
equities.sid as sid,
|
|
sym.symbol as symbol,
|
|
sym.company_symbol as company_symbol,
|
|
sym.share_class_symbol as share_class_symbol,
|
|
sym.company_symbol || sym.share_class_symbol as fuzzy_symbol,
|
|
equities.asset_name as asset_name,
|
|
equities.start_date as start_date,
|
|
equities.end_date as end_date,
|
|
equities.first_traded as first_traded,
|
|
equities.auto_close_date as auto_close_date,
|
|
equities.exchange as exchange,
|
|
equities.exchange_full as exchange_full
|
|
from
|
|
equities
|
|
inner join
|
|
-- Nested select here to take the most recently held ticker
|
|
-- for each sid. The group by with no aggregation function will
|
|
-- take the last element in the group, so we first order by
|
|
-- the end date ascending to ensure that the groupby takes
|
|
-- the last ticker.
|
|
(select
|
|
*
|
|
from
|
|
(select
|
|
*
|
|
from
|
|
equity_symbol_mappings
|
|
order by
|
|
equity_symbol_mappings.end_date asc)
|
|
group by
|
|
sid) sym
|
|
on
|
|
equities.sid == sym.sid
|
|
""",
|
|
)
|
|
op.drop_table('equity_symbol_mappings')
|
|
op.drop_table('equities')
|
|
op.rename_table('_new_equities', 'equities')
|
|
# we need to make sure the indicies have the proper names after the rename
|
|
op.create_index(
|
|
'ix_equities_company_symbol',
|
|
'equities',
|
|
['company_symbol'],
|
|
)
|
|
op.create_index(
|
|
'ix_equities_fuzzy_symbol',
|
|
'equities',
|
|
['fuzzy_symbol'],
|
|
)
|