Repository: SFDO-Tooling/CumulusCI
Pull request #3718 (2023-12-11 16:36:43)

Rollback Operation for Inserts and Upserts
[W-14522337](https://gus.lightning.force.com/lightning/r/ADM_Work__c/a07EE00001et8KzYAI/view) and [W-14522327](https://gus.lightning.force.com/lightning/r/ADM_Work__c/a07EE00001etNspYAE/view)

The implementation is as follows:
1. For inserts, we maintain a separate table `{sobject}_insert_rollback` for records that have been created (either through insert or upsert) where we enter only the `sf_id`
2.  For upserts, before we load the records, we store the previous values of the records that will get updated (not upserted) in a table called `{sobject}_upsert_rollback`. It will have all the fields that are given in the `mapping.yml` file along with the `sf_id`.
3. When rollback occurs, it is as follows, in reverse order of tables created (and only considering the rollback tables):
    a. For insert, we delete all the `sf_id`
    b. For upsert, we upsert again the previous values



Patch:
diff --git a/cumulusci/tasks/bulkdata/load.py b/cumulusci/tasks/bulkdata/load.py
index 9a4f82ecf..9fa66c2a4 100644
--- a/cumulusci/tasks/bulkdata/load.py
+++ b/cumulusci/tasks/bulkdata/load.py
@@ -9,6 +9,7 @@ from sqlalchemy import Column, MetaData, Table, Unicode, create_engine, func, in
 from sqlalchemy.ext.automap import automap_base
 from sqlalchemy.orm import Session
 
+from cumulusci.core.enums import StrEnum
 from cumulusci.core.exceptions import BulkDataException, TaskOptionsError
 from cumulusci.core.utils import process_bool_arg
 from cumulusci.salesforce_api.org_schema import get_org_schema
@@ -28,9 +29,11 @@ from cumulusci.tasks.bulkdata.query_transformers import (
 )
 from cumulusci.tasks.bulkdata.step import (
     DEFAULT_BULK_BATCH_SIZE,
+    DataApi,
     DataOperationJobResult,
     DataOperationStatus,
     DataOperationType,
+    RestApiDmlOperation,
     get_dml_operation,
 )
 from cumulusci.tasks.bulkdata.upsert_utils import (
@@ -88,6 +91,9 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
         "org_shape_match_only": {
             "description": "When True, all path options are ignored and only a dataset matching the org shape name will be loaded. Defaults to False."
         },
+        "enable_rollback": {
+            "description": "When True, performs rollback operation incase of error. Defaults to False"
+        },
     }
     row_warning_limit = 10
 
@@ -115,6 +121,9 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
         self.options["set_recently_viewed"] = process_bool_arg(
             self.options.get("set_recently_viewed", True)
         )
+        self.options["enable_rollback"] = process_bool_arg(
+            self.options.get("enable_rollback", False)
+        )
 
     def _init_dataset(self):
         """Find the dataset paths to use with the following sequence:
@@ -261,13 +270,33 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
         step, query = self.configure_step(mapping)
 
         with tempfile.TemporaryFile(mode="w+t") as local_ids:
+            # Store the previous values of the records before upsert
+            # This is so that we can perform rollback
+            if (
+                mapping.action
+                in [
+                    DataOperationType.ETL_UPSERT,
+                    DataOperationType.UPSERT,
+                    DataOperationType.UPDATE,
+                ]
+                and self.options["enable_rollback"]
+            ):
+                UpdateRollback.prepare_for_rollback(
+                    self, step, self._stream_queried_data(mapping, local_ids, query)
+                )
             step.start()
             step.load_records(self._stream_queried_data(mapping, local_ids, query))
             step.end()
 
+            # Process Job Results
             if step.job_result.status is not DataOperationStatus.JOB_FAILURE:
                 local_ids.seek(0)
                 self._process_job_results(mapping, step, local_ids)
+            elif (
+                step.job_result.status is DataOperationStatus.JOB_FAILURE
+                and self.options["enable_rollback"]
+            ):
+                Rollback._perform_rollback(self)
 
             return step.job_result
 
@@ -454,7 +483,7 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
             id_table_name = self._initialize_id_table(mapping, self.reset_oids)
             conn = self.session.connection()
 
-        results_generator = self._generate_results_id_map(step, local_ids)
+        sf_id_results = self._generate_results_id_map(step, local_ids)
 
         # If we know we have no successful inserts, don't attempt to persist Ids.
         # Do, however, drain the generator to get error-checking behavior.
@@ -465,11 +494,8 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
                 connection=conn,
                 table=self.metadata.tables[id_table_name],
                 columns=("id", "sf_id"),
-                record_iterable=results_generator,
+                record_iterable=sf_id_results,
             )
-        else:
-            for r in results_generator:
-                pass  # Drain generator to validate results
 
         # Contact records for Person Accounts are inserted during an Account
         # sf_object step.  Insert records into the Contact ID table for
@@ -496,16 +522,37 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
 
     def _generate_results_id_map(self, step, local_ids):
         """Consume results from load and prepare rows for id table.
-        Raise BulkDataException on row errors if configured to do so."""
+        Raise BulkDataException on row errors if configured to do so.
+        Adds created records into insert_rollback Table
+        Performs rollback in case of any errors if enable_rollback is True"""
         error_checker = RowErrorChecker(
             self.logger, self.options["ignore_row_errors"], self.row_warning_limit
         )
         local_ids = (lid.strip("\n") for lid in local_ids)
+        sf_id_results = []
+        created_results = []
+        failed_results = []
         for result, local_id in zip(step.get_results(), local_ids):
             if result.success:
-                yield (local_id, result.id)
+                sf_id_results.append([local_id, result.id])
+                if result.created:
+                    created_results.append([result.id])
             else:
+                failed_results.append([result, local_id])
+
+        # We record failed_results separately since if a unsuccesful record
+        # was in between, it would not store all the successful ids
+        for result, local_id in failed_results:
+            try:
                 error_checker.check_for_row_error(result, local_id)
+            except Exception as e:
+                if self.options["enable_rollback"]:
+                    CreateRollback.prepare_for_rollback(self, step, created_results)
+                    Rollback._perform_rollback(self)
+                raise e
+        if self.options["enable_rollback"]:
+            CreateRollback.prepare_for_rollback(self, step, created_results)
+        return sf_id_results
 
     def _initialize_id_table(self, mapping, should_reset_table):
         """initalize or find table to hold the inserted SF Ids
@@ -568,6 +615,9 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
                 self.metadata.bind = connection
                 self.inspector = inspect(parent_engine)
 
+                # empty the record of initalized tables
+                Rollback._initialized_rollback_tables_api = {}
+
                 # initialize the automap mapping
                 self.base = automap_base(bind=connection, metadata=self.metadata)
                 self.base.prepare(connection, reflect=True)
@@ -810,6 +860,141 @@ class LoadData(SqlAlchemyMixin, BaseSalesforceApiTask):
         return results
 
 
+class RollbackType(StrEnum):
+    """Enum to specify type of rollback"""
+
+    UPSERT = "upsert_rollback"
+    INSERT = "insert_rollback"
+
+
+class Rollback:
+    # Store the table name and it's corresponding API (rest or bulk)
+    _initialized_rollback_tables_api = {}
+
+    @staticmethod
+    def _create_tables_for_rollback(context, step, rollback_type: RollbackType) -> str:
+        """Create the tables required for upsert and insert rollback"""
+        table_name = f"{step.sobject}_{rollback_type}"
+
+        if table_name not in Rollback._initialized_rollback_tables_api:
+            common_columns = [Column("Id", Unicode(255), primary_key=True)]
+
+            additional_columns = (
+                [Column(field, Unicode(255)) for field in step.fields if field != "Id"]
+                if rollback_type is RollbackType.UPSERT
+                else []
+            )
+
+            columns = common_columns + additional_columns
+
+            # Create the table
+            rollback_table = Table(table_name, context.metadata, *columns)
+            rollback_table.create()
+
+            # Store the API in the initialized tables dictionary
+            if isinstance(step, RestApiDmlOperation):
+                Rollback._initialized_rollback_tables_api[table_name] = DataApi.REST
+            else:
+                Rollback._initialized_rollback_tables_api[table_name] = DataApi.BULK
+
+        return table_name
+
+    @staticmethod
+    def _perform_rollback(context):
+        """Perform total rollback"""
+        context.logger.info("--Initiated Rollback Procedure--")
+        for table in reversed(context.metadata.sorted_tables):
+            if table.name.endswith(RollbackType.INSERT):
+                CreateRollback._perform_rollback(context, table)
+            elif table.name.endswith(RollbackType.UPSERT):
+                UpdateRollback._perform_rollback(context, table)
+        context.logger.info("--Finished Rollback Procedure--")
+
+
+class UpdateRollback:
+    @staticmethod
+    def prepare_for_rollback(context, step, records):
+        """Retrieve previous values for records being updated"""
+        results, columns = step.get_prev_record_values(records)
+        if results:
+            table_name = Rollback._create_tables_for_rollback(
+                context, step, RollbackType.UPSERT
+            )
+            conn = context.session.connection()
+            sql_bulk_insert_from_records(
+                connection=conn,
+                table=context.metadata.tables[table_name],
+                columns=columns,
+                record_iterable=results,
+            )
+
+    @staticmethod
+    def _perform_rollback(context, table: Table) -> None:
+        """Perform rollback for updated records"""
+        sf_object = table.name.split(f"_{RollbackType.UPSERT.value}")[0]
+        records = context.session.query(table).all()
+
+        if records:
+            context.logger.info(f"Reverting upserts for {sf_object}")
+            api_options = {"update_key": "Id"}
+
+            # Use get_dml_operation to create an UPSERT step
+            step = get_dml_operation(
+                sobject=sf_object,
+                operation=DataOperationType.UPSERT,
+                api_options=api_options,
+                context=context,
+                fields=[column.name for column in table.columns],
+                api=Rollback._initialized_rollback_tables_api[table.name],
+                volume=len(records),
+            )
+            step.start()
+            step.load_records(records)
+            step.end()
+            context.logger.info("Done")
+
+
+class CreateRollback:
+    @staticmethod
+    def prepare_for_rollback(context, step, records):
+        """Store the sf_ids of all records that were created
+        to prepare for rollback"""
+        if records:
+            table_name = Rollback._create_tables_for_rollback(
+                context, step, RollbackType.INSERT
+            )
+            conn = context.session.connection()
+            sql_bulk_insert_from_records(
+                connection=conn,
+                table=context.metadata.tables[table_name],
+                columns=["Id"],
+                record_iterable=records,
+            )
+
+    @staticmethod
+    def _perform_rollback(context, table: Table) -> None:
+        """Perform rollback for insert operation"""
+        sf_object = table.name.split(f"_{RollbackType.INSERT.value}")[0]
+        records = context.session.query(table).all()
+
+        if records:
+            context.logger.info(f"Deleting {sf_object} records")
+            # Perform DELETE operation using get_dml_operation
+            step = get_dml_operation(
+                sobject=sf_object,
+                operation=DataOperationType.DELETE,
+                fields=["Id"],
+                api_options={},
+                context=context,
+                api=Rollback._initialized_rollback_tables_api[table.name],
+                volume=len(records),
+            )
+            step.start()
+            step.load_records(records)
+            step.end()
+            context.logger.info("Done")
+
+
 class StepResultInfo(T.NamedTuple):
     """Represent a Step Result in a form easily convertible to JSON"""
 
diff --git a/cumulusci/tasks/bulkdata/step.py b/cumulusci/tasks/bulkdata/step.py
index a5da05301..edcb62afb 100644
--- a/cumulusci/tasks/bulkdata/step.py
+++ b/cumulusci/tasks/bulkdata/step.py
@@ -1,5 +1,6 @@
 import csv
 import io
+import json
 import os
 import pathlib
 import tempfile
@@ -9,6 +10,7 @@ from contextlib import contextmanager
 from typing import Any, Dict, List, NamedTuple, Optional
 
 import requests
+import salesforce_bulk
 
 from cumulusci.core.enums import StrEnum
 from cumulusci.core.exceptions import BulkDataException
@@ -58,6 +60,7 @@ class DataOperationResult(NamedTuple):
     id: str
     success: bool
     error: str
+    created: Optional[bool] = None
 
 
 class DataOperationJobResult(NamedTuple):
@@ -312,6 +315,11 @@ class BaseDmlOperation(BaseDataOperation, metaclass=ABCMeta):
         """Perform any required setup, such as job initialization, for the operation."""
         pass
 
+    @abstractmethod
+    def get_prev_record_values(self, records):
+        """Get the previous records values in case of UPSERT and UPDATE to prepare for rollback"""
+        pass
+
     @abstractmethod
     def load_records(self, records):
         """Perform the requested DML operation on the supplied row iterator."""
@@ -358,6 +366,56 @@ class BulkApiDmlOperation(BaseDmlOperation, BulkJobMixin):
         self.bulk.close_job(self.job_id)
         self.job_result = self._wait_for_job(self.job_id)
 
+    def get_prev_record_values(self, records):
+        """Get the previous values of the records based on the update key
+        to ensure rollback can be performed"""
+        # Function to be called only for UPSERT and UPDATE
+        assert self.operation in [DataOperationType.UPSERT, DataOperationType.UPDATE]
+
+        self.logger.info(f"Retrieving Previous Record Values of {self.sobject}")
+        prev_record_values = []
+        relevant_fields = set(self.fields + ["Id"])
+
+        # Set update key
+        update_key = (
+            self.api_options.get("update_key")
+            if self.operation == DataOperationType.UPSERT
+            else "Id"
+        )
+
+        for count, batch in enumerate(
+            self._batch(records, self.api_options["batch_size"])
+        ):
+            self.context.logger.info(f"Querying batch {count + 1}")
+
+            # Extract update key values from the batch
+            update_key_values = [
+                rec[update_key]
+                for rec in csv.DictReader([line.decode("utf-8") for line in batch])
+            ]
+
+            # Construct the SOQL query
+            query_fields = ", ".join(relevant_fields)
+            query_values = ", ".join(f"'{value}'" for value in update_key_values)
+            query = f"SELECT {query_fields} FROM {self.sobject} WHERE {update_key} IN ({query_values})"
+
+            # Execute the query using Bulk API
+            job_id = self.bulk.create_query_job(self.sobject, contentType="JSON")
+            batch_id = self.bulk.query(job_id, query)
+            self.bulk.wait_for_batch(job_id, batch_id)
+            self.bulk.close_job(job_id)
+            results = self.bulk.get_all_results_for_query_batch(batch_id)
+
+            # Extract relevant fields from results and append to the respective lists
+            for result in results:
+                result = json.load(salesforce_bulk.util.IteratorBytesIO(result))
+                prev_record_values.extend(
+                    [[res[key] for key in relevant_fields] for res in result]
+                )
+
+        self.logger.info("Done")
+        return prev_record_values, tuple(relevant_fields)
+
     def load_records(self, records):
         self.batch_ids = []
 
@@ -429,10 +487,12 @@ class BulkApiDmlOperation(BaseDmlOperation, BulkJobMixin):
 
                     for row in reader:
                         success = process_bool_arg(row[1])
+                        created = process_bool_arg(row[2])
                         yield DataOperationResult(
                             row[0] if success else None,
                             success,
                             row[3] if not success else None,
+                            created,
                         )
             except Exception as e:
                 raise BulkDataException(
@@ -489,6 +549,43 @@ class RestApiDmlOperation(BaseDmlOperation):
         result["attributes"] = {"type": self.sobject}
         return result
 
+    def get_prev_record_values(self, records):
+        """Get the previous values of the records based on the update key
+        to ensure rollback can be performed"""
+        # Function to be called only for UPSERT and UPDATE
+        assert self.operation in [DataOperationType.UPSERT, DataOperationType.UPDATE]
+
+        self.logger.info(f"Retrieving Previous Record Values of {self.sobject}")
+        prev_record_values = []
+        relevant_fields = set(self.fields + ["Id"])
+
+        # Set update key
+        update_key = (
+            self.api_options.get("update_key")
+            if self.operation == DataOperationType.UPSERT
+            else "Id"
+        )
+
+        for chunk in iterate_in_chunks(self.api_options.get("batch_size"), records):
+            update_key_values = tuple(
+                filter(None, (self._record_to_json(rec)[update_key] for rec in chunk))
+            )
+
+            # Construct the query string
+            query_fields = ", ".join(relevant_fields)
+            query = f"SELECT {query_fields} FROM {self.sobject} WHERE {update_key} IN {update_key_values}"
+
+            # Execute the query
+            results = self.sf.query(query)
+
+            # Extract relevant fields from results and extend the list
+            prev_record_values.extend(
+                [[res[key] for key in relevant_fields] for res in results["records"]]
+            )
+
+        self.logger.info("Done")
+        return prev_record_values, tuple(relevant_fields)
+
     def load_records(self, records):
         """Load, update, upsert or delete records into the org"""
 
@@ -547,7 +644,14 @@ class RestApiDmlOperation(BaseDmlOperation):
             else:
                 errors = ""
 
-            return DataOperationResult(res.get("id"), res["success"], errors)
+            if self.operation == DataOperationType.INSERT:
+                created = True
+            elif self.operation == DataOperationType.UPDATE:
+                created = False
+            else:
+                created = res.get("created")
+
+            return DataOperationResult(res.get("id"), res["success"], errors, created)
 
         yield from (_convert(res) for res in self.results)
 


Test patch:
diff --git a/cumulusci/tasks/bulkdata/tests/test_load.py b/cumulusci/tasks/bulkdata/tests/test_load.py
index c6e85fdcb..a4293910d 100644
--- a/cumulusci/tasks/bulkdata/tests/test_load.py
+++ b/cumulusci/tasks/bulkdata/tests/test_load.py
@@ -6,6 +6,7 @@ import random
 import shutil
 import string
 import tempfile
+from collections import namedtuple
 from contextlib import nullcontext
 from datetime import date, timedelta
 from pathlib import Path
@@ -18,6 +19,12 @@ from sqlalchemy import Column, Table, Unicode, create_engine
 from cumulusci.core.exceptions import BulkDataException, TaskOptionsError
 from cumulusci.salesforce_api.org_schema import get_org_schema
 from cumulusci.tasks.bulkdata import LoadData
+from cumulusci.tasks.bulkdata.load import (
+    CreateRollback,
+    Rollback,
+    RollbackType,
+    UpdateRollback,
+)
 from cumulusci.tasks.bulkdata.mapping_parser import MappingLookup, MappingStep
 from cumulusci.tasks.bulkdata.step import (
     BulkApiDmlOperation,
@@ -125,6 +132,132 @@ class TestLoadData:
                 hh_ids = next(c.execute("SELECT * from households_sf_ids"))
                 assert hh_ids == ("1", "001000000000000")
 
+    @responses.activate
+    @mock.patch("cumulusci.tasks.bulkdata.load.get_dml_operation")
+    def test__insert_rollback(self, dml_mock):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "start_step": "Insert Contacts",
+                    "set_recently_viewed": False,
+                    "enable_rollback": True,
+                }
+            },
+        )
+        table = mock.Mock()
+        p = mock.PropertyMock(return_value=f"Contact_{RollbackType.INSERT}")
+        type(table).name = p
+        task._initialized_rollback_tables_api = {
+            f"Contact_{RollbackType.INSERT}": "rest"
+        }
+        task.metadata = mock.Mock(sorted_tables=[table])
+        task.session = mock.Mock(
+            query=mock.Mock(
+                return_value=mock.Mock(
+                    all=mock.Mock(return_value=[mock.Mock(sf_id="00001111")]),
+                )
+            )
+        )
+
+        Rollback._initialized_rollback_tables_api = {"Contact_insert_rollback": "rest"}
+        CreateRollback._perform_rollback(task, table)
+
+        dml_mock.assert_called_once_with(
+            sobject="Contact",
+            operation=(DataOperationType.DELETE),
+            fields=["Id"],
+            api_options={},
+            context=task,
+            api="rest",
+            volume=1,
+        )
+        dml_mock.return_value.start.assert_called_once()
+        dml_mock.return_value.end.assert_called_once()
+
+    @responses.activate
+    @mock.patch("cumulusci.tasks.bulkdata.load.get_dml_operation")
+    def test__upsert_rollback(self, dml_mock):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "start_step": "Upsert Contacts",
+                    "set_recently_viewed": False,
+                    "enable_rollback": True,
+                }
+            },
+        )
+        table = mock.Mock()
+        p = mock.PropertyMock(return_value=f"Contact_{RollbackType.UPSERT}")
+        type(table).name = p
+        Column = namedtuple("Column", ["name"])
+        type(table).columns = [Column("Id"), Column("LastName")]
+        task._initialized_rollback_tables_api = {
+            f"Contact_{RollbackType.UPSERT}": "rest"
+        }
+        task.metadata = mock.Mock(sorted_tables=[table])
+        task.session = mock.Mock(
+            query=mock.Mock(
+                return_value=mock.Mock(
+                    all=mock.Mock(
+                        return_value=[mock.Mock(Id="00001111", LastName="TestName")]
+                    ),
+                )
+            )
+        )
+
+        Rollback._initialized_rollback_tables_api = {"Contact_upsert_rollback": "rest"}
+        UpdateRollback._perform_rollback(task, table)
+
+        dml_mock.assert_called_once_with(
+            sobject="Contact",
+            operation=(DataOperationType.UPSERT),
+            fields=["Id", "LastName"],
+            api_options={"update_key": "Id"},
+            context=task,
+            api="rest",
+            volume=1,
+        )
+        dml_mock.return_value.start.assert_called_once()
+        dml_mock.return_value.end.assert_called_once()
+
+    def test__perform_rollback(self):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "start_step": "Insert Contacts",
+                    "set_recently_viewed": False,
+                    "enable_rollback": True,
+                }
+            },
+        )
+        table_insert = mock.Mock()
+        p = mock.PropertyMock(return_value=f"Contact_{RollbackType.INSERT}")
+        type(table_insert).name = p
+        table_upsert = mock.Mock()
+        p = mock.PropertyMock(return_value=f"Account_{RollbackType.UPSERT}")
+        type(table_upsert).name = p
+        task.metadata = mock.Mock()
+        task.metadata.sorted_tables = [table_insert, table_upsert]
+
+        with mock.patch.object(
+            CreateRollback, "_perform_rollback"
+        ) as mock_insert_rollback, mock.patch.object(
+            UpdateRollback, "_perform_rollback"
+        ) as mock_upsert_rollback:
+            Rollback._perform_rollback(task)
+
+            mock_insert_rollback.assert_called_once_with(task, table_insert)
+            mock_upsert_rollback.assert_called_once_with(task, table_upsert)
+
     def test_run_task__start_step(self):
         task = _make_task(
             LoadData,
@@ -1099,6 +1232,7 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
         )
 
         task.session = mock.Mock()
+        task.metadata = mock.MagicMock()
         task._initialize_id_table = mock.Mock()
         task.bulk = mock.Mock()
         task.sf = mock.Mock()
@@ -1176,6 +1310,7 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
         )
 
         task.session = mock.Mock()
+        task.metadata = mock.MagicMock()
         task._initialize_id_table = mock.Mock()
         task.bulk = mock.Mock()
         task.sf = mock.Mock()
@@ -1529,37 +1664,46 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
             {"options": {"database_url": "sqlite://", "mapping": "mapping.yml"}},
         )
 
+        task.session = mock.MagicMock()
+        task.metadata = mock.MagicMock()
         step = mock.Mock()
         step.get_results.return_value = iter(
             [
-                DataOperationResult("001000000000000", True, None),
-                DataOperationResult("001000000000001", True, None),
-                DataOperationResult("001000000000002", True, None),
+                DataOperationResult("001000000000000", True, None, True),
+                DataOperationResult("001000000000001", True, None, True),
+                DataOperationResult("001000000000002", True, None, True),
             ]
         )
 
-        generator = task._generate_results_id_map(
+        sf_id_list = task._generate_results_id_map(
             step, ["001000000000009", "001000000000010", "001000000000011"]
         )
 
-        assert list(generator) == [
-            ("001000000000009", "001000000000000"),
-            ("001000000000010", "001000000000001"),
-            ("001000000000011", "001000000000002"),
+        assert sf_id_list == [
+            ["001000000000009", "001000000000000"],
+            ["001000000000010", "001000000000001"],
+            ["001000000000011", "001000000000002"],
         ]
 
-    def test_generate_results_id_map__exception_failure(self):
+    def test_generate_results_id_map__exception_failure_without_rollback(self):
         task = _make_task(
             LoadData,
-            {"options": {"database_url": "sqlite://", "mapping": "mapping.yml"}},
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "enable_rollback": False,
+                }
+            },
         )
 
+        task.metadata = mock.MagicMock()
         step = mock.Mock()
         step.get_results.return_value = iter(
             [
-                DataOperationResult("001000000000000", True, None),
-                DataOperationResult(None, False, "error"),
-                DataOperationResult("001000000000002", True, None),
+                DataOperationResult("001000000000000", True, None, True),
+                DataOperationResult(None, False, "error", False),
+                DataOperationResult("001000000000002", True, None, True),
             ]
         )
 
@@ -1573,6 +1717,43 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
         assert "Error on record" in str(e.value)
         assert "001000000000010" in str(e.value)
 
+    def test_generate_results_id_map__exception_failure_with_rollback(self):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "enable_rollback": True,
+                }
+            },
+        )
+
+        task.metadata = mock.MagicMock()
+        task.session = mock.MagicMock()
+        step = mock.Mock()
+        step.get_results.return_value = iter(
+            [
+                DataOperationResult("001000000000000", True, None, True),
+                DataOperationResult(None, False, "error", False),
+                DataOperationResult("001000000000002", True, None, True),
+            ]
+        )
+
+        with pytest.raises(BulkDataException) as e, mock.patch(
+            "cumulusci.tasks.bulkdata.load.Rollback._perform_rollback"
+        ) as mock_rollback, mock.patch(
+            "cumulusci.tasks.bulkdata.load.sql_bulk_insert_from_records"
+        ) as mock_insert_records:
+            task._generate_results_id_map(
+                step, ["001000000000009", "001000000000010", "001000000000011"]
+            )
+
+        mock_rollback.assert_called_once()
+        mock_insert_records.assert_called_once()
+        assert "Error on record" in str(e.value)
+        assert "001000000000010" in str(e.value)
+
     def test_generate_results_id_map__respects_silent_error_flag(self):
         task = _make_task(
             LoadData,
@@ -1584,17 +1765,17 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
                 }
             },
         )
-
+        task.metadata = mock.MagicMock()
         step = mock.Mock()
         step.get_results.return_value = iter(
             [DataOperationResult(None, False, None)] * 15
         )
 
         with mock.patch.object(task.logger, "warning") as warning:
-            generator = task._generate_results_id_map(
+            sf_id_list = task._generate_results_id_map(
                 step, ["001000000000009", "001000000000010", "001000000000011"] * 15
             )
-            _ = list(generator)  # generate the errors
+            _ = sf_id_list  # generate the errors
 
         assert len(warning.mock_calls) == task.row_warning_limit + 1 == 11
         assert "warnings suppressed" in str(warning.mock_calls[-1])
@@ -1608,15 +1789,106 @@ FROM accounts LEFT OUTER JOIN accounts_sf_ids AS accounts_sf_ids_1 ON accounts_s
             ]
         )
 
-        generator = task._generate_results_id_map(
+        sf_id_list = task._generate_results_id_map(
             step, ["001000000000009", "001000000000010", "001000000000011"]
         )
 
-        assert list(generator) == [
-            ("001000000000009", "001000000000000"),
-            ("001000000000011", "001000000000002"),
+        assert sf_id_list == [
+            ["001000000000009", "001000000000000"],
+            ["001000000000011", "001000000000002"],
         ]
 
+    @mock.patch("cumulusci.tasks.bulkdata.load.get_dml_operation")
+    def test__execute_step__prev_record_values(self, mock_dml):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "enable_rollback": True,
+                }
+            },
+        )
+
+        ret_prev_records = [["TestName1", "Id1"], ["TestName2", "Id2"]]
+        ret_columns = ("Name", "Id")
+        conn = mock.Mock()
+        task.session = mock.Mock()
+        task.session.connection.return_value = conn
+        task.metadata = mock.MagicMock()
+        tables = {f"Account_{RollbackType.UPSERT}": "AccountUpsertTable"}
+        task.metadata.tables = tables
+        step = mock.Mock()
+        step.fields = ["Name"]
+        step.sobject = "Account"
+        query = mock.Mock()
+        task.configure_step = mock.Mock()
+        task.configure_step.return_value = (step, query)
+        step.get_prev_record_values.return_value = (ret_prev_records, ret_columns)
+        task._load_record_types = mock.Mock()
+        task._process_job_results = mock.Mock()
+        task._query_db = mock.Mock()
+        with mock.patch(
+            "cumulusci.tasks.bulkdata.load.sql_bulk_insert_from_records"
+        ) as mock_insert_records:
+            task._execute_step(
+                MappingStep(
+                    **{
+                        "sf_object": "Account",
+                        "fields": {"Name": "Name"},
+                        "action": DataOperationType.UPSERT,
+                        "api": "rest",
+                        "update_key": ["Name"],
+                    }
+                )
+            )
+
+            mock_insert_records.assert_called_once_with(
+                connection=conn,
+                table="AccountUpsertTable",
+                columns=ret_columns,
+                record_iterable=ret_prev_records,
+            )
+
+    @mock.patch("cumulusci.tasks.bulkdata.load.get_dml_operation")
+    def test__execute_step__job_failure_rollback(self, mock_dml):
+        task = _make_task(
+            LoadData,
+            {
+                "options": {
+                    "database_url": "sqlite://",
+                    "mapping": "mapping.yml",
+                    "enable_rollback": True,
+                }
+            },
+        )
+
+        task.session = mock.Mock()
+        task.metadata = mock.MagicMock()
+        step = mock.Mock()
+        query = mock.Mock()
+        step.job_result.status = DataOperationStatus.JOB_FAILURE
+        task.configure_step = mock.Mock()
+        task.configure_step.return_value = (step, query)
+        task._load_record_types = mock.Mock()
+        task._process_job_results = mock.Mock()
+        task._query_db = mock.Mock()
+        with mock.patch(
+            "cumulusci.tasks.bulkdata.load.Rollback._perform_rollback"
+        ) as mock_rollback:
+            task._execute_step(
+                MappingStep(
+                    **{
+                        "sf_object": "Account",
+                        "action": DataOperationType.INSERT,
+                        "fields": {"Name": "Name"},
+                        "api": "rest",
+                    }
+                )
+            )
+            mock_rollback.assert_called()
+
     @mock.patch("cumulusci.tasks.bulkdata.load.get_dml_operation")
     def test_execute_step__record_type_mapping(self, dml_mock):
         task = _make_task(
@@ -2614,7 +2886,6 @@ class TestLoadDataIntegrationTests:
             {
                 "sql_path": cumulusci_test_repo_root / "datasets/bad_sample.sql",
                 "mapping": cumulusci_test_repo_root / "datasets/mapping.yml",
-                "ignore_row_errors": True,
             },
         )
         with mock.patch("cumulusci.tasks.bulkdata.step.DEFAULT_BULK_BATCH_SIZE", 3):
@@ -2681,7 +2952,6 @@ class TestLoadDataIntegrationTests:
                 {
                     "sql_path": sql_path,
                     "mapping": mapping_path,
-                    "ignore_row_errors": True,
                 },
             )
             task()
@@ -2699,7 +2969,6 @@ class TestLoadDataIntegrationTests:
             {
                 "sql_path": cumulusci_test_repo_root / "datasets/sample.sql",
                 "mapping": cumulusci_test_repo_root / "datasets/mapping.yml",
-                "ignore_row_errors": True,
             },
         )
         task.logger = mock.Mock()
diff --git a/cumulusci/tasks/bulkdata/tests/test_step.py b/cumulusci/tasks/bulkdata/tests/test_step.py
index 786c61c81..fc8cea701 100644
--- a/cumulusci/tasks/bulkdata/tests/test_step.py
+++ b/cumulusci/tasks/bulkdata/tests/test_step.py
@@ -498,6 +498,42 @@ class TestBulkApiDmlOperation:
         serialized = step._serialize_csv_record(record)
         assert serialized == b'"col1","multiline\ncol2"\r\n'
 
+    def test_get_prev_record_values(self):
+        context = mock.Mock()
+        step = BulkApiDmlOperation(
+            sobject="Contact",
+            operation=DataOperationType.UPSERT,
+            api_options={"batch_size": 10, "update_key": "LastName"},
+            context=context,
+            fields=["LastName"],
+        )
+        results = [
+            [{"LastName": "Test1", "Id": "Id1"}, {"LastName": "Test2", "Id": "Id2"}]
+        ]
+        expected_record_values = [["Test1", "Id1"], ["Test2", "Id2"]]
+        expected_relevant_fields = ("Id", "LastName")
+        step.bulk.create_query_job = mock.Mock()
+        step.bulk.create_query_job.return_value = "JOB_ID"
+        step.bulk.query = mock.Mock()
+        step.bulk.query.return_value = "BATCH_ID"
+        step.bulk.get_all_results_for_query_batch = mock.Mock()
+        step.bulk.get_all_results_for_query_batch.return_value = results
+
+        records = iter([["Test1"], ["Test2"], ["Test3"]])
+        with mock.patch("json.load", side_effect=lambda result: result), mock.patch(
+            "salesforce_bulk.util.IteratorBytesIO", side_effect=lambda result: result
+        ):
+            prev_record_values, relevant_fields = step.get_prev_record_values(records)
+
+        assert sorted(map(sorted, prev_record_values)) == sorted(
+            map(sorted, expected_record_values)
+        )
+        assert set(relevant_fields) == set(expected_relevant_fields)
+        step.bulk.create_query_job.assert_called_once_with(
+            "Contact", contentType="JSON"
+        )
+        step.bulk.get_all_results_for_query_batch.assert_called_once_with("BATCH_ID")
+
     def test_batch(self):
         context = mock.Mock()
 
@@ -586,9 +622,9 @@ class TestBulkApiDmlOperation:
         results = step.get_results()
 
         assert list(results) == [
-            DataOperationResult("003000000000001", True, None),
-            DataOperationResult("003000000000002", True, None),
-            DataOperationResult(None, False, "error"),
+            DataOperationResult("003000000000001", True, None, True),
+            DataOperationResult("003000000000002", True, None, True),
+            DataOperationResult(None, False, "error", False),
         ]
         download_mock.assert_has_calls(
             [
@@ -649,9 +685,9 @@ class TestBulkApiDmlOperation:
         results = step.get_results()
 
         assert list(results) == [
-            DataOperationResult("003000000000001", True, None),
-            DataOperationResult("003000000000002", True, None),
-            DataOperationResult(None, False, "error"),
+            DataOperationResult("003000000000001", True, None, True),
+            DataOperationResult("003000000000002", True, None, True),
+            DataOperationResult(None, False, "error", False),
         ]
 
 
@@ -781,11 +817,68 @@ class TestRestApiDmlOperation:
             DataOperationStatus.SUCCESS, [], 3, 0
         )
         assert list(dml_op.get_results()) == [
-            DataOperationResult("003000000000001", True, ""),
-            DataOperationResult("003000000000002", True, ""),
-            DataOperationResult("003000000000003", True, ""),
+            DataOperationResult("003000000000001", True, "", True),
+            DataOperationResult("003000000000002", True, "", True),
+            DataOperationResult("003000000000003", True, "", True),
         ]
 
+    @responses.activate
+    def test_get_prev_record_values(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"],
+        )
+
+        results = {
+            "records": [
+                {"LastName": "Test1", "Id": "Id1"},
+                {"LastName": "Test2", "Id": "Id2"},
+            ]
+        }
+        expected_record_values = [["Test1", "Id1"], ["Test2", "Id2"]]
+        expected_relevant_fields = ("Id", "LastName")
+        step.sf.query = mock.Mock()
+        step.sf.query.return_value = results
+        records = iter([["Test1"], ["Test2"], ["Test3"]])
+        prev_record_values, relevant_fields = step.get_prev_record_values(records)
+
+        assert sorted(map(sorted, prev_record_values)) == sorted(
+            map(sorted, expected_record_values)
+        )
+        assert set(relevant_fields) == set(expected_relevant_fields)
+
     @responses.activate
     def test_insert_dml_operation__boolean_conversion(self):
         mock_describe_calls()
@@ -1013,10 +1106,10 @@ class TestRestApiDmlOperation:
             DataOperationStatus.ROW_FAILURE, [], 3, 1
         )
         assert list(dml_op.get_results()) == [
-            DataOperationResult("003000000000001", True, ""),
-            DataOperationResult("003000000000002", True, ""),
+            DataOperationResult("003000000000001", True, "", True),
+            DataOperationResult("003000000000002", True, "", True),
             DataOperationResult(
-                "003000000000003", False, "VALIDATION_ERR: Bad data (FirstName)"
+                "003000000000003", False, "VALIDATION_ERR: Bad data (FirstName)", True
             ),
         ]
 
diff --git a/cumulusci/tasks/bulkdata/tests/test_upsert.py b/cumulusci/tasks/bulkdata/tests/test_upsert.py
index f9bf0a937..aa23c50fc 100644
--- a/cumulusci/tasks/bulkdata/tests/test_upsert.py
+++ b/cumulusci/tasks/bulkdata/tests/test_upsert.py
@@ -70,6 +70,7 @@ class TestUpsert:
                 / f"datasets/upsert/upsert_mapping_{api}.yml",
                 # "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
 
@@ -128,6 +129,7 @@ class TestUpsert:
                 / f"datasets/upsert/upsert_mapping_{api}.yml",
                 # "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task()
@@ -196,6 +198,7 @@ class TestUpsert:
                 / "datasets/upsert/upsert_mapping_rest.yml",
                 "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task._update_credentials = mock.Mock()
@@ -385,6 +388,7 @@ class TestUpsert:
                 / "datasets/upsert/upsert_mapping_bulk.yml",
                 "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task._update_credentials = mock.Mock()
@@ -437,6 +441,7 @@ class TestUpsert:
                 "mapping": cumulusci_test_repo_root
                 / f"datasets/upsert/upsert_mapping_{api}_complex.yml",
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task()
@@ -457,6 +462,7 @@ class TestUpsert:
                 "mapping": cumulusci_test_repo_root
                 / f"datasets/upsert/upsert_mapping_{api}_complex.yml",
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task()
@@ -528,6 +534,7 @@ class TestUpsert:
                 "mapping": cumulusci_test_repo_root
                 / "datasets/upsert/upsert_mapping_rest_complex.yml",
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task()
@@ -656,6 +663,7 @@ class TestUpsert:
                 / "datasets/upsert/upsert_smart__native_field.yml",
                 "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task._update_credentials = mock.Mock()
@@ -722,6 +730,7 @@ class TestUpsert:
                 / "datasets/upsert/upsert_smart__native_field.yml",
                 "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task._update_credentials = mock.Mock()
@@ -893,6 +902,7 @@ class TestUpsert:
                 / "datasets/upsert/upsert_smart__non_native_field.yml",
                 "ignore_row_errors": True,
                 "set_recently_viewed": False,
+                "enable_rollback": False,
             },
         )
         task._update_credentials = mock.Mock()
diff --git a/cumulusci/tasks/bulkdata/tests/utils.py b/cumulusci/tasks/bulkdata/tests/utils.py
index 164e26a5e..173f4c612 100644
--- a/cumulusci/tasks/bulkdata/tests/utils.py
+++ b/cumulusci/tasks/bulkdata/tests/utils.py
@@ -92,6 +92,9 @@ class FakeBulkAPIDmlOperation(BaseDmlOperation):
             DataOperationStatus.SUCCESS, [], records_processed, 0
         )
 
+    def get_prev_record_values(self, records):
+        pass
+
     def load_records(self, records):
         self.records.extend(records)