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")