def _execute_step(
self,
mapping: MappingStep) -> Union[DataOperationJobResult, MagicMock]:
"""Load data for a single step."""
if mapping.get("fields", {}).get("RecordTypeId"):
conn = self.session.connection()
self._load_record_types([mapping["sf_object"]], conn)
self.session.commit()
mapping["oid_as_pk"] = bool(mapping.get("fields", {}).get("Id"))
bulk_mode = mapping.get("bulk_mode") or self.bulk_mode or "Parallel"
step = BulkApiDmlOperation(
sobject=mapping["sf_object"],
operation=(DataOperationType.INSERT if mapping.get("action")
== "insert" else DataOperationType.UPDATE),
api_options={"bulk_mode": bulk_mode},
context=self,
fields=self._get_columns(mapping),
)
local_ids = []
step.start()
step.load_records(self._stream_queried_data(mapping, local_ids))
step.end()
if step.job_result.status is not DataOperationStatus.JOB_FAILURE:
self._process_job_results(mapping, step, local_ids)
return step.job_result
def test_end_to_end(self, download_mock):
context = mock.Mock()
context.bulk.endpoint = "https://test"
context.bulk.create_job.return_value = "JOB"
context.bulk.post_batch.side_effect = ["BATCH1", "BATCH2"]
download_mock.return_value = io.StringIO("""id,success,created,error
003000000000001,true,true,
003000000000002,true,true,
003000000000003,false,false,error""")
step = BulkApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={},
context=context,
fields=["LastName"],
)
step._wait_for_job = mock.Mock()
step._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
step.start()
step.load_records(iter([["Test"], ["Test2"], ["Test3"]]))
step.end()
assert step.job_result.status is DataOperationStatus.SUCCESS
results = step.get_results()
assert list(results) == [
DataOperationResult("003000000000001", True, None),
DataOperationResult("003000000000002", True, None),
DataOperationResult(None, False, "error"),
]
def _run_task(self):
self._validate_and_inject_namespace()
for obj in self.sobjects:
query = f"SELECT Id FROM {obj}"
if self.options["where"]:
query += f" WHERE {self.options['where']}"
self.logger.info(f"Querying for {obj} objects")
qs = BulkApiQueryOperation(sobject=obj,
api_options={},
context=self,
query=query)
qs.query()
if qs.job_result.status is not DataOperationStatus.SUCCESS:
raise BulkDataException(
f"Unable to query records for {obj}: {','.join(qs.job_result.job_errors)}"
)
if not qs.job_result.records_processed:
self.logger.info("No records found, skipping delete operation")
continue
self.logger.info(f"Deleting {self._object_description(obj)} ")
ds = BulkApiDmlOperation(
sobject=obj,
operation=(DataOperationType.HARD_DELETE
if self.options["hardDelete"] else
DataOperationType.DELETE),
api_options={},
context=self,
fields=["Id"],
)
ds.start()
ds.load_records(qs.get_results())
ds.end()
if ds.job_result.status not in [
DataOperationStatus.SUCCESS,
DataOperationStatus.ROW_FAILURE,
]:
raise BulkDataException(
f"Unable to delete records for {obj}: {','.join(qs.job_result.job_errors)}"
)
error_checker = RowErrorChecker(self.logger,
self.options["ignore_row_errors"],
self.row_warning_limit)
for result in ds.get_results():
error_checker.check_for_row_error(result, result.id)
