diff --git a/atdb/communication.py b/atdb/communication.py index a7cb8e86070569df1ffd6b2d3733b4f0e244a1b2..e58161f1f1535a2c1c165d53a66b21a623b99d93 100644 --- a/atdb/communication.py +++ b/atdb/communication.py @@ -1,8 +1,9 @@ """ This module is responsible for the communication to and from ATDB """ -from typing import List, Generator from argparse import Namespace +from typing import List, Generator + import requests @@ -122,6 +123,12 @@ class APIConnector: for item in drf_reply.results: yield item + def update_task(self, task_id, content): + """ + Change the whole task content + """ + self._request_path("PUT", f"tasks/{task_id}", content=content) + def change_task_status(self, task_id, status) -> None: """ Change the status of a task diff --git a/atdb/fix.py b/atdb/fix.py new file mode 100644 index 0000000000000000000000000000000000000000..ec55ada00fe56380c534adde24a4251a7ee2927d --- /dev/null +++ b/atdb/fix.py @@ -0,0 +1,35 @@ +import atdb.communication as com +import tqdm + +def aggregate_on_tree(tree, field): + if isinstance(tree, dict) and field in tree: + return tree[field] + elif isinstance(tree, dict): + total = 0 + for value in tree.values(): + total += aggregate_on_tree(value, field) + return total + elif isinstance(tree, list): + total = 0 + for item in tree: + total += aggregate_on_tree(item, field) + return total + else: + return 0 + +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') + else: + return 0 + +def fix_computed_sizes(connector): + for task in tqdm.tqdm(connector.list_iter('tasks')): + task_id = task['id'] + total_output_size = compute_output_sizes(task['outputs']) + task['size_processed'] = total_output_size + connector.update_task(task_id, task) + +def fix(args): + connector = com.APIConnector.from_args(args) + fix_computed_sizes(connector) \ No newline at end of file diff --git a/atdb/main.py b/atdb/main.py index 2268a3a9ed2a2457cc5aae4d26636e175ac674c0..b6d1e1d72231cd193cead91dfecec25501a51430 100644 --- a/atdb/main.py +++ b/atdb/main.py @@ -61,6 +61,9 @@ def parse_args() -> (Namespace, ArgumentParser): prune_parser = subparser.add_parser("prune") prune_parser.add_argument("--workflow_id", help="Filters by workflow id") prune_parser.add_argument("--status", help="Filter by status") + + _ = subparser.add_parser("fix") + return parser.parse_args(), parser