From f2d0d7817a6a1fd95198ccead4279260a33a5e78 Mon Sep 17 00:00:00 2001
From: mancini <mancini@astron.nl>
Date: Mon, 5 Dec 2022 14:16:09 +0100
Subject: [PATCH] Add various fixes
---
atdb/communication.py | 5 +++--
atdb/fix.py | 21 +++++++++++++++------
atdb/main.py | 3 +++
3 files changed, 21 insertions(+), 8 deletions(-)
diff --git a/atdb/communication.py b/atdb/communication.py
index e58161f..b466c20 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 ec55ada..196a00d 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 b6d1e1d..1637e3d 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()
--
GitLab