Repository: rucio/rucio

Issue:
Possibility to download DIDs by giving metadata filter
Motivation
----------
rucio list-dids allows to get a list of DIDs matching a pattern. It would be good to have an option to download by giving a metadata filter, e.g.:

rucio download --scope mc15_13TeV --filter guid=2e2232aafac8324db452070304f8d745

Modification
------------
Implement a --filter option in rucio download



Discussion:

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,11 @@
 # - 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
+#
+# PY3K COMPATIBLE
 
+from __future__ import division
 
 import copy
 import logging
@@ -28,6 +32,7 @@
 import signal
 import time
 
+from builtins import round
 try:
     from Queue import Queue, Empty, deque
 except ImportError:
@@ -262,7 +267,7 @@ def download_pfns(self, items, num_threads=2, trace_custom_fields={}):
 
         return self._check_output(output_items)
 
-    def download_dids(self, items, num_threads=2, trace_custom_fields={}):
+    def download_dids(self, items, num_threads=2, trace_custom_fields={}, filters={}):
         """
         Download items with given DIDs. This function can also download datasets and wildcarded DIDs.
 
@@ -277,7 +282,8 @@ def download_dids(self, items, num_threads=2, trace_custom_fields={}):
             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
+        :param trace_custom_fields: Custom key value pairs to send with the traces.
+        :param filters: dictionary containing filter options.
 
         :returns: a list of dictionaries with an entry for each file, containing the input options, the did, and the clientState
 
@@ -288,7 +294,7 @@ def download_dids(self, items, num_threads=2, trace_custom_fields={}):
         """
         trace_custom_fields['uuid'] = generate_uuid()
 
-        input_items = self._prepare_items_for_download(items)
+        input_items = self._prepare_items_for_download(items, filters)
 
         num_files_in = len(input_items)
         output_items = self._download_multithreaded(input_items, num_threads, trace_custom_fields)
@@ -569,7 +575,7 @@ def _download_item(self, item, trace, log_prefix=''):
             logger.info('%sFile %s successfully downloaded in %s seconds' % (log_prefix, did_str, duration))
         return item
 
-    def download_aria2c(self, items, trace_custom_fields={}):
+    def download_aria2c(self, items, trace_custom_fields={}, filters={}):
         """
         Uses aria2c to download the items with given DIDs. This function can also download datasets and wildcarded DIDs.
         It only can download files that are available via https/davs.
@@ -583,6 +589,7 @@ def download_aria2c(self, items, trace_custom_fields={}):
             nrandom             - Optional: if the DID addresses a dataset, nrandom files will be randomly choosen for download from the dataset
             ignore_checksum     - Optional: If true, skips the checksum validation between the downloaded file and the rucio catalouge. (Default: False)
         :param trace_custom_fields: Custom key value pairs to send with the traces
+        :param filters: dictionary containing filter options
 
         :returns: a list of dictionaries with an entry for each file, containing the input options, the did, and the clientState
 
@@ -599,7 +606,7 @@ def download_aria2c(self, items, trace_custom_fields={}):
 
         for item in items:
             item['force_scheme'] = ['https', 'davs']
-        input_items = self._prepare_items_for_download(items)
+        input_items = self._prepare_items_for_download(items, filters)
 
         try:
             output_items = self._download_items_aria2c(input_items, aria_rpc, rpc_auth, trace_custom_fields)
@@ -835,13 +842,14 @@ def _download_items_aria2c(self, items, aria_rpc, rpc_auth, trace_custom_fields=
 
         return items
 
-    def _prepare_items_for_download(self, items):
+    def _prepare_items_for_download(self, items, filters={}):
         """
         Resolves wildcarded DIDs, get DID details (e.g. type), and collects
         the available replicas for each DID
         (This function is meant to be used as class internal only)
 
         :param items: list of dictionaries containing the items to prepare
+        :param filters: dictionary containing filters
 
         :returns: list of dictionaries, one dict for each file to download
 
@@ -854,7 +862,7 @@ def _prepare_items_for_download(self, items):
         # resolve input: extend rse expression, resolve wildcards, get did type
         for item in items:
             did_str = item.get('did')
