Adds selection_filter and RANDOM selection strategy

Author: aditya-balachander

cumulusci/tasks/bulkdata/load.py

@@ -289,7 +289,12 @@ def _execute_step(
                     self, step, self._stream_queried_data(mapping, local_ids, query)
                 )
             step.start()
-            step.load_records(self._stream_queried_data(mapping, local_ids, query))
+            if mapping.action == DataOperationType.SELECT:
+                step.select_records(
+                    self._stream_queried_data(mapping, local_ids, query)
+                )
+            else:
+                step.load_records(self._stream_queried_data(mapping, local_ids, query))
             step.end()
 
             # Process Job Results
@@ -336,6 +341,8 @@ def configure_step(self, mapping):
             self.check_simple_upsert(mapping)
             api_options["update_key"] = mapping.update_key[0]
             action = DataOperationType.UPSERT
+        elif mapping.action == DataOperationType.SELECT:
+            action = DataOperationType.QUERY
         else:
             action = mapping.action
 
@@ -349,6 +356,8 @@ def configure_step(self, mapping):
             fields=fields,
             api=mapping.api,
             volume=query.count(),
+            selection_strategy=mapping.selection_strategy,
+            selection_filter=mapping.selection_filter,
         )
         return step, query
 
@@ -481,10 +490,11 @@ def _process_job_results(self, mapping, step, local_ids):
         """Get the job results and process the results. If we're raising for
         row-level errors, do so; if we're inserting, store the new Ids."""
 
-        is_insert_or_upsert = mapping.action in (
+        is_insert_upsert_or_select = mapping.action in (
             DataOperationType.INSERT,
             DataOperationType.UPSERT,
             DataOperationType.ETL_UPSERT,
+            DataOperationType.SELECT,
         )
 
         conn = self.session.connection()
@@ -500,7 +510,7 @@ def _process_job_results(self, mapping, step, local_ids):
                 break
         # If we know we have no successful inserts, don't attempt to persist Ids.
         # Do, however, drain the generator to get error-checking behavior.
-        if is_insert_or_upsert and (
+        if is_insert_upsert_or_select and (
             step.job_result.records_processed - step.job_result.total_row_errors
         ):
             table = self.metadata.tables[self.ID_TABLE_NAME]
@@ -516,7 +526,7 @@ def _process_job_results(self, mapping, step, local_ids):
         # person account Contact records so lookups to
         # person account Contact records get populated downstream as expected.
         if (
-            is_insert_or_upsert
+            is_insert_upsert_or_select
             and mapping.sf_object == "Contact"
             and self._can_load_person_accounts(mapping)
         ):
@@ -531,7 +541,7 @@ def _process_job_results(self, mapping, step, local_ids):
                     ),
                 )
 
-        if is_insert_or_upsert:
+        if is_insert_upsert_or_select:
             self.session.commit()
 
     def _generate_results_id_map(self, step, local_ids):

cumulusci/tasks/bulkdata/mapping_parser.py

@@ -15,6 +15,7 @@
 from cumulusci.core.enums import StrEnum
 from cumulusci.core.exceptions import BulkDataException
 from cumulusci.tasks.bulkdata.dates import iso_to_date
+from cumulusci.tasks.bulkdata.select_utils import SelectStrategy
 from cumulusci.tasks.bulkdata.step import DataApi, DataOperationType
 from cumulusci.utils import convert_to_snake_case
 from cumulusci.utils.yaml.model_parser import CCIDictModel
@@ -84,7 +85,7 @@ class BulkMode(StrEnum):
 
 ENUM_VALUES = {
     v.value.lower(): v.value
-    for enum in [BulkMode, DataApi, DataOperationType]
+    for enum in [BulkMode, DataApi, DataOperationType, SelectStrategy]
     for v in enum.__members__.values()
 }
 
@@ -107,9 +108,13 @@ class MappingStep(CCIDictModel):
     ] = None  # default should come from task options
     anchor_date: Optional[Union[str, date]] = None
     soql_filter: Optional[str] = None  # soql_filter property
+    selection_strategy: SelectStrategy = SelectStrategy.STANDARD  # selection strategy
+    selection_filter: Optional[
+        str
+    ] = None  # filter to be added at the end of select query
     update_key: T.Union[str, T.Tuple[str, ...]] = ()  # only for upserts
 
-    @validator("bulk_mode", "api", "action", pre=True)
+    @validator("bulk_mode", "api", "action", "selection_strategy", pre=True)
     def case_normalize(cls, val):
         if isinstance(val, Enum):
             return val

cumulusci/tasks/bulkdata/select_utils.py

