From 69ee26a3df42f70b4355c345c02d3bb5620f55e0 Mon Sep 17 00:00:00 2001
From: Robert Avram <robertavram@users.noreply.github.com>
Date: Tue, 2 Apr 2024 12:06:16 -0400
Subject: [PATCH] add rebuild tree functionality

---
 django_api/etools_prp/apps/core/admin.py          | 9 ++++++---
 django_api/etools_prp/apps/core/cartodb.py        | 9 ++++++++-
 django_api/etools_prp/apps/core/locations_sync.py | 6 +++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/django_api/etools_prp/apps/core/admin.py b/django_api/etools_prp/apps/core/admin.py
index 45eb7d3a0..d3c495df0 100644
--- a/django_api/etools_prp/apps/core/admin.py
+++ b/django_api/etools_prp/apps/core/admin.py
@@ -3,7 +3,7 @@
 from leaflet.admin import LeafletGeoAdmin
 from unicef_locations.models import CartoDBTable
 
-from .cartodb import import_locations
+from .cartodb import import_locations, rebuild_tree
 from .forms import CartoDBTableForm
 from .models import Location, PRPRoleOld, Realm, ResponsePlan, Workspace
 
@@ -48,7 +48,7 @@ class CartoDBTableAdmin(admin.ModelAdmin):
         'parent_table_name',
     )
 
-    actions = ('import_sites',)
+    actions = ('import_sites', 'rebuild_tree')
     raw_id_fields = ('parent', )
 
     def parent_table_name(self, obj):
@@ -56,7 +56,10 @@ def parent_table_name(self, obj):
 
     def import_sites(self, request, queryset):
         for table in queryset:
-            import_locations(table.pk)
+            import_locations.delay(table.pk)
+
+    def rebuild_tree(self, request, queryset):
+        rebuild_tree.delay()
 
 
 class WorkspaceAdmin(admin.ModelAdmin):
diff --git a/django_api/etools_prp/apps/core/cartodb.py b/django_api/etools_prp/apps/core/cartodb.py
index 42a1be61c..75c1fe342 100644
--- a/django_api/etools_prp/apps/core/cartodb.py
+++ b/django_api/etools_prp/apps/core/cartodb.py
@@ -10,6 +10,7 @@
 from celery.utils.log import get_task_logger
 from pyrestcli.auth import BaseAuthClient
 from unicef_locations.models import CartoDBTable
+from unicef_locations.utils import get_location_model
 
 from etools_prp.apps.core.locations_sync import PRPLocationSynchronizer
 from etools_prp.apps.core.models import Location
@@ -156,10 +157,16 @@ def create_location(pcode,
 
 @shared_task
 def import_locations(carto_table_pk):
-    print("import locations started")
+    logger.info("import locations started")
     PRPLocationSynchronizer(carto_table_pk).sync()
 
 
+@shared_task
+def rebuild_tree():
+    logger.info("rebuild tree started")
+    get_location_model().objects.rebuild()
+
+
 @shared_task
 def update_sites_from_cartodb(carto_table_pk):
     """
diff --git a/django_api/etools_prp/apps/core/locations_sync.py b/django_api/etools_prp/apps/core/locations_sync.py
index 27fff694c..3d87fc215 100644
--- a/django_api/etools_prp/apps/core/locations_sync.py
+++ b/django_api/etools_prp/apps/core/locations_sync.py
@@ -146,9 +146,9 @@ def create_or_update_locations(self, batch_size=500):
                 for loc in newly_created:
                     loc.workspaces.add(self.workspace)
 
-        logger.info("Rebuilding the tree, have patience")
-        get_location_model().objects.rebuild()
-        logger.info("Rebuilt")
+        logger.info("Rebuilding the tree, have patience -deferred for now. done separately")
+        # get_location_model().objects.rebuild()
+        # logger.info("Rebuilt")
         return new, updated, skipped, error
 
     def clean_upper_level(self):