From 6aac54544e95799d8551da27de94fc77720e02f8 Mon Sep 17 00:00:00 2001 From: jfkirk Date: Mon, 9 Nov 2015 14:37:01 -0500 Subject: [PATCH] MAINT: Only write asset db version on db creation --- tests/test_assets.py | 13 +++++---- zipline/assets/asset_writer.py | 52 ++++++++++++++++++++-------------- zipline/assets/assets.py | 5 +--- 3 files changed, 39 insertions(+), 31 deletions(-) diff --git a/tests/test_assets.py b/tests/test_assets.py index 6954db47..3ca3b4ae 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -1097,7 +1097,8 @@ class TestAssetDBVersioning(TestCase): env = TradingEnvironment(load=noop_load) version_table = env.asset_finder.version - self.assertTrue(check_version_info(version_table, ASSET_DB_VERSION)) + # 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): @@ -1116,12 +1117,14 @@ class TestAssetDBVersioning(TestCase): # Assert that the version is not present in the table self.assertIsNone(sa.select((version_table.c.version,)).scalar()) - # This should return false because there is no version info in the db - self.assertFalse(check_version_info(version_table, -2)) + # 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 return true because the version has been written + # This should not raise an error because the version has been written write_version_info(version_table, -2) - self.assertTrue(check_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) diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 2e0d2ba3..7d4b6011 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -33,6 +33,9 @@ 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') +# A list of the names of all tables in the assets db +table_names = ['version', 'equities', 'futures_exchanges', + 'futures_root_symbols', 'futures_contracts', 'asset_router'] # Default values for the equities DataFrame _equities_defaults = { @@ -182,12 +185,6 @@ def check_version_info(version_table, expected_version): expected_version : int The expected version of the asset database - Returns - ------- - bool - True if the version information is present and correct. - False if the version information is missing. - Raises ------ AssetDBVersionError @@ -197,19 +194,14 @@ def check_version_info(version_table, expected_version): # Read the version out of the table version_from_table = sa.select((version_table.c.version,)).scalar() - # If a version exists in the table... - if version_from_table is not None: + # 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) - - # A matching version was found - return True - - # No version info found - return False + # 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): @@ -352,6 +344,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 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. @@ -364,6 +371,8 @@ class AssetDBWriter(with_metaclass(ABCMeta)): """ metadata = sa.MetaData(bind=engine) + tables_already_exist = self.check_for_tables(engine=engine) + self.equities = sa.Table( 'equities', metadata, @@ -461,10 +470,9 @@ class AssetDBWriter(with_metaclass(ABCMeta)): # Create the SQL tables if they do not already exist. metadata.create_all(checkfirst=True) - # Check if the version is mismatched or, if it is not present, add it - version_found = check_version_info(self.version_table, - ASSET_DB_VERSION) - if not version_found: + if tables_already_exist: + check_version_info(self.version_table, ASSET_DB_VERSION) + else: write_version_info(self.version_table, ASSET_DB_VERSION) return metadata diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index f6a62cf8..f08c24d6 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -39,6 +39,7 @@ from zipline.assets.asset_writer import ( split_delimited_symbol, check_version_info, ASSET_DB_VERSION, + table_names, ) log = Logger('assets.py') @@ -81,10 +82,6 @@ class AssetFinder(object): self.engine = engine metadata = sa.MetaData(bind=engine) - - table_names = ['version', 'equities', 'futures_exchanges', - 'futures_root_symbols', 'futures_contracts', - 'asset_router'] metadata.reflect(only=table_names) for table_name in table_names: setattr(self, table_name, metadata.tables[table_name])