Repository: rucio/rucio

Issue:
rucio download should be able to use locally cached metalink
Motivation
----------

Instead of calling list_replicas use locally cached metalink file.

Ref: @PalNilsson


Discussion:
`rucio download --metalink=my_replicas.meta4 scope:name`

Resolving patch:
diff --git a/lib/rucio/client/downloadclient.py b/lib/rucio/client/downloadclient.py
--- a/lib/rucio/client/downloadclient.py
+++ b/lib/rucio/client/downloadclient.py
@@ -18,7 +18,7 @@
 # - Joaquin Bogado <jbogado@linti.unlp.edu.ar>, 2018
 # - Nicolo Magini <nicolo.magini@cern.ch>, 2018
 # - Tobias Wegner <tobias.wegner@cern.ch>, 2018
-# - Hannes Hansen <hannes.jakob.hansen@cern.ch>, 2018
+# - Hannes Hansen <hannes.jakob.hansen@cern.ch>, 2018-2019
 #
 # PY3K COMPATIBLE
 
@@ -42,7 +42,7 @@
 from rucio.client.client import Client
 from rucio.common.exception import (InputValidationError, NoFilesDownloaded, ServiceUnavailable,
                                     NotAllFilesDownloaded, RSENotFound, RucioException, SourceNotFound)
-from rucio.common.utils import adler32, md5, detect_client_location, generate_uuid, parse_replicas_from_string, send_trace, sizefmt, execute
+from rucio.common.utils import adler32, md5, detect_client_location, generate_uuid, parse_replicas_from_string, send_trace, sizefmt, execute, parse_replicas_from_file
 from rucio.rse import rsemanager as rsemgr
 from rucio import version
 
@@ -380,7 +380,53 @@ def download_dids(self, items, num_threads=2, trace_custom_fields={}):
         self.logger.debug('num_unmerged_items=%d; num_dids=%d; num_merged_items=%d' % (len(items), len(did_to_options), len(merged_items)))
 
         logger.info('Getting sources of DIDs')
-        input_items = self._prepare_items_for_download(did_to_options, merged_items)
+        # if one item wants to resolve archives we enable it for all items
+        resolve_archives = not all(item.get('no_resolve_archives') for item in merged_items)
+        merged_items_with_sources = self._get_sources(merged_items, resolve_archives=resolve_archives)
+        input_items = self._prepare_items_for_download(did_to_options, merged_items_with_sources, resolve_archives=resolve_archives)
+
+        num_files_in = len(input_items)
+        output_items = self._download_multithreaded(input_items, num_threads, trace_custom_fields)
+        num_files_out = len(output_items)
+
+        if num_files_in != num_files_out:
+            raise RucioException('%d items were in the input queue but only %d are in the output queue' % (num_files_in, num_files_out))
+
+        return self._check_output(output_items)
+
+    def download_from_metalink_file(self, item, metalink_file_path, num_threads=2, trace_custom_fields={}):
+        """
+        Download items using a given metalink file.
+
+        :param item: dictionary describing an item to download. Keys:
+            base_dir            - Optional: base directory where the downloaded files will be stored. (Default: '.')
+            no_subdir           - Optional: If true, files are written directly into base_dir and existing files are overwritten. (Default: False)
+            ignore_checksum     - Optional: If true, skips the checksum validation between the downloaded file and the rucio catalouge. (Default: False)
+            transfer_timeout    - Optional: Timeout time for the download protocols. (Default: None)
+        :param num_threads: Suggestion of number of threads to use for the download. It will be lowered if it's too high.
+        :param trace_custom_fields: Custom key value pairs to send with the traces.
+
+        :returns: a list of dictionaries with an entry for each file, containing the input options, the did, and the clientState
+
+        :raises InputValidationError: if one of the input items is in the wrong format
+        :raises NoFilesDownloaded: if no files could be downloaded
+        :raises NotAllFilesDownloaded: if not all files could be downloaded
+        :raises RucioException: if something unexpected went wrong during the download
+        """
+        logger = self.logger
+
+        logger.info('Getting sources from metalink file')
+        metalinks = parse_replicas_from_file(metalink_file_path)
+
+        trace_custom_fields['uuid'] = generate_uuid()
+
+        did_to_options = {}
+        item.setdefault('destinations', set()).add((item['base_dir'], item['no_subdir']))
+        for metalink in metalinks:
+            did_to_options[metalink['did']] = item
+
+        metalinks = [metalinks]
+        input_items = self._prepare_items_for_download(did_to_options, metalinks)
 
         num_files_in = len(input_items)
         output_items = self._download_multithreaded(input_items, num_threads, trace_custom_fields)
