Add first implementation

atdb/communication.py

 """
 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

atdb/fix.py

+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

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