Modify functionality to return records in order of the user query to support ORDER BY operation

Author: aditya-balachander

cumulusci/tasks/bulkdata/step.py

@@ -491,7 +491,11 @@ def select_records(self, records):
                 user_query_records = user_query_executor.get_results()
 
                 # Find intersection based on 'Id'
-                user_query_ids = set(record[0] for record in user_query_records)
+                user_query_ids = (
+                    list(record[0] for record in user_query_records)
+                    if user_query_records
+                    else []
+                )
 
             # Execute the main select query using Bulk API
             select_query_records = self._execute_select_query(
@@ -500,10 +504,16 @@ def select_records(self, records):
 
             # If user_query_ids exist, filter select_query_records based on the intersection of Ids
             if self.selection_filter:
+                # Create a dictionary to map IDs to their corresponding records
+                id_to_record_map = {
+                    record[query_fields.index("Id")]: record
+                    for record in select_query_records
+                }
+                # Extend query_records in the order of user_query_ids
                 query_records.extend(
                     record
-                    for record in select_query_records
-                    if record[query_fields.index("Id")] in user_query_ids
+                    for id in user_query_ids
+                    if (record := id_to_record_map.get(id)) is not None
                 )
             else:
                 query_records.extend(select_query_records)
@@ -928,12 +938,17 @@ def convert(rec, fields):
                     f"{sub_response['body'][0]['errorCode']}: {sub_response['body'][0]['message']}"
                 )
         # Find intersection based on 'Id'
-        user_query_ids = set(record[0] for record in user_query_records)
+        user_query_ids = list(record[0] for record in user_query_records)
+        # Create a dictionary to map IDs to their corresponding records
+        id_to_record_map = {
+            record[query_fields.index("Id")]: record for record in select_query_records
+        }
 
+        # Extend query_records in the order of user_query_ids
         return [
             record
-            for record in select_query_records
-            if record[query_fields.index("Id")] in user_query_ids
+            for id in user_query_ids
+            if (record := id_to_record_map.get(id)) is not None
         ]
 
     def get_results(self):

cumulusci/tasks/bulkdata/tests/test_step.py

@@ -701,6 +701,70 @@ def test_select_records_user_selection_filter_success(self, download_mock):
                 == 3
             )
 
+    @mock.patch("cumulusci.tasks.bulkdata.step.download_file")
+    def test_select_records_user_selection_filter_order_success(self, download_mock):
+        # Set up mock context and BulkApiDmlOperation
+        context = mock.Mock()
+        step = BulkApiDmlOperation(
+            sobject="Contact",
+            operation=DataOperationType.QUERY,
+            api_options={"batch_size": 10, "update_key": "LastName"},
+            context=context,
+            fields=["LastName"],
+            selection_strategy=SelectStrategy.STANDARD,
+            selection_filter="ORDER BY CreatedDate",
+        )
+
+        # Mock Bulk API responses
+        step.bulk.endpoint = "https://test"
+        step.bulk.create_query_job.return_value = "JOB"
+        step.bulk.query.return_value = "BATCH"
+        step.bulk.get_query_batch_result_ids.return_value = ["RESULT"]
+
+        # Mock the downloaded CSV content with a single record
+        download_mock.return_value = io.StringIO(
+            """Id
+003000000000001
+003000000000002
+003000000000003"""
+        )
+        # Mock the query operation
+        with mock.patch(
+            "cumulusci.tasks.bulkdata.step.get_query_operation"
+        ) as query_operation_mock:
+            query_operation_mock.return_value = mock.Mock()
+            query_operation_mock.return_value.query = mock.Mock()
+            query_operation_mock.return_value.get_results = mock.Mock()
+            query_operation_mock.return_value.get_results.return_value = [
+                ["003000000000003"],
+                ["003000000000001"],
+                ["003000000000002"],
+            ]
+
+            # Mock the _wait_for_job method to simulate a successful job
+            step._wait_for_job = mock.Mock()
+            step._wait_for_job.return_value = DataOperationJobResult(
+                DataOperationStatus.SUCCESS, [], 0, 0
+            )
+
+            # Prepare input records
+            records = iter([["Test1"], ["Test2"], ["Test3"]])
+
+            # Execute the select_records operation
+            step.start()
+            step.select_records(records)
+            step.end()
+
+            # Get the results and assert their properties
+            results = list(step.get_results())
+            assert (
+                len(results) == 3
+            )  # Expect 3 results (matching the input records count)
+            # Assert that all results are in the order given by user query
+            assert results[0].id == "003000000000003"
+            assert results[1].id == "003000000000001"
+            assert results[2].id == "003000000000002"
+
     @mock.patch("cumulusci.tasks.bulkdata.step.download_file")
     def test_select_records_user_selection_filter_failure(self, download_mock):
         # Set up mock context and BulkApiDmlOperation
@@ -1428,6 +1492,87 @@ def test_select_records_user_selection_filter_success(self):
             == 3
         )
 
+    @responses.activate
+    def test_select_records_user_selection_filter_order_success(self):
+        mock_describe_calls()
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite:///test.db",
+                    "mapping": "mapping.yml",
+                }
+            },
+        )
+        task.project_config.project__package__api_version = CURRENT_SF_API_VERSION
+        task._init_task()
+
+        responses.add(
+            responses.POST,
+            url=f"https://example.com/services/data/v{CURRENT_SF_API_VERSION}/composite/sobjects",
+            json=[
+                {"id": "003000000000001", "success": True},
+                {"id": "003000000000002", "success": True},
+            ],
+            status=200,
+        )
+        responses.add(
+            responses.POST,
+            url=f"https://example.com/services/data/v{CURRENT_SF_API_VERSION}/composite/sobjects",
+            json=[{"id": "003000000000003", "success": True}],
+            status=200,
+        )
+        step = RestApiDmlOperation(
+            sobject="Contact",
+            operation=DataOperationType.UPSERT,
+            api_options={"batch_size": 10, "update_key": "LastName"},
+            context=task,
+            fields=["LastName"],
+            selection_strategy=SelectStrategy.STANDARD,
+            selection_filter="ORDER BY CreatedDate",
+        )
+
+        results = {
+            "compositeResponse": [
+                {
+                    "body": {
+                        "records": [
+                            {"Id": "003000000000001"},
+                            {"Id": "003000000000002"},
+                            {"Id": "003000000000003"},
+                        ]
+                    },
+                    "referenceId": "select_query",
+                    "httpStatusCode": 200,
+                },
+                {
+                    "body": {
+                        "records": [
+                            {"Id": "003000000000003"},
+                            {"Id": "003000000000001"},
+                            {"Id": "003000000000002"},
+                        ]
+                    },
+                    "referenceId": "user_query",
+                    "httpStatusCode": 200,
+                },
+            ]
+        }
+        step.sf.restful = mock.Mock()
+        step.sf.restful.return_value = results
+        records = iter([["Test1"], ["Test2"], ["Test3"]])
+        step.start()
+        step.select_records(records)
+        step.end()
+
+        # Get the results and assert their properties
+        results = list(step.get_results())
+        assert len(results) == 3  # Expect 3 results (matching the input records count)
+        # Assert that all results are in the order of user_query
+        assert results[0].id == "003000000000003"
+        assert results[1].id == "003000000000001"
+        assert results[2].id == "003000000000002"
+
     @responses.activate
     def test_select_records_user_selection_filter_failure(self):
         mock_describe_calls()