def test_query(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
query.query()
assert query.job_result.status is DataOperationStatus.SUCCESS
context.bulk.create_query_job.assert_called_once_with(
"Contact", contentType="CSV")
context.bulk.query.assert_called_once_with(
context.bulk.create_query_job.return_value,
"SELECT Id FROM Contact")
query._wait_for_job.assert_called_once_with(
context.bulk.create_query_job.return_value)
context.bulk.close_job.assert_called_once_with(
context.bulk.create_query_job.return_value)
def test_query__contextmanager(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
with query:
assert query.job_result.status is DataOperationStatus.SUCCESS
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)
def test_query__failure(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0)
query.query()
assert query.job_result.status is DataOperationStatus.JOB_FAILURE
def _run_query(self, soql, mapping):
"""Execute a Bulk API query job and store the results."""
step = BulkApiQueryOperation(sobject=mapping["sf_object"],
api_options={},
context=self,
query=soql)
self.logger.info(f"Extracting data for sObject {mapping['sf_object']}")
step.query()
if step.job_result.status is DataOperationStatus.SUCCESS:
if step.job_result.records_processed:
self._import_results(mapping, step)
else:
self.logger.info(
f"No records found for sObject {mapping['sf_object']}")
else:
raise BulkDataException(
f"Unable to execute query: {','.join(step.job_result.job_errors)}"
)
def test_get_results__no_results(self, download_mock):
context = mock.Mock()
context.bulk.endpoint = "https://test"
context.bulk.create_query_job.return_value = "JOB"
context.bulk.query.return_value = "BATCH"
context.bulk.get_query_batch_result_ids.return_value = ["RESULT"]
download_mock.return_value = io.StringIO(
"Records not found for this query")
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
query.query()
results = list(query.get_results())
context.bulk.get_query_batch_result_ids.assert_called_once_with(
"BATCH", job_id="JOB")
download_mock.assert_called_once_with(
"https://test/job/JOB/batch/BATCH/result/RESULT", context.bulk)
assert list(results) == []