-            if not did_str:
+            if not did_str and not filters:
                 raise InputValidationError('The key did is mandatory')
 
             logger.debug('Processing item %s' % did_str)
@@ -867,26 +875,38 @@ def _prepare_items_for_download(self, items):
                 new_item['rse'] = 'istape=False' if not rse else '(%s)&istape=False' % rse
                 logger.debug('RSE-Expression: %s' % new_item['rse'])
 
-            # resolve any wildcards in the input dids
-            did_scope, did_name = self._split_did_str(did_str)
-            logger.debug('Splitted DID: %s:%s' % (did_scope, did_name))
-            new_item['scope'] = did_scope
-            if '*' in did_name:
-                logger.debug('Resolving wildcarded DID %s' % did_str)
-                for dids in self.client.list_dids(did_scope, filters={'name': did_name}, type='all', long=True):
-                    logger.debug('%s - %s:%s' % (dids['did_type'], did_scope, dids['name']))
+            if not did_str:
+                logger.debug('Resolving DIDs by using filter options.')
+                scope = filters['scope']
+                del filters['scope']
+                for dids in self.client.list_dids(scope, filters=filters, type='all', long=True):
+                    logger.debug('%s - %s:%s' % (dids['did_type'], scope, dids['name']))
+                    new_item['scope'] = scope
                     new_item['type'] = dids['did_type'].upper()
                     new_item['name'] = dids['name']
-                    new_item['did'] = '%s:%s' % (did_scope, dids['name'])
+                    new_item['did'] = '%s:%s' % (scope, dids['name'])
                     resolved_items.append(copy.deepcopy(new_item))
             else:
-                new_item['type'] = self.client.get_did(did_scope, did_name)['type'].upper()
-                new_item['name'] = did_name
-                resolved_items.append(new_item)
+                # resolve any wildcards in the input dids
+                did_scope, did_name = self._split_did_str(did_str)
+                logger.debug('Splitted DID: %s:%s' % (did_scope, did_name))
+                new_item['scope'] = did_scope
+                if '*' in did_name:
+                    logger.debug('Resolving wildcarded DID %s' % did_str)
+                    filters['name'] = did_name
+                    for dids in self.client.list_dids(did_scope, filters=filters, type='all', long=True):
+                        logger.debug('%s - %s:%s' % (dids['did_type'], did_scope, dids['name']))
+                        new_item['type'] = dids['did_type'].upper()
+                        new_item['name'] = dids['name']
+                        new_item['did'] = '%s:%s' % (did_scope, dids['name'])
+                        resolved_items.append(copy.deepcopy(new_item))
+                else:
+                    new_item['type'] = self.client.get_did(did_scope, did_name)['type'].upper()
+                    new_item['name'] = did_name
+                    resolved_items.append(new_item)
 
         # this list will have one dict for each file to download
         file_items = []
-
         # get replicas for every file of the given dids
         logger.debug('%d DIDs after processing input' % len(resolved_items))
         for item in resolved_items:
@@ -948,7 +968,6 @@ def _prepare_items_for_download(self, items):
                 file_item['temp_file_path'] = dest_file_path + '.part'
 
                 file_items.append(file_item)
-
         return file_items
 
     def _split_did_str(self, did_str):
diff --git a/lib/rucio/common/utils.py b/lib/rucio/common/utils.py
--- a/lib/rucio/common/utils.py
+++ b/lib/rucio/common/utils.py
@@ -75,7 +75,7 @@
     import urllib.parse as urlparse
 
 from rucio.common.config import config_get
