diff --git a/atdb_result_gatherer/atdb_result_gatherer/result_gatherer.py b/atdb_result_gatherer/atdb_result_gatherer/result_gatherer.py index 90ae703ec2b9940891d7e93a2747993cacb19de5..f74479fd47c5d361374f51d5bc8b28b6c963b9e1 100644 --- a/atdb_result_gatherer/atdb_result_gatherer/result_gatherer.py +++ b/atdb_result_gatherer/atdb_result_gatherer/result_gatherer.py @@ -9,8 +9,11 @@ import re def parse_args(): parser = ArgumentParser(description='Gather the results of tasks with a give SAS ID') parser.add_argument('--atdb_url', help='url to the atdb website', default='https://sdc.astron.nl:5554/atdb') - parser.add_argument('sas_id', help='sas id of the observation') + parser.add_argument('--sas_id', help='sas id of the observation') + parser.add_argument('--filter', help='query by filter', default=None) + parser.add_argument('--status', help='query by status', default=None) parser.add_argument('--print_first', help='print first result', action='store_true') + parser.add_argument('--to_gftp_path', help='convert srm url to gftp', action='store_true') parser.add_argument('paths', help='path to the tasks to extract (e.g. outputs.uv_coverage)', nargs='+') return parser.parse_args() @@ -31,7 +34,12 @@ def path_to_nested_dict(content: dict, path: str): raise MissingFieldException(item) if isinstance(leaf, list): return list(map(lambda x: path_to_nested_dict(x, uri[-1]), leaf)) - return leaf[uri[-1]] + + try: + result = leaf[uri[-1]] + except KeyError: + raise MissingFieldException + return result def _get_paginated(url, params, partial_result=None): @@ -51,29 +59,50 @@ def _get_paginated(url, params, partial_result=None): return partial_result -def get_tasks(url: str, sas_id: Dict): - tasks = _get_paginated(f'{url}/tasks', params=dict(sas_id=sas_id)) +def get_tasks(url: str, sas_id: int = None, filter: str = None, status: str = None): + parameters = {} + if sas_id is not None: + parameters['sas_id'] = sas_id + if filter is not None: + parameters['filter'] = filter + if status is not None: + parameters['status'] = status + tasks = _get_paginated(f'{url}/tasks', params=parameters) return tasks +def flatten_nested_list(nested_list): + copy_of = list(nested_list) + while isinstance(copy_of[0], list): + copy_of = reduce(lambda a, b: a + b, copy_of, []) + return copy_of + + def extract_result_from_tasks(tasks, path): - result = reduce( - lambda a, b: a + b, - map( + result = flatten_nested_list( + list(map( lambda item: path_to_nested_dict(item, path), tasks - ), - []) + ))) return result +def convert_srm_to_gftp(path): + if 'srm' in path: + path = re.sub('^srm://srm.grid.sara.nl[:0-9]*', 'gsiftp://gridftp.grid.sara.nl', path) + return path + + def main(): args = parse_args() - tasks = get_tasks(args.atdb_url, args.sas_id) + tasks = get_tasks(args.atdb_url, sas_id=args.sas_id, filter=args.filter, status=args.status) if args.print_first: print(json.dumps(tasks[0], indent=4)) else: results = [extract_result_from_tasks(tasks, path) for path in args.paths] for result in zip(*results): + if args.to_gftp_path: + result = list(map(convert_srm_to_gftp, result)) + print(*result)