@@ -491,7 +537,6 @@ def _download_item(self, item, trace, log_prefix=''):
         did_scope = item['scope']
         did_name = item['name']
         did_str = '%s:%s' % (did_scope, did_name)
-
         logger.info('%sPreparing download of %s' % (log_prefix, did_str))
 
         trace['scope'] = did_scope
@@ -733,7 +778,6 @@ def download_aria2c(self, items, trace_custom_fields={}, filters={}):
 
         for item in items:
             item['force_scheme'] = ['https', 'davs']
-            item['no_resolve_archives'] = True
 
         logger.info('Processing %d item(s) for input' % len(items))
         download_info = self._resolve_and_merge_input_items(copy.deepcopy(items))
@@ -743,7 +787,8 @@ def download_aria2c(self, items, trace_custom_fields={}, filters={}):
         self.logger.debug('num_unmerged_items=%d; num_dids=%d; num_merged_items=%d' % (len(items), len(did_to_options), len(merged_items)))
 
         logger.info('Getting sources of DIDs')
-        input_items = self._prepare_items_for_download(did_to_options, merged_items)
+        merged_items_with_sources = self._get_sources(merged_items)
+        input_items = self._prepare_items_for_download(did_to_options, merged_items_with_sources, resolve_archives=False)
 
         try:
             output_items = self._download_items_aria2c(input_items, aria_rpc, rpc_auth, trace_custom_fields)
@@ -1084,47 +1129,16 @@ def _resolve_and_merge_input_items(self, items):
                 merged_items.append(item)
         return download_info
 
-    def _prepare_items_for_download(self, did_to_options, merged_items):
+    def _get_sources(self, merged_items, resolve_archives=True):
         """
-        Get sources (PFNs) of the DIDs and optimises the amount of files to download
-        (This function is meant to be used as class internal only)
+        Get sources (PFNs) of the DIDs.
 
-        :param did_to_options: dictionary that maps each input DID to some input options
         :param merged_items: list of dictionaries. Each dictionary describes a bunch of DIDs to download
 
-        :returns: list of dictionaries. Each dictionary describes an element to download
-
-        :raises InputValidationError: if the given input is not valid or incomplete
+        :returns: list of list of dictionaries.
         """
         logger = self.logger
-
-        # if one item wants to resolve archives we enable it for all items
-        resolve_archives = not all(item.get('no_resolve_archives') for item in merged_items)
-        if resolve_archives:
-            # perhaps we'll need an extraction tool so check what is installed
-            self.extraction_tools = [tool for tool in self.extraction_tools if tool.is_useable()]
-            if len(self.extraction_tools) < 1:
-                logger.warning('Archive resolution is enabled but no extraction tool is available. '
-                               'Sources whose protocol doesnt support extraction wont be considered for download.')
-
-        # maps file item IDs (fiid) to the file item object
-        fiid_to_file_item = {}
-
-        # list of all file item objects
-        all_file_items = []
-
-        # cea -> client_extract archives to avoid confusion with archives that dont need explicit extraction
-        # this dict will contain all ids of cea's that definitely will be downloaded
-        cea_id_pure_to_fiids = {}
-
-        # this dict will contain ids of cea's that have higher prioritised non cea sources
-        cea_id_mixed_to_fiids = {}
-
-        all_input_dids = set(did_to_options.keys())
-        all_dest_file_paths = set()
-
-        # get replicas for every file of the given dids
-        logger.debug('num list_replicas calls: %d' % len(merged_items))
+        merged_items_with_sources = []
         for item in merged_items:
             # since we're using metalink we need to explicitly give all schemes
             force_scheme = item.get('force_scheme')
@@ -1142,6 +1156,7 @@ def _prepare_items_for_download(self, did_to_options, merged_items):
 
             # get PFNs of files and datasets
             logger.debug('num DIDs for list_replicas call: %d' % len(item['dids']))
+
             metalink_str = self.client.list_replicas(item['dids'],
                                                      schemes=schemes,
                                                      rse_expression=rse_expression,
@@ -1150,7 +1165,6 @@ def _prepare_items_for_download(self, did_to_options, merged_items):
                                                      resolve_parents=True,
                                                      metalink=True)
             file_items = parse_replicas_from_string(metalink_str)
-
             logger.debug('num resolved files: %s' % len(file_items))
 
             nrandom = item.get('nrandom')
@@ -1158,9 +1172,51 @@ def _prepare_items_for_download(self, did_to_options, merged_items):
                 logger.info('Selecting %d random replicas from DID(s): %s' % (nrandom, item['dids']))
                 random.shuffle(file_items)
                 file_items = file_items[0:nrandom]
