From a277998e1de770b1979074b9732bf18c8a69e73f Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Mon, 5 Dec 2022 10:48:45 +0100
Subject: [PATCH] Add first implementation
---
atdb/communication.py | 9 ++++++++-
atdb/fix.py | 35 +++++++++++++++++++++++++++++++++++
atdb/main.py | 3 +++
3 files changed, 46 insertions(+), 1 deletion(-)
create mode 100644 atdb/fix.py
diff --git a/atdb/communication.py b/atdb/communication.py
index a7cb8e8..e58161f 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 0000000..ec55ada
--- /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 2268a3a..b6d1e1d 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
--
GitLab