DID Parameter Parsing (#17)

* Implement parameter parsing as in this [ServiceX story](ssl-hep/ServiceX#395)

Individual CommitsL

* Parse a `files` count thing

* Get rid of post-3.6 dependencies

* Search for the parsed DID, not the original parameterized one!

* Fix up flake8 errors

* Test for parsed parameter removal

* Update docs and comments

* Convert so everything is multi-path

* Make sure multi-path works ok

* Fix flake8

* Obey all vs avail get modes

* Get logic of using get and files correct

Author: gordonwatts

README.md

@@ -128,3 +128,13 @@ In the end, all DID finders for ServiceX will run under Kubernetes. ServiceX com
 Note the parameter `request-id`: this marks the log messages with the request id that triggered this DID request. This will enable the system to track all log messages across all containers connected with this particular request id - making debugging a lot easier.
 
 The `start_did_finder` will configure the python root logger properly to dump messages with a request ID in them.
+
+## URI Format
+
+All the incoming DID's are expected to be URI's without the schema. As such, there are two parameters that are currently parsed by the library. The rest are let through and routed to the callback:
+
+* `files` - Number of files to report back to ServiceX. All files from the dataset are found, and then sorted in order. The first n files are then
+    sent back. Default is all files.
+* `get` - If the value is `all` (the default) then all files in the dataset must be returned. If the value is `available`, then only files that are accessible need be returned.
+
+As am example, if the following URI is given to ServiceX, "rucio://dataset_name?files=20&get=available", then the first 20 available files of the dataset will be processed by the rest of servicex.

src/servicex_did_finder_lib/communication.py

@@ -3,14 +3,15 @@
 import json
 import logging
 import time
-from typing import Any, AsyncGenerator, Callable, Dict, Optional
+from typing import Any, AsyncGenerator, Callable, Dict, List, Optional
 import sys
 
 import pika
 from make_it_sync import make_sync
 
 from servicex_did_finder_lib.did_summary import DIDSummary
 from servicex_did_finder_lib.did_logging import initialize_root_logger
+from servicex_did_finder_lib.util_uri import parse_did_uri
 from .servicex_adaptor import ServiceXAdapter
 
 # The type for the callback method to handle DID's, supplied by the user.
@@ -26,17 +27,52 @@
 __logging.addHandler(logging.NullHandler())
 
 
+class _accumulator:
+    'Track or cache files depending on the mode we are operating in'
+    def __init__(self, sx: ServiceXAdapter, sum: DIDSummary, hold_till_end: bool):
+        self._servicex = sx
+        self._summary = sum
+        self._hold_till_end = hold_till_end
+        self._file_cache: List[Dict[str, Any]] = []
+
+    def add(self, file_info: Dict[str, Any]):
+        'Track and inject the file back into the system'
+        if self._hold_till_end:
+            self._file_cache.append(file_info)
+        else:
+            self._summary.add_file(file_info)
+            if self._summary.file_count == 1:
+                self._servicex.post_transform_start()
+            self._servicex.put_file_add(file_info)
+
+    def send_on(self, count):
+        'Send the accumulated files'
+        if self._hold_till_end:
+            self._hold_till_end = False
+            files = sorted(self._file_cache, key=lambda x: x['paths'])
+            for file_info in files[0:count]:
+                self.add(file_info)
+
+
 async def run_file_fetch_loop(did: str, servicex: ServiceXAdapter, info: Dict[str, Any],
                               user_callback: UserDIDHandler):
-    summary = DIDSummary(did)
     start_time = datetime.now()
-    async for file_info in user_callback(did, info):
 
-        # Track the file, inject back into the system
-        summary.add_file(file_info)
-        if summary.file_count == 1:
-            servicex.post_transform_start()
-        servicex.put_file_add(file_info)
+    summary = DIDSummary(did)
+    did_info = parse_did_uri(did)
+    hold_till_end = did_info.file_count != -1
+    acc = _accumulator(servicex, summary, hold_till_end)
+
+    try:
+        async for file_info in user_callback(did_info.did, info):
+            acc.add(file_info)
+
+    except Exception:
+        if did_info.get_mode == 'all':
+            raise
+
+    # If we've been holding onto any files, we need to send them now.
+    acc.send_on(did_info.file_count)
 
     # Simple error checking and reporting
     if summary.file_count == 0:
@@ -54,7 +90,7 @@ async def run_file_fetch_loop(did: str, servicex: ServiceXAdapter, info: Dict[st
         }
     )
 
-    servicex.post_status_update(f'Completed load of file in {elapsed_time} seconds')
+    servicex.post_status_update(f'Completed load of files in {elapsed_time} seconds')
 
 
 def rabbit_mq_callback(user_callback: UserDIDHandler, channel, method, properties, body,
@@ -109,7 +145,7 @@ def rabbit_mq_callback(user_callback: UserDIDHandler, channel, method, propertie
 def init_rabbit_mq(user_callback: UserDIDHandler,
                    rabbitmq_url: str, queue_name: str, retries: int,
                    retry_interval: float,
-                   file_prefix: str = None):
+                   file_prefix: str = None):  # type: ignore
     rabbitmq = None
     retry_count = 0
 

src/servicex_did_finder_lib/util_uri.py

@@ -0,0 +1,60 @@
+from typing import Dict, List
+import urllib
+
+
+class ParsedDIDInfo:
+    def __init__(self, did: str, get_mode: str, file_count: int):
+        self.did = did
+        self.get_mode = get_mode
+        self.file_count = file_count
+
+    # The did to pass into the library
+    did: str
+
+    # Mode to get the files (default 'all')
+    get_mode: str
+
+    # Number of files to fetch (default '-1')
+    file_count: int
+
+
+def parse_did_uri(uri: str) -> ParsedDIDInfo:
+    '''Parse the uri that is given to us from ServiceX, pulling out
+    the components we care about, and keeping the DID that needs to
+    be passed down.
+
+    URI arguments parsed:
+
+    * `files` - Number of files to fetch (default is all)
+    * `get` - Mode to get the files (default is 'all'). Only "available" is also supported.
+
+    Args:
+        uri (str): DID from ServiceX
+
+    Returns:
+        ParsedDIDInfo: The URI parsed into parts
+    '''
+    info = urllib.parse.urlparse(uri)  # type: ignore
+
+    params = urllib.parse.parse_qs(info.query)  # type: ignore
+    get_string = 'all' if 'get' not in params else params['get'][-1]
+    file_count = -1 if 'files' not in params else int(params['files'][0])
+
+    if get_string not in ['all', 'available']:
+        raise ValueError('Bad value for "get" string in DID - must be "all" or "available", not '
+                         f'"{get_string}"')
+
+    for k in ['get', 'files']:
+        if k in params:
+            del params[k]
+
+    def unwind_params(ps: Dict[str, List[str]]):
+        for k, values in ps.items():
+            for v in values:
+                yield k, v
+
+    new_query = "&".join(f'{k}={v}' for k, v in unwind_params(params))
+    if len(new_query) > 0:
+        new_query = "?" + new_query
+
+    return ParsedDIDInfo(info.path + new_query, get_string, file_count)

tests/servicex_did_finder_lib/test_communication.py

@@ -116,6 +116,39 @@ async def my_callback(did_name: str, info: Dict[str, Any]):
     # Make sure callback was called
     assert seen_name == 'hi-there'
 
+    # Make sure the file was sent along, along with the completion
+    SXAdaptor.put_file_add.assert_called_once_with({
+        'paths': ['fork/it/over'],
+        'adler32': 'no clue',
+        'file_size': 22323,
+        'file_events': 0
+        })
+    SXAdaptor.put_fileset_complete.assert_called_once()
+
+
+def test_one_file_call_with_param(rabbitmq, SXAdaptor):
+    'Test a working, simple, one file call with parameter'
+
+    seen_name = None
+
+    async def my_callback(did_name: str, info: Dict[str, Any]):
+        nonlocal seen_name
+        seen_name = did_name
+        yield {
+            'paths': ["fork/it/over"],
+            'adler32': 'no clue',
+            'file_size': 22323,
+            'file_events': 0,
+        }
+
+    init_rabbit_mq(my_callback, 'http://myrabbit.com', 'test_queue_name', retries=12,
+                   retry_interval=10)
+
+    rabbitmq.send_did_request('hithere?files=10')
+
+    # Make sure callback was called
+    assert seen_name == 'hithere'
+
     # Make sure the file was sent along, along with the completion
     SXAdaptor.put_file_add.assert_called_once()
     SXAdaptor.put_fileset_complete.assert_called_once()
@@ -142,6 +175,81 @@ async def my_callback(did_name: str, info: Dict[str, Any]) \
     SXAdaptor.post_status_update.assert_any_call(ANY, severity='fatal')
 
 
+def test_failed_file_after_good(rabbitmq, SXAdaptor):
+    'Test a callback that fails before any files are sent'
+
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
+        yield {
+            'paths': ["fork/it/over"],
+            'adler32': 'no clue',
+            'file_size': 22323,
+            'file_events': 0,
+        }
+        raise Exception('that did not work')
+
+    init_rabbit_mq(my_callback, 'http://myrabbit.com', 'test_queue_name', retries=12,
+                   retry_interval=10)
+    rabbitmq.send_did_request('hi-there')
+
+    # Make sure the file was sent along, along with the completion
+    SXAdaptor.put_file_add.assert_called_once()
+    SXAdaptor.put_fileset_complete.assert_not_called()
+    SXAdaptor.post_status_update.assert_any_call(ANY, severity='fatal')
+
+
+def test_failed_file_after_good_with_avail(rabbitmq, SXAdaptor):
+    'Test a callback that fails before any files are sent'
+
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
+        yield {
+            'paths': ["fork/it/over"],
+            'adler32': 'no clue',
+            'file_size': 22323,
+            'file_events': 0,
+        }
+        raise Exception('that did not work')
+
+    init_rabbit_mq(my_callback, 'http://myrabbit.com', 'test_queue_name', retries=12,
+                   retry_interval=10)
+    rabbitmq.send_did_request('hi-there?get=available')
+
+    # Make sure the file was sent along, along with the completion
+    SXAdaptor.put_file_add.assert_called_once()
+    SXAdaptor.put_fileset_complete.assert_called_once()
+    SXAdaptor.post_status_update.assert_any_call("Completed load of files in 0 seconds")
+
+
+def test_failed_file_after_good_with_avail_limited_number(rabbitmq, SXAdaptor):
+    'Files are sent back, only available allowed, but want a certian number'
+
+    async def my_callback(did_name: str, info: Dict[str, Any]) \
+            -> AsyncGenerator[Dict[str, Any], None]:
+        yield {
+            'paths': ["fork/it/over1"],
+            'adler32': 'no clue',
+            'file_size': 22323,
+            'file_events': 0,
+        }
+        yield {
+            'paths': ["fork/it/over2"],
+            'adler32': 'no clue',
+            'file_size': 22323,
+            'file_events': 0,
+        }
+        raise Exception('that did not work')
+
+    init_rabbit_mq(my_callback, 'http://myrabbit.com', 'test_queue_name', retries=12,
+                   retry_interval=10)
+    rabbitmq.send_did_request('hi-there?get=available&files=1')
+
+    # Make sure the file was sent along, along with the completion
+    SXAdaptor.put_file_add.assert_called_once()
+    SXAdaptor.put_fileset_complete.assert_called_once()
+    SXAdaptor.post_status_update.assert_any_call("Completed load of files in 0 seconds")
+
+
 def test_no_files_returned(rabbitmq, SXAdaptor):
     'Test a callback that fails before any files are sent'
 
@@ -258,6 +366,103 @@ async def my_user_callback(did, info):
     assert SXAdaptor.post_status_update.called_once()
 
 
+@pytest.mark.asyncio
+async def test_run_file_fetch_one(SXAdaptor, mocker):
+    async def my_user_callback(did, info):
+        return_values = [
+            {
+                'paths': ['/tmp/foo'],
+                'adler32': '13e4f',
+                'file_size': 1024,
+                'file_events': 128
+            },
+            {
+                'paths': ['/tmp/bar'],
+                'adler32': 'f33d',
+                'file_size': 2046,
+                'file_events': 64
+            }
+        ]
+        for v in return_values:
+            yield v
+
+    await run_file_fetch_loop("123-456?files=1", SXAdaptor, {}, my_user_callback)
+    SXAdaptor.post_transform_start.assert_called_once()
+
+    assert SXAdaptor.put_file_add.call_count == 1
+    assert SXAdaptor.put_file_add.call_args_list[0][0][0]['paths'] == ['/tmp/bar']
+
+    SXAdaptor.put_fileset_complete.assert_called_once
+    assert SXAdaptor.put_fileset_complete.call_args[0][0]['files'] == 1
+    assert SXAdaptor.post_status_update.called_once()
+
+
+@pytest.mark.asyncio
+async def test_run_file_fetch_one_reverse(SXAdaptor, mocker):
+    'The files should be sorted so they return the same'
+    async def my_user_callback(did, info):
+        return_values = [
+            {
+                'paths': ['/tmp/bar'],
+                'adler32': 'f33d',
+                'file_size': 2046,
+                'file_events': 64
+            },
+            {
+                'paths': ['/tmp/foo'],
+                'adler32': '13e4f',
+                'file_size': 1024,
+                'file_events': 128
+            },
+        ]
+        for v in return_values:
+            yield v
+
+    await run_file_fetch_loop("123-456?files=1", SXAdaptor, {}, my_user_callback)
+    SXAdaptor.post_transform_start.assert_called_once()
+
+    assert SXAdaptor.put_file_add.call_count == 1
+    assert SXAdaptor.put_file_add.call_args_list[0][0][0]['paths'][0] == '/tmp/bar'
+
+    SXAdaptor.put_fileset_complete.assert_called_once
+    assert SXAdaptor.put_fileset_complete.call_args[0][0]['files'] == 1
+    assert SXAdaptor.post_status_update.called_once()
+
+
+@pytest.mark.asyncio
+async def test_run_file_fetch_one_multi(SXAdaptor, mocker):
+    async def my_user_callback(did, info):
+        return_values = [
+            {
+                'paths': ['/tmp/foo', 'others:/tmp/foo'],
+                'adler32': '13e4f',
+                'file_size': 1024,
+                'file_events': 128
+            },
+            {
+                'paths': ['/tmp/bar', 'others:/tmp/bar'],
+                'adler32': 'f33d',
+                'file_size': 2046,
+                'file_events': 64
+            }
+        ]
+        for v in return_values:
+            yield v
+
+    await run_file_fetch_loop("123-456?files=1", SXAdaptor, {}, my_user_callback)
+    SXAdaptor.post_transform_start.assert_called_once()
+
+    assert SXAdaptor.put_file_add.call_count == 1
+    assert SXAdaptor.put_file_add.call_args_list[0][0][0]['paths'] == [
+        '/tmp/bar',
+        'others:/tmp/bar'
+        ]
+
+    SXAdaptor.put_fileset_complete.assert_called_once
+    assert SXAdaptor.put_fileset_complete.call_args[0][0]['files'] == 1
+    assert SXAdaptor.post_status_update.called_once()
+
+
 @pytest.mark.asyncio
 async def test_run_file_fetch_loop_bad_did(SXAdaptor, mocker):
     async def my_user_callback(did, info):