diff --git a/atdb/communication.py b/atdb/communication.py index e58161f1f1535a2c1c165d53a66b21a623b99d93..b466c2027a112fa859ec8e33de51f53b34c04bc4 100644 --- a/atdb/communication.py +++ b/atdb/communication.py @@ -123,11 +123,12 @@ class APIConnector: for item in drf_reply.results: yield item - def update_task(self, task_id, content): + def update_task_processed_size(self, task_id, processed_size): """ Change the whole task content """ - self._request_path("PUT", f"tasks/{task_id}", content=content) + return self._request_path("PUT", f"tasks/{task_id}", + content={"size_processed": processed_size}) def change_task_status(self, task_id, status) -> None: """ diff --git a/atdb/fix.py b/atdb/fix.py index ec55ada00fe56380c534adde24a4251a7ee2927d..196a00d7b38534e54ad1442a9a762f6d3957a617 100644 --- a/atdb/fix.py +++ b/atdb/fix.py @@ -1,6 +1,7 @@ import atdb.communication as com -import tqdm +import logging +logger = logging.getLogger('fix') def aggregate_on_tree(tree, field): if isinstance(tree, dict) and field in tree: return tree[field] @@ -19,17 +20,25 @@ def aggregate_on_tree(tree, field): def compute_output_sizes(outputs): if outputs is not None: - return aggregate_on_tree({key: value for key, value in outputs.items if key != 'ingest'}, 'size') + return aggregate_on_tree({key: value for key, value in outputs.items() if key != 'ingest'}, 'size') else: return 0 -def fix_computed_sizes(connector): - for task in tqdm.tqdm(connector.list_iter('tasks')): +def fix_computed_sizes(connector, dry_run=True): + for task in connector.list_iter('tasks'): task_id = task['id'] + size_before = task['size_processed'] total_output_size = compute_output_sizes(task['outputs']) task['size_processed'] = total_output_size - connector.update_task(task_id, task) + if not dry_run: + if size_before != total_output_size: + connector.update_task_processed_size(task_id, total_output_size) + else: + if size_before != total_output_size: + logger.info('Dry run: Size updated for %s from %s to %s', task_id, + size_before, total_output_size) + def fix(args): connector = com.APIConnector.from_args(args) - fix_computed_sizes(connector) \ No newline at end of file + fix_computed_sizes(connector, dry_run=args.dry_run) \ No newline at end of file diff --git a/atdb/main.py b/atdb/main.py index b6d1e1d72231cd193cead91dfecec25501a51430..1637e3dd66a06cb4e7f9e5e832c091042f9a40f5 100644 --- a/atdb/main.py +++ b/atdb/main.py @@ -8,6 +8,7 @@ from argparse import ArgumentParser, Namespace from configparser import ConfigParser from atdb.prune import prune +from atdb.fix import fix DEFAULT_PATH = os.path.expanduser("~/.config/ldv/services.cfg") logging.basicConfig( @@ -78,6 +79,8 @@ def main(): if args.operation == "prune": prune(args) + elif args.operation == "fix": + fix(args) else: parser.print_help()