+                merged_items_with_sources.append(file_items)
+            else:
+                merged_items_with_sources.append(file_items)
+        return merged_items_with_sources
 
-            all_file_items.extend(file_items)
+    def _prepare_items_for_download(self, did_to_options, merged_items_with_sources, resolve_archives=True):
+        """
+        Optimises the amount of files to download
+        (This function is meant to be used as class internal only)
 
+        :param did_to_options: dictionary that maps each input DID to some input options
+        :param merged_items_with_sources: list of dictionaries. Each dictionary describes a bunch of DIDs to download
+
+        :returns: list of dictionaries. Each dictionary describes an element to download
+
+        :raises InputValidationError: if the given input is not valid or incomplete
+        """
+        logger = self.logger
+        if resolve_archives:
+            # perhaps we'll need an extraction tool so check what is installed
+            self.extraction_tools = [tool for tool in self.extraction_tools if tool.is_useable()]
+            if len(self.extraction_tools) < 1:
+                logger.warning('Archive resolution is enabled but no extraction tool is available. '
+                               'Sources whose protocol doesnt support extraction wont be considered for download.')
+
+        # maps file item IDs (fiid) to the file item object
+        fiid_to_file_item = {}
+
+        # list of all file item objects
+        all_file_items = []
+
+        # cea -> client_extract archives to avoid confusion with archives that dont need explicit extraction
+        # this dict will contain all ids of cea's that definitely will be downloaded
+        cea_id_pure_to_fiids = {}
+
+        # this dict will contain ids of cea's that have higher prioritised non cea sources
+        cea_id_mixed_to_fiids = {}
+
+        all_input_dids = set(did_to_options.keys())
+        all_dest_file_paths = set()
+
+        # get replicas for every file of the given dids
+        logger.debug('num list_replicas calls: %d' % len(merged_items_with_sources))
+        for file_items in merged_items_with_sources:
+            all_file_items.extend(file_items)
             for file_item in file_items:
                 # parent_dids contains all parents, so we take the intersection with the input dids
                 dataset_did_strs = file_item.setdefault('parent_dids', set())


Test patch:
diff --git a/lib/rucio/tests/test_bin_rucio.py b/lib/rucio/tests/test_bin_rucio.py
--- a/lib/rucio/tests/test_bin_rucio.py
+++ b/lib/rucio/tests/test_bin_rucio.py
@@ -23,7 +23,7 @@
 # - Martin Barisits <martin.barisits@cern.ch>, 2015-2016
 # - Frank Berghaus <frank.berghaus@cern.ch>, 2017-2018
 # - Tobias Wegner <twegner@cern.ch>, 2018
-# - Hannes Hansen <hannes.jakob.hansen@cern.ch>, 2018
+# - Hannes Hansen <hannes.jakob.hansen@cern.ch>, 2018-2019
 #
 # PY3K COMPATIBLE
 
@@ -623,6 +623,35 @@ def test_download_filter(self):
         print(out, err)
         nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
 
+    def test_download_metalink_file(self):
+        """CLIENT(USER): Rucio download with metalink file"""
+        metalink_file_path = generate_uuid()
+        scope = self.user
+
+        # Use filter and metalink option
+        cmd = 'rucio download --scope mock --filter size=1 --metalink=test'
+        exitcode, out, err = execute(cmd)
+        nose.tools.assert_in('Arguments filter and metalink cannot be used together', err)
+
+        # Use did and metalink option
+        cmd = 'rucio download --metalink=test mock:test'
+        exitcode, out, err = execute(cmd)
+        nose.tools.assert_in('Arguments dids and metalink cannot be used together', err)
+
+        # Download only with metalink file
+        tmp_file = file_generator()
+        tmp_file_name = tmp_file[5:]
+        cmd = 'rucio upload --rse {0} --scope {1} {2}'.format(self.def_rse, scope, tmp_file)
+        exitcode, out, err = execute(cmd)
+        replica_file = ReplicaClient().list_replicas([{'scope': scope, 'name': tmp_file_name}], metalink=True)
+        with open(metalink_file_path, 'w+') as metalink_file:
+            metalink_file.write(replica_file)
+        cmd = 'rucio download --dir /tmp --metalink {0}'.format(metalink_file_path)
+        exitcode, out, err = execute(cmd)
+        cmd = 'ls /tmp/{0}'.format(scope)
+        exitcode, out, err = execute(cmd)
+        nose.tools.assert_not_equal(re.search(tmp_file_name, out), None)
+
     def test_download_succeeds_md5only(self):
         """CLIENT(USER): Rucio download succeeds MD5 only"""
         # user has a file to upload