commands - ishtar_maintenance: task - fix operation missing parcels

1 files changed, 42 insertions, 4 deletions

diff --git a/ishtar_common/management/commands/ishtar_maintenance.py b/ishtar_common/management/commands/ishtar_maintenance.py
index 1b9a9be53..1eeab179a 100644
--- a/ishtar_common/management/commands/ishtar_maintenance.py
+++ b/ishtar_common/management/commands/ishtar_maintenance.py
@@ -113,6 +113,32 @@ def task_main_image(quiet=False, log=False):
             sys.stdout.write(f"{nb} main image fixed for {model.__name__}\n")
 
 
+def task_missing_parcels(quiet=False):
+    Parcel = apps.get_model("archaeological_operations", "Parcel")
+    q = Parcel.objects.filter(context_record__isnull=False, operation=None)
+    nb = q.count()
+    if not nb:
+        if not quiet:
+            sys.stdout.write("No parcel to fix.\n")
+            return
+    for idx, parcel in enumerate(q.all()):
+        if not quiet:
+            sys.stdout.write(f"\r[{percent(idx, nb)}] {idx + 1}/{nb}")
+        # assume all context record associated to this parcel are of the same operation
+        cr = parcel.context_record.all()[0]
+        if parcel.operation:
+            continue
+        parcel.operation = cr.operation
+        parcel.skip_history_when_saving = True
+        parcel.save()
+    if not quiet:
+        sys.stdout.write("\n")
+
+
+def percent(current, total):
+    return f"{(current + 1) / total * 100:.1f}".rjust(4, "0") + "%"
+
+
 def get_time():
     return datetime.datetime.now().isoformat().split(".")[0]
 
@@ -126,21 +152,33 @@ TASKS = {
         "help": "regenerate cached label on all tables if necessary",
         "action": task_check_cached_label,
     },
+    "operation_missing_parcels": {
+        "help": "fix lost parcel association on operation from context records.",
+        "action": task_missing_parcels,
+    },
 }
 
 
 class Command(BaseCommand):
-    help = "Launch maintenance tasks: \n" + "\n".join(
-        [f"* {t}: {TASKS[t]['help']}" for t in TASKS.keys()]
-    )
+    help = "Launch maintenance tasks"
+
+    def parser_error(self, message=""):
+        sys.stderr.write(f"{message}\n")
+        self.parser.print_help()
+        sys.exit(2)
 
     def create_parser(self, *args, **kwargs):
         parser = super(Command, self).create_parser(*args, **kwargs)
         parser.formatter_class = RawTextHelpFormatter
+        self.parser = parser
+        parser.error = self.parser_error
         return parser
 
     def add_arguments(self, parser):
-        parser.add_argument("task")
+        task_help = "\n".join(
+            [f"{k} - {TASKS[k]['help']}" for k in sorted(TASKS.keys())]
+        )
+        parser.add_argument("task", help=task_help)
         parser.add_argument(
             "--quiet", dest="quiet", action="store_true", help="Quiet output"
         )