Trigger Geopoint ES Index on Geospatial Feature Flag Enable

corehq/apps/geospatial/const.py

@@ -133,3 +133,5 @@
         }
     }
 }
+
+INDEX_ES_TASK_HELPER_BASE_KEY = 'geo_cases_index_cases'

corehq/apps/geospatial/management/commands/index_geolocation_case_properties.py

@@ -41,15 +41,23 @@ def index_case_docs(domain, query_limit=DEFAULT_QUERY_LIMIT, chunk_size=DEFAULT_
     query = _es_case_query(domain, geo_case_property, case_type)
     count = query.count()
     print(f'{count} case(s) to process')
-    batch_count = 1
-    if query_limit:
-        batch_count = math.ceil(count / query_limit)
+    batch_count = get_batch_count(count, query_limit)
     print(f"Cases will be processed in {batch_count} batches")
     for i in range(batch_count):
         print(f'Processing {i+1}/{batch_count}')
-        query = _es_case_query(domain, geo_case_property, case_type, size=query_limit)
-        case_ids = query.get_ids()
-        _index_case_ids(domain, case_ids, chunk_size)
+        process_batch(domain, geo_case_property, case_type, query_limit, chunk_size)
+
+
+def get_batch_count(doc_count, query_limit):
+    if not query_limit:
+        return 1
+    return math.ceil(doc_count / query_limit)
+
+
+def process_batch(domain, geo_case_property, case_type, query_limit, chunk_size):
+    query = _es_case_query(domain, geo_case_property, case_type, size=query_limit)
+    case_ids = query.get_ids()
+    _index_case_ids(domain, case_ids, chunk_size)
 
 
 def _index_case_ids(domain, case_ids, chunk_size):

corehq/apps/geospatial/tasks.py

@@ -1,5 +1,24 @@
+from django.utils.translation import gettext as _
+
+from corehq.util.decorators import serial_task
+
 from corehq.apps.celery import task
-from corehq.apps.geospatial.utils import CeleryTaskTracker, update_cases_owner
+from corehq.apps.geospatial.const import INDEX_ES_TASK_HELPER_BASE_KEY
+from corehq.apps.geospatial.utils import (
+    get_celery_task_tracker,
+    CeleryTaskTracker,
+    update_cases_owner,
+    get_geo_case_property,
+)
+from corehq.apps.geospatial.management.commands.index_geolocation_case_properties import (
+    _es_case_query,
+    get_batch_count,
+    process_batch,
+    DEFAULT_QUERY_LIMIT,
+    DEFAULT_CHUNK_SIZE,
+)
+
+from settings import MAX_GEOSPATIAL_INDEX_DOC_LIMIT
 
 
 @task(queue="background_queue", ignore_result=True)
@@ -9,3 +28,38 @@ def geo_cases_reassignment_update_owners(domain, case_owner_updates_dict, task_k
     finally:
         celery_task_tracker = CeleryTaskTracker(task_key)
         celery_task_tracker.mark_completed()
+
+
+@serial_task('async-index-es-docs', timeout=30 * 60, queue='background_queue', ignore_result=True)
+def index_es_docs_with_location_props(domain):
+    celery_task_tracker = get_celery_task_tracker(domain, INDEX_ES_TASK_HELPER_BASE_KEY)
+    if celery_task_tracker.is_active():
+        return
+
+    geo_case_prop = get_geo_case_property(domain)
+    query = _es_case_query(domain, geo_case_prop)
+    doc_count = query.count()
+    if doc_count > MAX_GEOSPATIAL_INDEX_DOC_LIMIT:
+        celery_task_tracker.set_message(
+            _('This domain contains too many cases and so they will not be made available '
+              'for use by this feature. Please reach out to support.')
+        )
+        return
+
+    celery_task_tracker.mark_requested()
+    batch_count = get_batch_count(doc_count, DEFAULT_QUERY_LIMIT)
+    try:
+        for i in range(batch_count):
+            progress = (i / batch_count) * 100
+            celery_task_tracker.set_message(
+                _(f'Cases are being made ready for use by this feature. Please be patient. ({progress}%)')
+            )
+            process_batch(
+                domain,
+                geo_case_prop,
+                case_type=None,
+                query_limit=DEFAULT_QUERY_LIMIT,
+                chunk_size=DEFAULT_CHUNK_SIZE,
+            )
+    finally:
+        celery_task_tracker.mark_completed()

corehq/apps/geospatial/utils.py

@@ -33,6 +33,12 @@ def get_geo_user_property(domain):
     return config.user_location_property_name
 
 
+def get_celery_task_tracker(domain, base_key):
+    task_key = f'{base_key}_{domain}'
+    message_key = f'{base_key}_message_{domain}'
+    return CeleryTaskTracker(task_key, message_key)
+
+
 def _format_coordinates(lat, lon):
     return f"{lat} {lon} 0.0 0.0"
 
@@ -221,8 +227,9 @@ class CeleryTaskTracker(object):
     Simple Helper class using redis to track if a celery task was requested and is not completed yet.
     """
 
-    def __init__(self, task_key):
+    def __init__(self, task_key, message_key=None):
         self.task_key = task_key
+        self.message_key = message_key
         self._client = get_redis_client()
 
     def mark_requested(self, timeout=ONE_DAY):
@@ -234,4 +241,14 @@ def is_active(self):
         return self._client.has_key(self.task_key)
 
     def mark_completed(self):
+        self.clear_message()
         return self._client.delete(self.task_key)
+
+    def get_message(self):
+        return self._client.get(self.message_key)
+
+    def set_message(self, message, timeout=ONE_DAY * 3):
+        return self._client.set(self.message_key, message, timeout=timeout)
+
+    def clear_message(self):
+        return self._client.delete(self.message_key)

corehq/toggles/__init__.py

@@ -2555,13 +2555,22 @@ def _handle_attendance_tracking_role(domain, is_enabled):
     save_fn=_handle_attendance_tracking_role,
 )
 
+
+def _handle_geospatial_es_index(domain, is_enabled):
+    from corehq.apps.geospatial.es import index_es_docs_with_location_props
+
+    if is_enabled:
+        index_es_docs_with_location_props.delay(domain)
+
+
 GEOSPATIAL = StaticToggle(
     'geospatial',
     'Allows access to GIS functionality',
     TAG_SOLUTIONS_LIMITED,
     namespaces=[NAMESPACE_DOMAIN],
     description='Additional views will be added allowing for visually viewing '
-                'and assigning cases on a map.'
+                'and assigning cases on a map.',
+    save_fn=_handle_geospatial_es_index,
 
 )