-from rucio.common.exception import MissingModuleException
+from rucio.common.exception import MissingModuleException, InvalidType
 
 # Extra modules: Only imported if available
 EXTRA_MODULES = {'web': False,
@@ -759,3 +759,58 @@ def get_bytes_value_from_string(input_string):
         return value
     else:
         return False
+
+
+def parse_did_filter_from_string(input_string):
+    """
+    Parse DID filter options in format 'length<3,type=all' from string.
+
+    :param input_string: String containing the filter options.
+    :return: filter dictionary and type as string.
+    """
+    filters = {}
+    type = 'collection'
+    if input_string:
+        filter_options = input_string.replace(' ', '').split(',')
+        for option in filter_options:
+            value = None
+            key = None
+
+            if '>=' in option:
+                key, value = option.split('>=')
+                if key == 'length':
+                    key = 'length.gte'
+            elif '>' in option:
+                key, value = option.split('>')
+                if key == 'length':
+                    key = 'length.gt'
+            elif '<=' in option:
+                key, value = option.split('<=')
+                if key == 'length':
+                    key = 'length.lte'
+            elif '<' in option:
+                key, value = option.split('<')
+                if key == 'length':
+                    key = 'length.lt'
+            elif '=' in option:
+                key, value = option.split('=')
+
+            if key == 'type':
+                if value.upper() in ['ALL', 'COLLECTION', 'CONTAINER', 'DATASET', 'FILE']:
+                    type = value.lower()
+                else:
+                    raise InvalidType('{0} is not a valid type. Valid types are {1}'.format(value, ['ALL', 'COLLECTION', 'CONTAINER', 'DATASET', 'FILE']))
+            elif key in ('length.gt', 'length.lt', 'length.gte', 'length.lte', 'length'):
+                try:
+                    value = int(value)
+                    filters[key] = value
+                except ValueError:
+                    raise ValueError('Length has to be an integer value.')
+                filters[key] = value
+            else:
+                if value.lower() == 'true':
+                    value = '1'
+                elif value.lower() == 'false':
+                    value = '0'
+                filters[key] = value
+    return filters, type


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
@@ -32,6 +32,7 @@
 import nose.tools
 import re
 
+from rucio.db.sqla import session, models
 from rucio.client.accountlimitclient import AccountLimitClient
 from rucio.client.didclient import DIDClient
 from rucio.client.replicaclient import ReplicaClient
@@ -412,6 +413,25 @@ def test_download_file(self):
         exitcode, out, err = execute(cmd)
         print(out, err)
         nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
+        tmp_file1 = file_generator()
+        # add files
+        cmd = 'rucio upload --rse {0} --scope {1} {2}'.format(self.def_rse, self.user, tmp_file1)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        # download files
+        cmd = 'rucio -v download --dir /tmp {0}:{1}'.format(self.user, tmp_file1[5:-2] + '*')  # triming '/tmp/' from filename
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        # search for the files with ls
+        cmd = 'ls /tmp/'    # search in /tmp/
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
         try:
             for i in listdir('data13_hip'):
                 unlink('data13_hip/%s' % i)
@@ -419,6 +439,118 @@ def test_download_file(self):
         except Exception:
             pass
 
+    def test_download_filter(self):
+        """CLIENT(USER): Rucio download with filter options"""
+        # Use filter option to download file with wildcarded name
+        tmp_file1 = file_generator()
+        uuid = generate_uuid()
+        cmd = 'rucio upload --rse {0} --scope {1} --guid {2} {3}'.format(self.def_rse, self.user, uuid, tmp_file1)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        remove(tmp_file1)
+        wrong_guid = generate_uuid()
+        cmd = 'rucio -v download --dir /tmp {0}:{1} --filter guid={2}'.format(self.user, '*', wrong_guid)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(self.user)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_equal(re.search(tmp_file1[5:], out), None)
+        cmd = 'rucio -v download --dir /tmp {0}:{1} --filter guid={2}'.format(self.user, '*', uuid)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(self.user)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
+        # Only use filter option to download file
+        tmp_file1 = file_generator()
+        uuid = generate_uuid()
+        cmd = 'rucio upload --rse {0} --scope {1} --guid {2} {3}'.format(self.def_rse, self.user, uuid, tmp_file1)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        remove(tmp_file1)
+        wrong_guid = generate_uuid()
+        cmd = 'rucio -v download --dir /tmp --scope {0} --filter guid={1}'.format(self.user, wrong_guid)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(self.user)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_equal(re.search(tmp_file1[5:], out), None)
+        cmd = 'rucio -v download --dir /tmp --scope {0} --filter guid={1}'.format(self.user, uuid)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(self.user)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
+        # Only use filter option to download dataset
+        tmp_file1 = file_generator()
+        dataset_name = 'dataset_%s' % generate_uuid()
+        cmd = 'rucio upload --rse {0} --scope {1} {2} {1}:{3}'.format(self.def_rse, self.user, tmp_file1, dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        remove(tmp_file1)
+        db_session = session.get_session()
+        db_session.query(models.DataIdentifier).filter_by(scope=self.user, name=dataset_name).one().length = 15
+        db_session.commit()
+        cmd = 'rucio download --dir /tmp --scope {0} --filter length=100'.format(self.user)
+        exitcode, out, err = execute(cmd)
+        cmd = 'ls /tmp/{0}'.format(dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_equal(re.search(tmp_file1[5:], out), None)
+        cmd = 'rucio download --dir /tmp --scope {0} --filter length=15'.format(self.user)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
+        # Use filter option to download dataset with wildcarded name
+        tmp_file1 = file_generator()
+        dataset_name = 'dataset_%s' % generate_uuid()
+        cmd = 'rucio upload --rse {0} --scope {1} {2} {1}:{3}'.format(self.def_rse, self.user, tmp_file1, dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        remove(tmp_file1)
+        db_session = session.get_session()
+        db_session.query(models.DataIdentifier).filter_by(scope=self.user, name=dataset_name).one().length = 1
+        db_session.commit()
+        cmd = 'rucio download --dir /tmp {0}:{1} --filter length=10'.format(self.user, dataset_name[0:-1] + '*')
+        exitcode, out, err = execute(cmd)
+        cmd = 'ls /tmp/{0}'.format(dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_equal(re.search(tmp_file1[5:], out), None)
+        cmd = 'rucio download --dir /tmp {0}:{1} --filter length=1'.format(self.user, dataset_name[0:-1] + '*')
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        cmd = 'ls /tmp/{0}'.format(dataset_name)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        nose.tools.assert_not_equal(re.search(tmp_file1[5:], out), None)
+
     def test_download_succeeds_md5only(self):
         """CLIENT(USER): Rucio download succeeds MD5 only"""
         # user has a file to upload
diff --git a/lib/rucio/tests/test_utils.py b/lib/rucio/tests/test_utils.py
--- a/lib/rucio/tests/test_utils.py
+++ b/lib/rucio/tests/test_utils.py
@@ -17,7 +17,8 @@
 
 from nose.tools import assert_raises, assert_equal, assert_is_instance, assert_is_not_none
 from re import match
-from rucio.common.utils import md5, adler32
+from rucio.common.exception import InvalidType
+from rucio.common.utils import md5, adler32, parse_did_filter_from_string
 
 
 class TestUtils(unittest.TestCase):
@@ -55,3 +56,28 @@ def test_utils_adler32(self):
         with assert_raises(Exception) as e:
             adler32('no_file')
         assert_equal('FATAL - could not get Adler32 checksum of file no_file - [Errno 2] No such file or directory: \'no_file\'', e.exception.message)
+
+    def test_parse_did_filter_string(self):
+        """(COMMON/UTILS): test parsing of did filter string"""
+        test_cases = [{
+            'input': 'type=all,length=3,length>4,length>=6,length<=7,  test=b',
+            'expected_filter': {'length': 3, 'length.gt': 4, 'length.gte': 6, 'length.lte': 7, 'test': 'b'},
+            'expected_type': 'all'
+        }, {
+            'input': 'type=FILE',
+            'expected_filter': {},
+            'expected_type': 'file'
+        }, {
+            'input': '',
+            'expected_filter': {},
+            'expected_type': 'collection'
+        }]
+
+        for test_case in test_cases:
+            filters, type = parse_did_filter_from_string(test_case['input'])
+            assert_equal(test_case['expected_filter'], filters)
+            assert_equal(test_case['expected_type'], type)
+
+        with assert_raises(InvalidType):
+            input = 'type=g'
+            parse_did_filter_from_string(input)