diff --git a/docs/source/whatsnew/0.8.4.txt b/docs/source/whatsnew/0.8.4.txt index 6acb0c89..fb053b13 100644 --- a/docs/source/whatsnew/0.8.4.txt +++ b/docs/source/whatsnew/0.8.4.txt @@ -46,7 +46,8 @@ then directs `lookup_symbol` to these dictionaries to find matching equities Maintenance and Refactorings ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -None +* Asset databases now contain version information to ensure compatibility + with current Zipline version (:issue:`815`). Build ~~~~~ diff --git a/tests/test_assets.py b/tests/test_assets.py index f081d939..1c8b6a33 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -29,25 +29,32 @@ from pandas.tseries.tools import normalize_date from pandas.util.testing import assert_frame_equal from numpy import full +import sqlalchemy as sa from zipline.assets import ( Asset, Equity, Future, AssetFinder, - AssetFinderCachedEquities + AssetFinderCachedEquities, ) - from zipline.assets.futures import ( cme_code_to_month, FutureChain, month_to_cme_code ) +from zipline.assets.asset_writer import ( + check_version_info, + write_version_info, + ASSET_DB_VERSION, + _version_table_schema, +) from zipline.errors import ( SymbolNotFound, MultipleSymbolsFound, SidAssignmentError, RootSymbolNotFound, + AssetDBVersionError, ) from zipline.finance.trading import TradingEnvironment, noop_load from zipline.utils.test_utils import ( @@ -1082,3 +1089,46 @@ class TestFutureChain(TestCase): } for key in codes: self.assertEqual(codes[key], month_to_cme_code(key)) + + +class TestAssetDBVersioning(TestCase): + + def test_check_version(self): + env = TradingEnvironment(load=noop_load) + version_table = env.asset_finder.version_info + + # This should not raise an error + check_version_info(version_table, ASSET_DB_VERSION) + + # This should fail because the version is too low + with self.assertRaises(AssetDBVersionError): + check_version_info(version_table, ASSET_DB_VERSION - 1) + + # This should fail because the version is too high + with self.assertRaises(AssetDBVersionError): + check_version_info(version_table, ASSET_DB_VERSION + 1) + + def test_write_version(self): + env = TradingEnvironment(load=noop_load) + metadata = sa.MetaData(bind=env.engine) + version_table = _version_table_schema(metadata) + version_table.delete().execute() + + # Assert that the version is not present in the table + self.assertIsNone(sa.select((version_table.c.version,)).scalar()) + + # This should fail because the table has no version info and is, + # therefore, consdered v0 + with self.assertRaises(AssetDBVersionError): + check_version_info(version_table, -2) + + # This should not raise an error because the version has been written + write_version_info(version_table, -2) + check_version_info(version_table, -2) + + # Assert that the version is in the table and correct + self.assertEqual(sa.select((version_table.c.version,)).scalar(), -2) + + # Assert that trying to overwrite the version fails + with self.assertRaises(sa.exc.IntegrityError): + write_version_info(version_table, -3) diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 5d6fc776..328b02e1 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -25,7 +25,7 @@ import numpy as np from six import with_metaclass import sqlalchemy as sa -from zipline.errors import SidAssignmentError +from zipline.errors import SidAssignmentError, AssetDBVersionError from zipline.assets._assets import Asset SQLITE_MAX_VARIABLE_NUMBER = 999 @@ -33,7 +33,6 @@ SQLITE_MAX_VARIABLE_NUMBER = 999 # Define a namedtuple for use with the load_data and _load_data methods AssetData = namedtuple('AssetData', 'equities futures exchanges root_symbols') - # Default values for the equities DataFrame _equities_defaults = { 'symbol': None, @@ -165,6 +164,194 @@ def _generate_output_dataframe(data_subset, defaults): return output +# Define a version number for the database generated by these writers +# Increment this version number any time a breaking change is made to the +# schema and readers of the database +ASSET_DB_VERSION = 0 + + +def check_version_info(version_table, expected_version): + """ + Checks for a version value in the version table. + + Parameters + ---------- + version_table : sa.Table + The version table of the asset database + expected_version : int + The expected version of the asset database + + Raises + ------ + AssetDBVersionError + If the version is in the table and not equal to ASSET_DB_VERSION. + """ + + # Read the version out of the table + version_from_table = sa.select((version_table.c.version,)).scalar() + + # A db without a version is considered v0 + if version_from_table is None: + version_from_table = 0 + + # Raise an error if the versions do not match + if (version_from_table != expected_version): + raise AssetDBVersionError(db_version=version_from_table, + expected_version=expected_version) + + +def write_version_info(version_table, version_value): + """ + Inserts the version value in to the version table. + + Parameters + ---------- + version_table : sa.Table + The version table of the asset database + version_value : int + The version to write in to the database + + """ + sa.insert(version_table, values={'version': version_value}).execute() + + +# A list of the names of all tables in the assets db +asset_db_table_names = ['version_info', 'equities', 'futures_exchanges', + 'futures_root_symbols', 'futures_contracts', + 'asset_router'] + + +def _equities_table_schema(metadata): + return sa.Table( + 'equities', + metadata, + sa.Column( + 'sid', + sa.Integer, + unique=True, + nullable=False, + primary_key=True, + ), + sa.Column('symbol', sa.Text), + sa.Column('company_symbol', sa.Text, index=True), + sa.Column('share_class_symbol', sa.Text), + sa.Column('fuzzy_symbol', sa.Text, index=True), + 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('exchange', sa.Text), + ) + + +def _futures_exchanges_schema(metadata): + return sa.Table( + 'futures_exchanges', + metadata, + sa.Column( + 'exchange', + sa.Text, + unique=True, + nullable=False, + primary_key=True, + ), + sa.Column('timezone', sa.Text), + ) + + +def _futures_root_symbols_schema(metadata, futures_exchanges): + return sa.Table( + 'futures_root_symbols', + metadata, + sa.Column( + 'root_symbol', + sa.Text, + unique=True, + nullable=False, + primary_key=True, + ), + sa.Column('root_symbol_id', sa.Integer), + sa.Column('sector', sa.Text), + sa.Column('description', sa.Text), + sa.Column( + 'exchange', + sa.Text, + sa.ForeignKey(futures_exchanges.c.exchange), + ), + ) + + +def _futures_contracts_schema(metadata, futures_root_symbols, + futures_exchanges): + return sa.Table( + 'futures_contracts', + metadata, + sa.Column( + 'sid', + sa.Integer, + unique=True, + nullable=False, + primary_key=True, + ), + sa.Column('symbol', sa.Text, unique=True, index=True), + sa.Column( + 'root_symbol', + sa.Text, + sa.ForeignKey(futures_root_symbols.c.root_symbol), + index=True + ), + 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( + 'exchange', + sa.Text, + sa.ForeignKey(futures_exchanges.c.exchange), + ), + sa.Column('notice_date', sa.Integer, nullable=False), + sa.Column('expiration_date', sa.Integer, nullable=False), + sa.Column('auto_close_date', sa.Integer, nullable=False), + sa.Column('contract_multiplier', sa.Float), + ) + + +def _asset_router_schema(metadata): + return sa.Table( + 'asset_router', + metadata, + sa.Column( + 'sid', + sa.Integer, + unique=True, + nullable=False, + primary_key=True), + sa.Column('asset_type', sa.Text), + ) + + +def _version_table_schema(metadata): + return sa.Table( + 'version_info', + metadata, + sa.Column( + 'id', + sa.Integer, + unique=True, + nullable=False, + primary_key=True, + ), + sa.Column( + 'version', + sa.Integer, + unique=True, + nullable=False, + ), + # This constraint ensures a single entry in this table + sa.CheckConstraint('id <= 1'), + ) + + class AssetDBWriter(with_metaclass(ABCMeta)): """ Class used to write arbitrary data to SQLite database. @@ -268,6 +455,21 @@ class AssetDBWriter(with_metaclass(ABCMeta)): def _write_equities(self, equities, bind): self._write_assets(equities, self.equities, 'equity', bind) + def check_for_tables(self, engine): + """ + Checks if any tables are present in the current assets database. + + Returns + ------- + bool + True if any tables are present, otherwise False. + """ + conn = engine.connect() + for table_name in asset_db_table_names: + if engine.dialect.has_table(conn, table_name): + return True + return False + def init_db(self, engine): """Connect to database and create tables. @@ -280,101 +482,29 @@ class AssetDBWriter(with_metaclass(ABCMeta)): """ metadata = sa.MetaData(bind=engine) - self.equities = sa.Table( - 'equities', - metadata, - sa.Column( - 'sid', - sa.Integer, - unique=True, - nullable=False, - primary_key=True, - ), - sa.Column('symbol', sa.Text), - sa.Column('company_symbol', sa.Text, index=True), - sa.Column('share_class_symbol', sa.Text), - sa.Column('fuzzy_symbol', sa.Text, index=True), - 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('exchange', sa.Text), + tables_already_exist = self.check_for_tables(engine=engine) + + self.equities = _equities_table_schema(metadata) + self.futures_exchanges = _futures_exchanges_schema(metadata) + self.futures_root_symbols = _futures_root_symbols_schema( + metadata=metadata, + futures_exchanges=self.futures_exchanges, ) - self.futures_exchanges = sa.Table( - 'futures_exchanges', - metadata, - sa.Column( - 'exchange', - sa.Text, - unique=True, - nullable=False, - primary_key=True, - ), - sa.Column('timezone', sa.Text), - ) - self.futures_root_symbols = sa.Table( - 'futures_root_symbols', - metadata, - sa.Column( - 'root_symbol', - sa.Text, - unique=True, - nullable=False, - primary_key=True, - ), - sa.Column('root_symbol_id', sa.Integer), - sa.Column('sector', sa.Text), - sa.Column('description', sa.Text), - sa.Column( - 'exchange', - sa.Text, - sa.ForeignKey(self.futures_exchanges.c.exchange), - ), - ) - self.futures_contracts = sa.Table( - 'futures_contracts', - metadata, - sa.Column( - 'sid', - sa.Integer, - unique=True, - nullable=False, - primary_key=True, - ), - sa.Column('symbol', sa.Text, unique=True, index=True), - sa.Column( - 'root_symbol', - sa.Text, - sa.ForeignKey(self.futures_root_symbols.c.root_symbol), - index=True - ), - 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( - 'exchange', - sa.Text, - sa.ForeignKey(self.futures_exchanges.c.exchange), - ), - sa.Column('notice_date', sa.Integer, nullable=False), - sa.Column('expiration_date', sa.Integer, nullable=False), - sa.Column('auto_close_date', sa.Integer, nullable=False), - sa.Column('contract_multiplier', sa.Float), - ) - self.asset_router = sa.Table( - 'asset_router', - metadata, - sa.Column( - 'sid', - sa.Integer, - unique=True, - nullable=False, - primary_key=True), - sa.Column('asset_type', sa.Text), + self.futures_contracts = _futures_contracts_schema( + metadata=metadata, + futures_root_symbols=self.futures_root_symbols, + futures_exchanges=self.futures_exchanges, ) + self.asset_router = _asset_router_schema(metadata) + self.version_info = _version_table_schema(metadata) # Create the SQL tables if they do not already exist. metadata.create_all(checkfirst=True) + + if tables_already_exist: + check_version_info(self.version_info, ASSET_DB_VERSION) + else: + write_version_info(self.version_info, ASSET_DB_VERSION) + return metadata def load_data(self): diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index fee1b35c..8d059072 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -37,6 +37,9 @@ from zipline.assets import ( ) from zipline.assets.asset_writer import ( split_delimited_symbol, + check_version_info, + ASSET_DB_VERSION, + asset_db_table_names, ) log = Logger('assets.py') @@ -79,15 +82,15 @@ class AssetFinder(object): self.engine = engine metadata = sa.MetaData(bind=engine) - - table_names = ['equities', 'futures_exchanges', 'futures_root_symbols', - 'futures_contracts', 'asset_router'] - metadata.reflect(only=table_names) - for table_name in table_names: + metadata.reflect(only=asset_db_table_names) + for table_name in asset_db_table_names: setattr(self, table_name, metadata.tables[table_name]) - # Cache for lookup of assets by sid, the objects in the asset lookp may - # be shared with the results from equity and future lookup caches. + # Check the version info of the db for compatibility + check_version_info(self.version_info, ASSET_DB_VERSION) + + # Cache for lookup of assets by sid, the objects in the asset lookup + # may be shared with the results from equity and future lookup caches. # # The top level cache exists to minimize lookups on the asset type # routing. diff --git a/zipline/errors.py b/zipline/errors.py index eacfcc64..aab1f414 100644 --- a/zipline/errors.py +++ b/zipline/errors.py @@ -439,3 +439,15 @@ class PositionTrackerMissingAssetFinder(ZiplineError): "not have an AssetFinder. This may be caused by a failure to properly " "de-serialize a TradingAlgorithm." ) + + +class AssetDBVersionError(ZiplineError): + """ + Raised by an AssetDBWriter or AssetFinder if the version number in the + versions table does not match the ASSET_DB_VERSION in asset_writer.py. + """ + msg = ( + "The existing Asset database has an incorrect version: {db_version}. " + "Expected version: {expected_version}. Try rebuilding your asset " + "database or updating your version of Zipline." + )