From b1135e7edc13baee0f11c2ce2a2056e0933e53a4 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Thu, 5 Apr 2018 23:27:02 +0300 Subject: [PATCH] dvc: allow specifying '--remote' param for push/pull/fetch/status Signed-off-by: Ruslan Kuprieiev --- dvc/cli.py | 14 ++++++++++- dvc/cloud/data_cloud.py | 51 ++++++++++++++++++++++++---------------- dvc/command/data_sync.py | 6 ++--- dvc/command/status.py | 4 +--- dvc/config.py | 2 +- dvc/logger.py | 3 ++- dvc/project.py | 20 ++++++++-------- 7 files changed, 61 insertions(+), 39 deletions(-) diff --git a/dvc/cli.py b/dvc/cli.py index 8457b3463..61953d265 100644 --- a/dvc/cli.py +++ b/dvc/cli.py @@ -144,6 +144,9 @@ def parse_args(argv=None): 'pull', parents=[parent_cache_parser], help='Pull data files from the cloud') + pull_parser.add_argument('-r', + '--remote', + help='Remote repository to pull from') pull_parser.set_defaults(func=CmdDataPull) # Push @@ -151,6 +154,9 @@ def parse_args(argv=None): 'push', parents=[parent_cache_parser], help='Push data files to the cloud') + push_parser.add_argument('-r', + '--remote', + help='Remote repository to push to') push_parser.set_defaults(func=CmdDataPush) # Fetch @@ -158,6 +164,9 @@ def parse_args(argv=None): 'fetch', parents=[parent_cache_parser], help='Fetch data files from the cloud') + fetch_parser.add_argument('-r', + '--remote', + help='Remote repository to fetch from') fetch_parser.set_defaults(func=CmdDataFetch) # Status @@ -169,7 +178,10 @@ def parse_args(argv=None): '--cloud', action='store_true', default=False, - help='Show status of a local cache compared to a cloud') + help='Show status of a local cache compared to a remote repository') + status_parser.add_argument('-r', + '--remote', + help='Remote repository to compare local cache to') status_parser.set_defaults(func=CmdDataStatus) # Repro diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index a50bbd517..ded861a63 100644 --- a/dvc/cloud/data_cloud.py +++ b/dvc/cloud/data_cloud.py @@ -36,6 +36,7 @@ class DataCloud(object): } def __init__(self, cache, config): + self._cache = cache self._config = config remote = self._config[Config.SECTION_CORE].get(Config.SECTION_CORE_REMOTE, '') @@ -43,11 +44,14 @@ class DataCloud(object): if config[Config.SECTION_CORE].get(Config.SECTION_CORE_CLOUD, None): # backward compatibility Logger.warn('Using obsoleted config format. Consider updating.') - self.__init__compat(cache) + self._cloud = self.__init__compat() else: self._cloud = None return + self._cloud = self._init_remote(remote) + + def _init_remote(self, remote): section = Config.SECTION_REMOTE_FMT.format(remote) cloud_config = self._config.get(section, None) if not cloud_config: @@ -59,9 +63,9 @@ class DataCloud(object): if not cloud_type: raise ConfigError("Unsupported scheme '{}' in '{}'".format(scheme, url)) - self._init_cloud(cache, cloud_config, cloud_type) + return self._init_cloud(self._cache, cloud_config, cloud_type) - def __init__compat(self, cache): + def __init__compat(self): cloud_type = self._config[Config.SECTION_CORE].get(Config.SECTION_CORE_CLOUD, '').strip().lower() if cloud_type == '': self._cloud = None @@ -73,15 +77,16 @@ class DataCloud(object): if not cloud_config: raise ConfigError('Can\'t find cloud section \'[%s]\' in config' % cloud_type) - self._init_cloud(cache, cloud_config, cloud_type) + return self._init_cloud(self._cache, cloud_config, cloud_type) def _init_cloud(self, cache, cloud_config, cloud_type): cloud_settings = self.get_cloud_settings(cache, self._config, cloud_config) - self._cloud = self.CLOUD_MAP[cloud_type](cloud_settings) - self._cloud.sanity_check() + cloud = self.CLOUD_MAP[cloud_type](cloud_settings) + cloud.sanity_check() + return cloud @staticmethod def get_cloud_settings(cache, config, cloud_config): @@ -95,48 +100,54 @@ class DataCloud(object): cloud_settings = CloudSettings(cache, global_storage_path, cloud_config) return cloud_settings - def _collect(self, targets, jobs, local): + def _collect(self, cloud, targets, jobs, local): collected = set() pool = ThreadPool(processes=jobs) args = zip(targets, [local]*len(targets)) - ret = pool.map(self._cloud.collect, args) + ret = pool.map(cloud.collect, args) for r in ret: collected |= set(r) return collected - def _map_targets(self, func, targets, jobs, collect_local=False, collect_cloud=False): + def _map_targets(self, func, targets, jobs, collect_local=False, collect_cloud=False, remote=None): """ Process targets as data items in parallel. """ - if not self._cloud: + + if not remote: + cloud = self._cloud + else: + cloud = self._init_remote(remote) + + if not cloud: return - self._cloud.connect() + cloud.connect() collected = set() if collect_local: - collected |= self._collect(targets, jobs, True) + collected |= self._collect(cloud, targets, jobs, True) if collect_cloud: - collected |= self._collect(targets, jobs, False) + collected |= self._collect(cloud, targets, jobs, False) - return map_progress(func, list(collected), jobs) + return map_progress(getattr(cloud, func), list(collected), jobs) - def push(self, targets, jobs=1): + def push(self, targets, jobs=1, remote=None): """ Push data items in a cloud-agnostic way. """ - return self._map_targets(self._cloud.push, targets, jobs, collect_local=True) + return self._map_targets('push', targets, jobs, collect_local=True, remote=remote) - def pull(self, targets, jobs=1): + def pull(self, targets, jobs=1, remote=None): """ Pull data items in a cloud-agnostic way. """ - return self._map_targets(self._cloud.pull, targets, jobs, collect_cloud=True) + return self._map_targets('pull', targets, jobs, collect_cloud=True, remote=remote) - def status(self, targets, jobs=1): + def status(self, targets, jobs=1, remote=None): """ Check status of data items in a cloud-agnostic way. """ - return self._map_targets(self._cloud.status, targets, jobs, True, True) + return self._map_targets('status', targets, jobs, True, True, remote=remote) diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py index f174c3b63..e03f8eb5d 100644 --- a/dvc/command/data_sync.py +++ b/dvc/command/data_sync.py @@ -23,7 +23,7 @@ class CmdDataBase(CmdBase): class CmdDataPull(CmdDataBase): def do_run(self, target=None): try: - self.project.pull(target=target, jobs=self.args.jobs) + self.project.pull(target=target, jobs=self.args.jobs, remote=self.args.remote) except Exception as exc: self.project.logger.error('Failed to pull data from the cloud', exc) return 1 @@ -33,7 +33,7 @@ class CmdDataPull(CmdDataBase): class CmdDataPush(CmdDataBase): def do_run(self, target=None): try: - self.project.push(target=target, jobs=self.args.jobs) + self.project.push(target=target, jobs=self.args.jobs, remote=self.args.remote) except Exception as exc: self.project.logger.error('Failed to push data to the cloud', exc) return 1 @@ -43,7 +43,7 @@ class CmdDataPush(CmdDataBase): class CmdDataFetch(CmdDataBase): def do_run(self, target=None): try: - self.project.fetch(target=target, jobs=self.args.jobs) + self.project.fetch(target=target, jobs=self.args.jobs, remote=self.args.remote) except Exception as exc: self.project.logger.error('Failed to fetch data from the cloud', exc) return 1 diff --git a/dvc/command/status.py b/dvc/command/status.py index dd3c63c49..712e65181 100644 --- a/dvc/command/status.py +++ b/dvc/command/status.py @@ -1,7 +1,5 @@ import os -import dvc.cloud.base as cloud - from dvc.command.data_sync import CmdDataBase @@ -28,7 +26,7 @@ class CmdDataStatus(CmdDataBase): def do_run(self, target=None): indent = 1 if self.args.cloud else 0 try: - status = self.project.status(target=target, jobs=self.args.jobs, cloud=self.args.cloud) + status = self.project.status(target=target, jobs=self.args.jobs, cloud=self.args.cloud, remote=self.args.remote) self._show(status, indent) except Exception as exc: self.project.logger.error('Failed to obtain data status', exc) diff --git a/dvc/config.py b/dvc/config.py index d82b4041a..8b9b222bb 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -11,7 +11,7 @@ from dvc.exceptions import DvcException class ConfigError(DvcException): """ DVC config exception """ def __init__(self, ex=None): - DvcException.__init__(self, 'Config file error', ex) + super(ConfigError, self).__init__('Config file error', ex) class Config(object): diff --git a/dvc/logger.py b/dvc/logger.py index d4a7bf990..40bcad2f4 100644 --- a/dvc/logger.py +++ b/dvc/logger.py @@ -76,7 +76,8 @@ class Logger(object): str_exc = ': {}'.format(str(exc)) if exc else "" if exc and hasattr(exc, 'cause') and exc.cause: - cause_str_exc, cause_str_tb = Logger.parse_exc(exc.cause, exc.cause_tb) + cause_tb = exc.cause_tb if hasattr(exc, 'cause_tb') else None + cause_str_exc, cause_str_tb = Logger.parse_exc(exc.cause, cause_tb) str_tb = cause_str_tb str_exc = '{}{}'.format(str_exc, cause_str_exc) diff --git a/dvc/project.py b/dvc/project.py index 27f5287d1..e8ecdbb8e 100644 --- a/dvc/project.py +++ b/dvc/project.py @@ -199,14 +199,14 @@ class Project(object): os.unlink(cache) self.logger.info(u'\'{}\' was removed'.format(self.to_dvc_path(cache))) - def push(self, target=None, jobs=1): - return self.cloud.push(self._used_cache(target), jobs) + def push(self, target=None, jobs=1, remote=None): + return self.cloud.push(self._used_cache(target), jobs, remote=remote) - def fetch(self, target=None, jobs=1): - return self.cloud.pull(self._used_cache(target), jobs) + def fetch(self, target=None, jobs=1, remote=None): + return self.cloud.pull(self._used_cache(target), jobs, remote=remote) - def pull(self, target=None, jobs=1): - ret = self.fetch(target, jobs) + def pull(self, target=None, jobs=1, remote=None): + ret = self.fetch(target, jobs, remote=remote) self.checkout() return ret @@ -223,9 +223,9 @@ class Project(object): return status - def _cloud_status(self, target=None, jobs=1): + def _cloud_status(self, target=None, jobs=1, remote=None): status = {} - for target, ret in self.cloud.status(self._used_cache(target), jobs): + for target, ret in self.cloud.status(self._used_cache(target), jobs, remote=remote): if ret == cloud.STATUS_UNKNOWN or ret == cloud.STATUS_OK: continue @@ -241,9 +241,9 @@ class Project(object): return status - def status(self, target=None, jobs=1, cloud=False): + def status(self, target=None, jobs=1, cloud=False, remote=None): if cloud: - return self._cloud_status(target, jobs) + return self._cloud_status(target, jobs, remote=remote) return self._local_status(target) def graph(self):