def test_get_query_operation(self, rest_query, bulk_query):
context = mock.Mock()
context.sf.sf_version = "42.0"
op = get_query_operation(
sobject="Test",
fields=["Id"],
api_options={},
context=context,
query="SELECT Id FROM Test",
api=DataApi.BULK,
)
assert op == bulk_query.return_value
bulk_query.assert_called_once_with(
sobject="Test",
api_options={},
context=context,
query="SELECT Id FROM Test",
)
op = get_query_operation(
sobject="Test",
fields=["Id"],
api_options={},
context=context,
query="SELECT Id FROM Test",
api=DataApi.REST,
)
assert op == rest_query.return_value
rest_query.assert_called_once_with(
sobject="Test",
fields=["Id"],
api_options={},
context=context,
query="SELECT Id FROM Test",
)
def _run_query(self, soql, mapping):
"""Execute a Bulk or REST API query job and store the results."""
step = get_query_operation(
sobject=mapping.sf_object,
api=mapping.api,
fields=list(
mapping.get_complete_field_map(include_id=True).keys()),
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.logger.info("Downloading and importing records")
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 _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']}"
qs = get_query_operation(
sobject=obj,
fields=["Id"],
api_options={},
context=self,
query=query,
api=self.options["api"],
)
self.logger.info(f"Querying for {obj} objects")
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(
f"No records found, skipping delete operation for {obj}"
)
continue
self.logger.info(f"Deleting {self._object_description(obj)} ")
ds = get_dml_operation(
sobject=obj,
operation=(
DataOperationType.HARD_DELETE
if self.options["hardDelete"]
else DataOperationType.DELETE
),
fields=["Id"],
api_options={},
context=self,
api=self.options["api"],
volume=qs.job_result.records_processed,
)
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_get_query_operation__old_api(self, rest_query, bulk_query):
context = mock.Mock()
context.sf.sf_version = "39.0"
op = get_query_operation(
sobject="Test",
fields=["Id"],
api_options={},
context=context,
query="SELECT Id FROM Test",
api=DataApi.SMART,
)
assert op == bulk_query.return_value
context.sf.restful.assert_not_called()
def test_get_query_operation__smart_to_rest(self, rest_query, bulk_query):
context = mock.Mock()
context.sf.restful.return_value = {"sObjects": [{"name": "Test", "count": 1}]}
context.sf.sf_version = "42.0"
op = get_query_operation(
sobject="Test",
fields=["Id"],
api_options={},
context=context,
query="SELECT Id FROM Test",
api=DataApi.SMART,
)
assert op == rest_query.return_value
bulk_query.assert_not_called()
context.sf.restful.assert_called_once_with("limits/recordCount?sObjects=Test")