From c2c4644f59f7814c09d9f45e3ca9987d8836eb35 Mon Sep 17 00:00:00 2001
From: Kelvin Muchiri <kelvinmuchiri63@gmail.com>
Date: Mon, 20 Jan 2025 10:38:18 +0300
Subject: [PATCH] create default ordered columns when registering xform repeats

---
 onadata/libs/tests/utils/test_logger_tools.py | 58 ++++++++++++-------
 onadata/libs/utils/logger_tools.py            | 17 ++++--
 2 files changed, 51 insertions(+), 24 deletions(-)

diff --git a/onadata/libs/tests/utils/test_logger_tools.py b/onadata/libs/tests/utils/test_logger_tools.py
index 9513e17a7e..58ec45b148 100644
--- a/onadata/libs/tests/utils/test_logger_tools.py
+++ b/onadata/libs/tests/utils/test_logger_tools.py
@@ -1342,6 +1342,26 @@ def setUp(self):
             object_id=self.xform.pk,
             content_type=ContentType.objects.get_for_model(self.xform),
         )
+        self.expected_columns = OrderedDict(
+            [
+                (
+                    "hospital_repeat",
+                    ["hospital_repeat[1]/hospital", "hospital_repeat[2]/hospital"],
+                ),
+                (
+                    "hospital_repeat/child_repeat",
+                    [
+                        "hospital_repeat[1]/child_repeat[1]/name",
+                        "hospital_repeat[1]/child_repeat[1]/birthweight",
+                        "hospital_repeat[1]/child_repeat[2]/name",
+                        "hospital_repeat[1]/child_repeat[2]/birthweight",
+                        "hospital_repeat[2]/child_repeat[1]/name",
+                        "hospital_repeat[2]/child_repeat[1]/birthweight",
+                    ],
+                ),
+                ("meta/instanceID", None),
+            ]
+        )
 
     def test_register(self):
         """Repeats from all instances are registered"""
@@ -1369,25 +1389,23 @@ def test_register(self):
         ordered_columns = json.loads(
             self.register.extra_data, object_pairs_hook=OrderedDict
         )
-        expected_columns = OrderedDict(
-            [
-                (
-                    "hospital_repeat",
-                    ["hospital_repeat[1]/hospital", "hospital_repeat[2]/hospital"],
-                ),
-                (
-                    "hospital_repeat/child_repeat",
-                    [
-                        "hospital_repeat[1]/child_repeat[1]/name",
-                        "hospital_repeat[1]/child_repeat[1]/birthweight",
-                        "hospital_repeat[1]/child_repeat[2]/name",
-                        "hospital_repeat[1]/child_repeat[2]/birthweight",
-                        "hospital_repeat[2]/child_repeat[1]/name",
-                        "hospital_repeat[2]/child_repeat[1]/birthweight",
-                    ],
-                ),
-                ("meta/instanceID", None),
-            ]
+
+        self.assertEqual(ordered_columns, self.expected_columns)
+
+    def test_register_not_found(self):
+        """Register is created if not found"""
+        self.register.delete()
+        register_xform_export_columns(self.xform)
+
+        exists = MetaData.objects.filter(data_type="export_columns_register").exists()
+
+        self.assertTrue(exists)
+
+        register = MetaData.objects.get(
+            data_type="export_columns_register",
+            object_id=self.xform.pk,
+            content_type=ContentType.objects.get_for_model(self.xform),
         )
+        ordered_columns = json.loads(register.extra_data, object_pairs_hook=OrderedDict)
 
-        self.assertEqual(ordered_columns, expected_columns)
+        self.assertEqual(ordered_columns, self.expected_columns)
diff --git a/onadata/libs/utils/logger_tools.py b/onadata/libs/utils/logger_tools.py
index 6071c9b804..5d4e551bb0 100644
--- a/onadata/libs/utils/logger_tools.py
+++ b/onadata/libs/utils/logger_tools.py
@@ -1564,7 +1564,7 @@ def _update_export_columns_register(instance: Instance, metadata: MetaData) -> N
 
 @transaction.atomic()
 def register_instance_export_columns(instance: Instance) -> None:
-    """Register an Instance's export columns
+    """Register an Instance's repeat columns for export
 
     :param instance: Instance object
     """
@@ -1589,12 +1589,21 @@ def register_xform_export_columns(xform: XForm) -> None:
 
     :param xform: XForm object
     """
-    content_type = ContentType.objects.get_for_model(xform)
+    # Avoid cyclic import by using importlib
+    csv_builder = importlib.import_module("onadata.libs.utils.csv_builder")
+    ordered_columns = OrderedDict()
+    # pylint: disable=protected-access
+    csv_builder.CSVDataFrameBuilder._build_ordered_columns(
+        xform._get_survey(), ordered_columns
+    )
     metadata, _ = MetaData.objects.get_or_create(
-        content_type=content_type,
+        content_type=ContentType.objects.get_for_model(xform),
         object_id=xform.pk,
         data_type=EXPORT_COLUMNS_REGISTER,
-        defaults={"data_value": ""},
+        defaults={
+            "data_value": "",
+            "extra_data": json.dumps(ordered_columns),
+        },
     )
     instance_qs = xform.instances.filter(deleted_at__isnull=True)