Add limit and offset to queries under batch processing

Author: aditya-balachander

cumulusci/tasks/bulkdata/select_utils.py

@@ -20,17 +20,23 @@ def __init__(self, strategy: SelectStrategy):
         self.strategy = strategy
 
     def select_generate_query(
-        self, sobject: str, fields: T.List[str], num_records: int
+        self,
+        sobject: str,
+        fields: T.List[str],
+        limit: T.Union[int, None],
+        offset: T.Union[int, None],
     ):
         # For STANDARD strategy
         if self.strategy == SelectStrategy.STANDARD:
-            return standard_generate_query(sobject=sobject, num_records=num_records)
+            return standard_generate_query(sobject=sobject, limit=limit, offset=offset)
         # For SIMILARITY strategy
         elif self.strategy == SelectStrategy.SIMILARITY:
-            return similarity_generate_query(sobject=sobject, fields=fields)
+            return similarity_generate_query(
+                sobject=sobject, fields=fields, limit=limit, offset=offset
+            )
         # For RANDOM strategy
         elif self.strategy == SelectStrategy.RANDOM:
-            return standard_generate_query(sobject=sobject, num_records=num_records)
+            return standard_generate_query(sobject=sobject, limit=limit, offset=offset)
 
     def select_post_process(
         self, load_records, query_records: list, num_records: int, sobject: str
@@ -53,7 +59,7 @@ def select_post_process(
 
 
 def standard_generate_query(
-    sobject: str, num_records: int
+    sobject: str, limit: T.Union[int, None], offset: T.Union[int, None]
 ) -> T.Tuple[str, T.List[str]]:
     """Generates the SOQL query for the standard (as well as random) selection strategy"""
     # Get the WHERE clause from DEFAULT_DECLARATIONS if available
@@ -66,8 +72,8 @@ def standard_generate_query(
     query = f"SELECT Id FROM {sobject}"
     if where_clause:
         query += f" WHERE {where_clause}"
-    query += f" LIMIT {num_records}"
-
+    query += f" LIMIT {limit}" if limit else ""
+    query += f" OFFSET {offset}" if offset else ""
     return query, ["Id"]
 
 
@@ -98,6 +104,8 @@ def standard_post_process(
 def similarity_generate_query(
     sobject: str,
     fields: T.List[str],
+    limit: T.Union[int, None],
+    offset: T.Union[int, None],
 ) -> T.Tuple[str, T.List[str]]:
     """Generates the SOQL query for the similarity selection strategy"""
     # Get the WHERE clause from DEFAULT_DECLARATIONS if available
@@ -114,7 +122,8 @@ def similarity_generate_query(
     query = f"SELECT {fields_to_query} FROM {sobject}"
     if where_clause:
         query += f" WHERE {where_clause}"
-
+    query += f" LIMIT {limit}" if limit else ""
+    query += f" OFFSET {offset}" if offset else ""
     return query, fields
 
 

cumulusci/tasks/bulkdata/step.py

@@ -462,16 +462,21 @@ def select_records(self, records):
             )
 
             # Generate and execute SOQL query
+            # (not passing offset as it is not supported in Bulk)
             (
                 select_query,
                 query_fields,
             ) = self.select_operation_executor.select_generate_query(
-                self.sobject, self.fields, num_records
+                sobject=self.sobject, fields=self.fields, limit=num_records, offset=None
             )
             if self.selection_filter:
                 # Generate user filter query if selection_filter is present (offset clause not supported)
                 user_query = generate_user_filter_query(
-                    self.selection_filter, self.sobject, ["Id"], num_records, None
+                    filter_clause=self.selection_filter,
+                    sobject=self.sobject,
+                    fields=["Id"],
+                    limit_clause=num_records,
+                    offset_clause=None,
                 )
                 # Execute the user query using Bulk API
                 user_query_executor = get_query_operation(
@@ -508,19 +513,22 @@ def select_records(self, records):
             selected_records,
             error_message,
         ) = self.select_operation_executor.select_post_process(
-            records, query_records, num_records, self.sobject
+            load_records=records,
+            query_records=query_records,
+            num_records=num_records,
+            sobject=self.sobject,
         )
         if not error_message:
             self.select_results.extend(selected_records)
 
         # Update job result based on selection outcome
         self.job_result = DataOperationJobResult(
-            DataOperationStatus.SUCCESS
+            status=DataOperationStatus.SUCCESS
             if len(self.select_results)
             else DataOperationStatus.JOB_FAILURE,
-            [error_message] if error_message else [],
-            len(self.select_results),
-            0,
+            job_errors=[error_message] if error_message else [],
+            records_processed=len(self.select_results),
+            total_row_errors=0,
         )
 
     def _execute_select_query(self, select_query: str, query_fields: List[str]):
@@ -814,13 +822,20 @@ def convert(rec, fields):
                 select_query,
                 query_fields,
             ) = self.select_operation_executor.select_generate_query(
-                self.sobject, self.fields, num_records
+                sobject=self.sobject,
+                fields=self.fields,
+                limit=num_records,
+                offset=offset,
             )
 
             # If user given selection filter present, create composite request
             if self.selection_filter:
                 user_query = generate_user_filter_query(
-                    self.selection_filter, self.sobject, ["Id"], num_records, offset
+                    filter_clause=self.selection_filter,
+                    sobject=self.sobject,
+                    fields=["Id"],
+                    limit_clause=num_records,
+                    offset_clause=offset,
                 )
                 query_records.extend(
                     self._execute_composite_query(
@@ -843,20 +858,23 @@ def convert(rec, fields):
             selected_records,
             error_message,
         ) = self.select_operation_executor.select_post_process(
-            records, query_records, total_num_records, self.sobject
+            load_records=records,
+            query_records=query_records,
+            num_records=total_num_records,
+            sobject=self.sobject,
         )
         if not error_message:
             # Add selected records from this batch to the overall results
             self.results.extend(selected_records)
 
         # Update the job result based on the overall selection outcome
         self.job_result = DataOperationJobResult(
-            DataOperationStatus.SUCCESS
+            status=DataOperationStatus.SUCCESS
             if len(self.results)  # Check the overall results length
             else DataOperationStatus.JOB_FAILURE,
-            [error_message] if error_message else [],
-            len(self.results),
-            0,
+            job_errors=[error_message] if error_message else [],
+            records_processed=len(self.results),
+            total_row_errors=0,
         )
 
     def _execute_composite_query(self, select_query, user_query, query_fields):

cumulusci/tasks/bulkdata/tests/test_select_utils.py

@@ -11,53 +11,61 @@
 def test_standard_generate_query_with_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.STANDARD)
     sobject = "Account"  # Assuming Account has a declaration in DEFAULT_DECLARATIONS
-    num_records = 5
+    limit = 5
+    offset = 2
     query, fields = select_operator.select_generate_query(
-        sobject=sobject, fields=[], num_records=num_records
+        sobject=sobject, fields=[], limit=limit, offset=offset
     )
 
     assert "WHERE" in query  # Ensure WHERE clause is included
-    assert f"LIMIT {num_records}" in query
+    assert f"LIMIT {limit}" in query
+    assert f"OFFSET {offset}" in query
     assert fields == ["Id"]
 
 
 def test_standard_generate_query_without_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.STANDARD)
     sobject = "Contact"  # Assuming no declaration for this object
-    num_records = 3
+    limit = 3
+    offset = None
     query, fields = select_operator.select_generate_query(
-        sobject=sobject, fields=[], num_records=num_records
+        sobject=sobject, fields=[], limit=limit, offset=offset
     )
 
     assert "WHERE" not in query  # No WHERE clause should be present
-    assert f"LIMIT {num_records}" in query
+    assert f"LIMIT {limit}" in query
+    assert "OFFSET" not in query
     assert fields == ["Id"]
 
 
 # Test Cases for random generate query
 def test_random_generate_query_with_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.RANDOM)
     sobject = "Account"  # Assuming Account has a declaration in DEFAULT_DECLARATIONS
-    num_records = 5
+    limit = 5
+    offset = 2
     query, fields = select_operator.select_generate_query(
-        sobject=sobject, fields=[], num_records=num_records
+        sobject=sobject, fields=[], limit=limit, offset=offset
     )
 
     assert "WHERE" in query  # Ensure WHERE clause is included
-    assert f"LIMIT {num_records}" in query
+    assert f"LIMIT {limit}" in query
+    assert f"OFFSET {offset}" in query
     assert fields == ["Id"]
 
 
 def test_random_generate_query_without_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.RANDOM)
     sobject = "Contact"  # Assuming no declaration for this object