@@ -1,3 +1,4 @@
+import random
 import typing as T
 
 from cumulusci.core.enums import StrEnum
@@ -9,14 +10,52 @@
 class SelectStrategy(StrEnum):
     """Enum defining the different selection strategies requested."""
 
-    RANDOM = "random"
+    STANDARD = "standard"
     SIMILARITY = "similarity"
+    RANDOM = "random"
 
 
-def random_generate_query(
-    sobject: str, fields: T.List[str], num_records: float
+class SelectOperationExecutor:
+    def __init__(self, strategy: SelectStrategy):
+        self.strategy = strategy
+
+    def select_generate_query(
+        self, sobject: str, fields: T.List[str], num_records: int
+    ):
+        # For STANDARD strategy
+        if self.strategy == SelectStrategy.STANDARD:
+            return standard_generate_query(sobject=sobject, num_records=num_records)
+        # For SIMILARITY strategy
+        elif self.strategy == SelectStrategy.SIMILARITY:
+            return similarity_generate_query(sobject=sobject, fields=fields)
+        # For RANDOM strategy
+        elif self.strategy == SelectStrategy.RANDOM:
+            return standard_generate_query(sobject=sobject, num_records=num_records)
+
+    def select_post_process(
+        self, load_records, query_records: list, num_records: int, sobject: str
+    ):
+        # For STANDARD strategy
+        if self.strategy == SelectStrategy.STANDARD:
+            return standard_post_process(
+                query_records=query_records, num_records=num_records, sobject=sobject
+            )
+        # For SIMILARITY strategy
+        elif self.strategy == SelectStrategy.SIMILARITY:
+            return similarity_post_process(
+                load_records=load_records, query_records=query_records, sobject=sobject
+            )
+        # For RANDOM strategy
+        elif self.strategy == SelectStrategy.RANDOM:
+            return random_post_process(
+                query_records=query_records, num_records=num_records, sobject=sobject
+            )
+
+
+def standard_generate_query(
+    sobject: str, num_records: int
 ) -> T.Tuple[str, T.List[str]]:
-    """Generates the SOQL query for the random selection strategy"""
+    """Generates the SOQL query for the standard (as well as random) selection strategy"""
     # Get the WHERE clause from DEFAULT_DECLARATIONS if available
     declaration = DEFAULT_DECLARATIONS.get(sobject)
     if declaration:
@@ -32,10 +71,10 @@ def random_generate_query(
     return query, ["Id"]
 
 
-def random_post_process(
-    load_records, query_records: list, num_records: float, sobject: str
+def standard_post_process(
+    query_records: list, num_records: int, sobject: str
 ) -> T.Tuple[T.List[dict], T.Union[str, None]]:
-    """Processes the query results for the random selection strategy"""
+    """Processes the query results for the standard selection strategy"""
     # Handle case where query returns 0 records
     if not query_records:
         error_message = f"No records found for {sobject} in the target org."
@@ -59,9 +98,8 @@ def random_post_process(
 def similarity_generate_query(
     sobject: str,
     fields: T.List[str],
-    num_records: float,
 ) -> T.Tuple[str, T.List[str]]:
-    """Generates the SOQL query for the random selection strategy"""
+    """Generates the SOQL query for the similarity selection strategy"""
     # Get the WHERE clause from DEFAULT_DECLARATIONS if available
     declaration = DEFAULT_DECLARATIONS.get(sobject)
     if declaration:
@@ -81,7 +119,7 @@ def similarity_generate_query(
 
 
 def similarity_post_process(
-    load_records, query_records: list, num_records: float, sobject: str
+    load_records: list, query_records: list, sobject: str
 ) -> T.Tuple[T.List[dict], T.Union[str, None]]:
     """Processes the query results for the similarity selection strategy"""
     # Handle case where query returns 0 records
@@ -100,6 +138,26 @@ def similarity_post_process(
     return closest_records, None
 
 
+def random_post_process(
+    query_records: list, num_records: int, sobject: str
+) -> T.Tuple[T.List[dict], T.Union[str, None]]:
+    """Processes the query results for the random selection strategy"""
+
+    if not query_records:
+        error_message = f"No records found for {sobject} in the target org."
+        return [], error_message
+
+    selected_records = []
+    for _ in range(num_records):  # Loop 'num_records' times
+        # Randomly select one record from query_records
+        random_record = random.choice(query_records)
+        selected_records.append(
+            {"id": random_record[0], "success": True, "created": False}
+        )
+
+    return selected_records, None
+
+
 def find_closest_record(load_record: list, query_records: list):
     closest_distance = float("inf")
     closest_record = query_records[0]