Context records: Optimize record relations - management command - tests

1 files changed, 21 insertions, 45 deletions

diff --git a/ishtar_common/management/commands/relations_update_cache_tables.py b/ishtar_common/management/commands/relations_update_cache_tables.py
index ab7f134ff..3e2dfaef5 100644
--- a/ishtar_common/management/commands/relations_update_cache_tables.py
+++ b/ishtar_common/management/commands/relations_update_cache_tables.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2013-2018 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2021 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -24,60 +24,36 @@ from django.core.management.base import BaseCommand
 from django.apps import apps
 
 
-APPS = ['ishtar_common', 'archaeological_operations',
-        'archaeological_context_records', 'archaeological_finds',
-        'archaeological_warehouse']
+CACHE_TABLES = [
+    "archaeological_context_records.ContextRecordTree"
+]
 
 
 class Command(BaseCommand):
     args = ''
-    help = 'Regenerate geo, cached labels and search vectors'
+    help = 'Update all relations for cache tables'
 
     def add_arguments(self, parser):
-        parser.add_argument('app_name', nargs='?', default=None,
-                            choices=APPS)
-        parser.add_argument('model_name', nargs='?', default=None)
+        parser.add_argument('table', nargs='?', default=None,
+                            choices=CACHE_TABLES)
         parser.add_argument(
             '--quiet', dest='quiet', action='store_true',
             help='Quiet output')
 
     def handle(self, *args, **options):
         quiet = options['quiet']
-        limit = options['app_name']
-        model_name = options['model_name']
-        if model_name:
-            model_name = model_name.lower()
-        for app in APPS:
-            if limit and app != limit:
-                continue
+        tables = CACHE_TABLES
+        if options.get("table", None):
+            table = options.get("table", None)
+            if table not in CACHE_TABLES:
+                sys.stdout.write("{} not a valid cache table\n".format(table))
+                return
+            tables = [table]
+        for table in tables:
             if not quiet:
-                print("* app: {}".format(app))
-            for model in apps.get_app_config(app).get_models():
-                if model_name and model.__name__.lower() != model_name:
-                    continue
-                if model.__name__.startswith('Historical'):
-                    continue
-                if not bool(
-                        [k for k in dir(model)
-                         if k.startswith('_generate_') or
-                            k == "search_vector"]):
-                    continue
-                msg = "-> processing {}: ".format(model._meta.verbose_name)
-                ln = model.objects.count()
-                for idx, obj_id in enumerate(model.objects.values('pk').all()):
-                    obj = model.objects.get(pk=obj_id['pk'])
-                    obj.skip_history_when_saving = True
-                    obj._no_move = True
-                    if hasattr(obj, "point_source") and obj.point_source in (
-                            "M", "T"):
-                        obj.point = None
-                        obj.point_2d = None
-                        obj.x = None
-                        obj.y = None
-                    cmsg = "\r{} {}/{}".format(msg, idx + 1, ln)
-                    if not quiet:
-                        sys.stdout.write(cmsg)
-                        sys.stdout.flush()
-                    obj.save()
-                if not quiet:
-                    sys.stdout.write("\n")
+                print("* table: {}".format(table))
+            app, tablename = table.split(".")
+            model = apps.get_app_config(app).get_model(tablename)
+            model.regenerate_all(quiet=quiet)
+            if not quiet:
+                sys.stdout.write("\n")