-    num_records = 3
+    limit = 3
+    offset = None
     query, fields = select_operator.select_generate_query(
-        sobject=sobject, fields=[], num_records=num_records
+        sobject=sobject, fields=[], limit=limit, offset=offset
     )
 
     assert "WHERE" not in query  # No WHERE clause should be present
-    assert f"LIMIT {num_records}" in query
+    assert f"LIMIT {limit}" in query
+    assert "OFFSET" not in query
     assert fields == ["Id"]
 
 
@@ -141,25 +149,31 @@ def test_random_post_process_with_no_records():
 def test_similarity_generate_query_with_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.SIMILARITY)
     sobject = "Account"  # Assuming Account has a declaration in DEFAULT_DECLARATIONS
-    num_records = 5
+    limit = 5
+    offset = 2
     query, fields = select_operator.select_generate_query(
-        sobject, ["Name"], num_records
+        sobject, ["Name"], limit, offset
     )
 
     assert "WHERE" in query  # Ensure WHERE clause is included
     assert fields == ["Id", "Name"]
+    assert f"LIMIT {limit}" in query
+    assert f"OFFSET {offset}" in query
 
 
 def test_similarity_generate_query_without_default_record_declaration():
     select_operator = SelectOperationExecutor(SelectStrategy.SIMILARITY)
     sobject = "Contact"  # Assuming no declaration for this object
-    num_records = 3
+    limit = 3
+    offset = None
     query, fields = select_operator.select_generate_query(
-        sobject, ["Name"], num_records
+        sobject, ["Name"], limit, offset
     )
 
     assert "WHERE" not in query  # No WHERE clause should be present
     assert fields == ["Id", "Name"]
+    assert f"LIMIT {limit}" in query
+    assert "OFFSET" not in query
 
 
 def test_levenshtein_distance():