From 419bd1e3b52e5c7d31d2af75bccc35f37323e804 Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Tue, 23 Aug 2016 12:16:50 -0400 Subject: [PATCH] DEV: zipline ingest can downgrade the assets db This lets us publish an "old" db version for the most recent release of zipline, using the latest code base. --- tests/data/bundles/test_core.py | 107 ++++++++++++++++++++++++++----- zipline/__main__.py | 21 +++--- zipline/data/bundles/__init__.py | 2 + zipline/data/bundles/core.py | 45 ++++++++++--- zipline/utils/cache.py | 6 +- 5 files changed, 143 insertions(+), 38 deletions(-) diff --git a/tests/data/bundles/test_core.py b/tests/data/bundles/test_core.py index 870945b4..e6cf65ad 100644 --- a/tests/data/bundles/test_core.py +++ b/tests/data/bundles/test_core.py @@ -2,12 +2,17 @@ import os from nose_parameterized import parameterized import pandas as pd +import sqlalchemy as sa from toolz import valmap import toolz.curried.operator as op +from zipline.assets import ASSET_DB_VERSION +from zipline.assets.asset_writer import check_version_info from zipline.assets.synthetic import make_simple_equity_info -from zipline.data.bundles import UnknownBundle, from_bundle_ingest_dirname -from zipline.data.bundles.core import _make_bundle_core, BadClean +from zipline.data.bundles import UnknownBundle, from_bundle_ingest_dirname, \ + ingestions_for_bundle +from zipline.data.bundles.core import _make_bundle_core, BadClean, \ + to_bundle_ingest_dirname, asset_db_path from zipline.lib.adjustment import Float64Multiply from zipline.pipeline.loaders.synthetic import ( make_bar_data, @@ -17,7 +22,8 @@ from zipline.testing import ( subtest, str_to_seconds, ) -from zipline.testing.fixtures import WithInstanceTmpDir, ZiplineTestCase +from zipline.testing.fixtures import WithInstanceTmpDir, ZiplineTestCase, \ + WithDefaultDateBounds from zipline.testing.predicates import ( assert_equal, assert_false, @@ -37,7 +43,13 @@ import zipline.utils.paths as pth _1_ns = pd.Timedelta(1, unit='ns') -class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase): +class BundleCoreTestCase(WithInstanceTmpDir, + WithDefaultDateBounds, + ZiplineTestCase): + + START_DATE = pd.Timestamp('2014-01-06', tz='utc') + END_DATE = pd.Timestamp('2014-01-10', tz='utc') + def init_instance_fixtures(self): super(BundleCoreTestCase, self).init_instance_fixtures() (self.bundles, @@ -111,18 +123,17 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase): assert_true(called[0]) def test_ingest(self): - start = pd.Timestamp('2014-01-06', tz='utc') - end = pd.Timestamp('2014-01-10', tz='utc') calendar = get_calendar('NYSE') - - sessions = calendar.sessions_in_range(start, end) - minutes = calendar.minutes_for_sessions_in_range(start, end) + sessions = calendar.sessions_in_range(self.START_DATE, self.END_DATE) + minutes = calendar.minutes_for_sessions_in_range( + self.START_DATE, self.END_DATE, + ) sids = tuple(range(3)) equities = make_simple_equity_info( sids, - start, - end, + self.START_DATE, + self.END_DATE, ) daily_bar_data = make_bar_data(equities, sessions) @@ -145,8 +156,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase): @self.register( 'bundle', calendar=calendar, - start_session=start, - end_session=end, + start_session=self.START_DATE, + end_session=self.END_DATE, ) def bundle_ingest(environ, asset_db_writer, @@ -193,8 +204,8 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase): actual = bundle.equity_daily_bar_reader.load_raw_arrays( columns, - start, - end, + self.START_DATE, + self.END_DATE, sids, ) for actual_column, colname in zip(actual, columns): @@ -253,6 +264,72 @@ class BundleCoreTestCase(WithInstanceTmpDir, ZiplineTestCase): msg='volume', ) + def test_ingest_assets_versions(self): + versions = (1, 2) + + called = [False] + + @self.register('bundle', create_writers=False) + def bundle_ingest_no_create_writers(*args, **kwargs): + called[0] = True + + now = pd.Timestamp.utcnow() + with self.assertRaisesRegexp(ValueError, + "ingest .* creates writers .* downgrade"): + self.ingest('bundle', self.environ, assets_versions=versions, + timestamp=now - pd.Timedelta(seconds=1)) + assert_false(called[0]) + assert_equal(len(ingestions_for_bundle('bundle', self.environ)), 1) + + @self.register('bundle', create_writers=True) + def bundle_ingest_create_writers( + environ, + asset_db_writer, + minute_bar_writer, + daily_bar_writer, + adjustment_writer, + calendar, + start_session, + end_session, + cache, + show_progress, + output_dir): + self.assertIsNotNone(asset_db_writer) + self.assertIsNotNone(minute_bar_writer) + self.assertIsNotNone(daily_bar_writer) + self.assertIsNotNone(adjustment_writer) + + equities = make_simple_equity_info( + tuple(range(3)), + self.START_DATE, + self.END_DATE, + ) + asset_db_writer.write(equities=equities) + called[0] = True + + # Explicitly use different timestamp; otherwise, test could run so fast + # that first ingestion is re-used. + self.ingest('bundle', self.environ, assets_versions=versions, + timestamp=now) + assert_true(called[0]) + + ingestions = ingestions_for_bundle('bundle', self.environ) + assert_equal(len(ingestions), 2) + for version in sorted(set(versions) | {ASSET_DB_VERSION}): + eng = sa.create_engine( + 'sqlite:///' + + asset_db_path( + 'bundle', + to_bundle_ingest_dirname(ingestions[0]), # most recent + self.environ, + version, + ) + ) + metadata = sa.MetaData() + metadata.reflect(eng) + version_table = metadata.tables['version_info'] + check_version_info(eng, version_table, version) + @parameterized.expand([('clean',), ('load',)]) def test_bundle_doesnt_exist(self, fnname): with assert_raises(UnknownBundle) as e: diff --git a/zipline/__main__.py b/zipline/__main__.py index 3dde2b59..93c26bf4 100644 --- a/zipline/__main__.py +++ b/zipline/__main__.py @@ -5,10 +5,10 @@ from functools import wraps import click import logbook import pandas as pd +from six import text_type from zipline.data import bundles as bundles_module from zipline.utils.cli import Date, Timestamp -import zipline.utils.paths as pth from zipline.utils.run_algo import _run, load_extensions try: @@ -284,19 +284,25 @@ def zipline_magic(line, cell=None): show_default=True, help='The data bundle to ingest.', ) +@click.option( + '--assets-version', + type=int, + multiple=True, + help='Version of the assets db to which to downgrade.', +) @click.option( '--show-progress/--no-show-progress', - is_flag=True, default=True, help='Print progress information to the terminal.' ) -def ingest(bundle, show_progress): +def ingest(bundle, assets_version, show_progress): """Ingest the data for the given bundle. """ bundles_module.ingest( bundle, os.environ, pd.Timestamp.utcnow(), + assets_version, show_progress, ) @@ -352,11 +358,8 @@ def bundles(): # hide the test data continue try: - ingestions = sorted( - (str(bundles_module.from_bundle_ingest_dirname(ing)) - for ing in os.listdir(pth.data_path([bundle])) - if not pth.hidden(ing)), - reverse=True, + ingestions = list( + map(text_type, bundles_module.ingestions_for_bundle(bundle)) ) except OSError as e: if e.errno != errno.ENOENT: @@ -367,7 +370,7 @@ def bundles(): # because there were no entries, print a single message indicating that # no ingestions have yet been made. for timestamp in ingestions or [""]: - print("%s %s" % (bundle, timestamp)) + click.echo("%s %s" % (bundle, timestamp)) if __name__ == '__main__': diff --git a/zipline/data/bundles/__init__.py b/zipline/data/bundles/__init__.py index beea33bf..40746c6c 100644 --- a/zipline/data/bundles/__init__.py +++ b/zipline/data/bundles/__init__.py @@ -6,6 +6,7 @@ from .core import ( clean, from_bundle_ingest_dirname, ingest, + ingestions_for_bundle, load, register, to_bundle_ingest_dirname, @@ -20,6 +21,7 @@ __all__ = [ 'clean', 'from_bundle_ingest_dirname', 'ingest', + 'ingestions_for_bundle', 'load', 'register', 'to_bundle_ingest_dirname', diff --git a/zipline/data/bundles/core.py b/zipline/data/bundles/core.py index 33b1acd0..00c30a4e 100644 --- a/zipline/data/bundles/core.py +++ b/zipline/data/bundles/core.py @@ -20,9 +20,11 @@ from ..minute_bars import ( BcolzMinuteBarWriter, ) from zipline.assets import AssetDBWriter, AssetFinder, ASSET_DB_VERSION +from zipline.assets.asset_db_migrations import downgrade from zipline.utils.cache import ( dataframe_cache, working_dir, + working_file, ) from zipline.utils.compat import mappingproxy from zipline.utils.input_validation import ensure_timestamp, optionally @@ -31,9 +33,9 @@ from zipline.utils.preprocess import preprocess from zipline.utils.calendars import get_calendar, register_calendar -def asset_db_path(bundle_name, timestr, environ=None): +def asset_db_path(bundle_name, timestr, environ=None, db_version=None): return pth.data_path( - asset_db_relative(bundle_name, timestr, environ), + asset_db_relative(bundle_name, timestr, environ, db_version), environ=environ, ) @@ -82,8 +84,10 @@ def minute_equity_relative(bundle_name, timestr, environ=None): return bundle_name, timestr, 'minute_equities.bcolz' -def asset_db_relative(bundle_name, timestr, environ=None): - return bundle_name, timestr, 'assets-%d.sqlite' % ASSET_DB_VERSION +def asset_db_relative(bundle_name, timestr, environ=None, db_version=None): + db_version = ASSET_DB_VERSION if db_version is None else db_version + + return bundle_name, timestr, 'assets-%d.sqlite' % db_version def to_bundle_ingest_dirname(ts): @@ -119,6 +123,15 @@ def from_bundle_ingest_dirname(cs): return pd.Timestamp(cs.replace(';', ':')) +def ingestions_for_bundle(bundle, environ=None): + return sorted( + (from_bundle_ingest_dirname(ing) + for ing in os.listdir(pth.data_path([bundle], environ)) + if not pth.hidden(ing)), + reverse=True, + ) + + _BundlePayload = namedtuple( '_BundlePayload', 'calendar start_session end_session minutes_per_day ingest create_writers', @@ -328,6 +341,7 @@ def _make_bundle_core(): def ingest(name, environ=os.environ, timestamp=None, + assets_versions=(), show_progress=False): """Ingest data for a given bundle. @@ -340,6 +354,8 @@ def _make_bundle_core(): timestamp : datetime, optional The timestamp to use for the load. By default this is the current time. + assets_versions : Iterable[int], optional + Versions of the assets db to which to downgrade. show_progress : bool, optional Tell the ingest function to display the progress where possible. """ @@ -389,11 +405,10 @@ def _make_bundle_core(): bundle.end_session, minutes_per_day=bundle.minutes_per_day, ) - asset_db_writer = AssetDBWriter( - wd.getpath(*asset_db_relative( - name, timestr, environ=environ, - )) - ) + assets_db_path = wd.getpath(*asset_db_relative( + name, timestr, environ=environ, + )) + asset_db_writer = AssetDBWriter(assets_db_path) adjustment_db_writer = stack.enter_context( SQLiteAdjustmentWriter( @@ -409,6 +424,10 @@ def _make_bundle_core(): minute_bar_writer = None asset_db_writer = None adjustment_db_writer = None + if assets_versions: + raise ValueError('Need to ingest a bundle that creates ' + 'writers in order to downgrade the assets' + ' db.') bundle.ingest( environ, asset_db_writer, @@ -423,6 +442,14 @@ def _make_bundle_core(): pth.data_path([name, timestr], environ=environ), ) + for version in sorted(set(assets_versions), reverse=True): + version_path = wd.getpath(*asset_db_relative( + name, timestr, environ=environ, db_version=version, + )) + with working_file(version_path) as wf: + shutil.copy2(assets_db_path, wf.path) + downgrade(wf.path, version) + def most_recent_data(bundle_name, timestamp, environ=None): """Get the path to the most recent data after ``date``for the given bundle. diff --git a/zipline/utils/cache.py b/zipline/utils/cache.py index 21aee7d4..427a67d6 100644 --- a/zipline/utils/cache.py +++ b/zipline/utils/cache.py @@ -290,11 +290,7 @@ class working_file(object): def _commit(self): """Sync the temporary file to the final path. """ - self._tmpfile.close() - move(self._name, self._final_path) - - def __getattr__(self, attr): - return getattr(self._tmpfile, attr) + move(self.path, self._final_path) def __enter__(self): self._tmpfile.__enter__()