package com.celnet.datadump.util; import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.sforce.async.*; import com.sforce.soap.partner.PartnerConnection; import com.sforce.ws.ConnectionException; import com.sforce.ws.ConnectorConfig; import com.sforce.ws.bind.CalendarCodec; import org.apache.http.Header; import org.apache.http.HttpEntity; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpGet; import org.apache.http.client.methods.HttpPost; import org.apache.http.entity.StringEntity; import org.apache.http.impl.client.CloseableHttpClient; import org.apache.http.impl.client.HttpClients; import org.apache.http.util.EntityUtils; import java.io.*; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Paths; import java.util.*; import static com.sforce.async.BulkConnection.JSON_CONTENT_TYPE; public class BulkUtil { private static final String URI_STEM_QUERY = "query/"; private static final String URI_STEM_INGEST = "ingest/"; private static final String AUTH_HEADER = "Authorization"; private static final String REQUEST_CONTENT_TYPE_HEADER = "Content-Type"; private static final String ACCEPT_CONTENT_TYPES_HEADER = "ACCEPT"; private static final String AUTH_HEADER_VALUE_PREFIX = "Bearer "; private static final String UTF_8 = StandardCharsets.UTF_8.name(); private static final String INGEST_RESULTS_SUCCESSFUL = "successfulResults"; private static final String INGEST_RESULTS_UNSUCCESSFUL = "failedResults"; private static final String INGEST_RECORDS_UNPROCESSED = "unprocessedrecords"; public static final String SFORCE_CALL_OPTIONS_HEADER = "Sforce-Call-Options"; public static void closeJob(BulkConnection connection, String jobId) throws AsyncApiException { JobInfo job = new JobInfo(); job.setId(jobId); job.setState(JobStateEnum.Closed); connection.updateJob(job); } public static int awaitCompletionOne(BulkConnection connection, JobInfo job, BatchInfo bi) throws AsyncApiException { long sleepTime = 0L; int numberRecordsProcessed = 0; Set incomplete = new HashSet(); incomplete.add(bi.getId()); while (!incomplete.isEmpty()) { try { Thread.sleep(sleepTime); } catch (InterruptedException e) {} System.out.println("Awaiting results..." + incomplete.size()); sleepTime = 10000L; BatchInfo[] statusList = connection.getBatchInfoList(job.getId()).getBatchInfo(); for (BatchInfo b : statusList) { if (b.getState() == BatchStateEnum.Completed || b.getState() == BatchStateEnum.Failed) { if (incomplete.remove(b.getId())) { System.out.println("BATCH STATUS:\n" + b); } } numberRecordsProcessed += b.getNumberRecordsProcessed(); } } return numberRecordsProcessed; } /** * Wait for a job to complete by polling the Bulk API. * * @param connection * BulkConnection used to check results. * @param job * The job awaiting completion. * List of batches for this job. * @throws AsyncApiException */ public static void awaitCompletion(BulkConnection connection, JobInfo job, List batchInfoList) throws AsyncApiException { long sleepTime = 0L; Set incomplete = new HashSet(); for (BatchInfo bi : batchInfoList) { incomplete.add(bi.getId()); } while (!incomplete.isEmpty()) { try { Thread.sleep(sleepTime); } catch (InterruptedException e) {} System.out.println("Awaiting results..." + incomplete.size()); sleepTime = 10000L; BatchInfo[] statusList = connection.getBatchInfoList(job.getId()).getBatchInfo(); for (BatchInfo b : statusList) { if (b.getState() == BatchStateEnum.Completed || b.getState() == BatchStateEnum.Failed) { if (incomplete.remove(b.getId())) { System.out.println("BATCH STATUS:\n" + b); } } } } } /** * Wait for a job to complete by polling the Bulk API. * * @param connection * BulkConnection used to check results. * @param job * The job awaiting completion. * List of batches for this job. * @throws AsyncApiException */ public static void awaitOneCompletion(BulkConnection connection, JobInfo job, BatchInfo bi) throws AsyncApiException { long sleepTime = 0L; Set incomplete = new HashSet(); incomplete.add(bi.getId()); while (!incomplete.isEmpty()) { try { Thread.sleep(sleepTime); } catch (InterruptedException e) {} System.out.println("Awaiting results..." + incomplete.size()); sleepTime = 10000L; BatchInfo[] statusList = connection.getBatchInfoList(job.getId()).getBatchInfo(); for (BatchInfo b : statusList) { if (b.getState() == BatchStateEnum.Completed || b.getState() == BatchStateEnum.Failed) { if (incomplete.remove(b.getId())) { System.out.println("BATCH STATUS:\n" + b); } } } } } /** * Create a new job using the Bulk API. * * @param sobjectType * The object type being loaded, such as "Account" * @param connection * BulkConnection used to create the new job. * @return The JobInfo for the new job. * @throws AsyncApiException */ public static JobInfo createJob( BulkConnection connection,String sobjectType,OperationEnum operation) throws AsyncApiException { JobInfo job = new JobInfo(); job.setObject(sobjectType); job.setOperation(operation); job.setContentType(ContentType.CSV); job = connection.createJob(job); return job; } /** * 创建Bulk API 2.0查询作业 ** @param soqlQuery SOQL查询语句 * @return 作业ID * @throws Exception 创建作业时发生错误 */ public static String createBulkV2Job(String url,String soqlQuery,BulkConnection connection) throws Exception { // 构建创建查询作业的 JSON 请求体 String jsonRequest = String.format("{\"operation\":\"query\", \"query\":\"%s\"}", soqlQuery.replace("\"", "\\\"")); CloseableHttpClient httpClient = HttpClients.createDefault(); HttpPost httpPost = new HttpPost(url + "services/data/v56.0/jobs/query"); httpPost.setHeader("Content-Type", "application/json"); httpPost.setHeader("Authorization", "Bearer " + connection.getConfig().getSessionId()); httpPost.setEntity(new StringEntity(jsonRequest, "UTF-8")); CloseableHttpResponse response = httpClient.execute(httpPost); try { int statusCode = response.getStatusLine().getStatusCode(); if (statusCode == 200) { // 从响应体中解析出 jobId String responseBody = EntityUtils.toString(response.getEntity()); // 简单的JSON解析提取id，实际项目中建议使用JSON库如Jackson String jobId = responseBody.split("\"id\":\"")[1].split("\"")[0]; System.out.println("Bulk API 2.0 Query job created with ID: " + jobId); return jobId; } else { String errorBody = EntityUtils.toString(response.getEntity()); throw new RuntimeException("Failed to create Bulk API 2.0 query job, status code: " + statusCode + ", error: " + errorBody); } } finally { response.close(); httpClient.close(); } } /** * Wait for a Bulk API 2.0 job to complete by polling the job status. * * @param jobId The ID of the job to check. * @param connection BulkConnection used to get session info. * @throws Exception If an error occurs during the polling process. */ public static void waitForBulkV2JobCompletion(String url, String jobId, BulkConnection connection) throws Exception { CloseableHttpClient httpClient = HttpClients.createDefault(); String jobStatusUrl = url + "services/data/v56.0/jobs/query/" + jobId; boolean jobCompleted = false; while (!jobCompleted) { HttpGet httpGet = new HttpGet(jobStatusUrl); httpGet.setHeader("Authorization", "Bearer " + connection.getConfig().getSessionId()); CloseableHttpResponse response = httpClient.execute(httpGet); try { int statusCode = response.getStatusLine().getStatusCode(); if (statusCode == 200) { String responseBody = EntityUtils.toString(response.getEntity()); // 简单解析响应体中的状态字段，实际项目中建议使用JSON库如Jackson if (responseBody.contains("\"state\":\"JobComplete\"")) { System.out.println("Job completed successfully."); jobCompleted = true; } else if (responseBody.contains("\"state\":\"Failed\"") || responseBody.contains("\"state\":\"Aborted\"")) { throw new RuntimeException("Job failed or was aborted. Response: " + responseBody); } else { System.out.println("Job is still running. Waiting..."); Thread.sleep(10000); // 等待10秒后再次检查 } } else { String errorBody = EntityUtils.toString(response.getEntity()); throw new RuntimeException("Failed to get job status, status code: " + statusCode + ", error: " + errorBody); } } finally { response.close(); } } httpClient.close(); } /** * 获取 Bulk API 2.0 查询作业的结果 * * @param url Salesforce实例URL * @param jobId 作业ID * @param connection BulkConnection used to get session info. * @return 查询结果的字符串内容 * @throws Exception 如果在获取结果过程中发生错误 */ public static String getBulkV2QueryJobResults(String url, String jobId, BulkConnection connection) throws Exception { CloseableHttpClient httpClient = HttpClients.createDefault(); String resultsUrl = url + "/services/data/v65.0/jobs/query/" + jobId + "/results" ; HttpPost httpGet = new HttpPost(resultsUrl); httpGet.setHeader("Authorization", "Bearer " + connection.getConfig().getSessionId()); httpGet.setHeader("Content-Type", "application/json"); httpGet.setHeader("Accept", "text/csv"); httpGet.setHeader("X-PrettyPrint", "1"); CloseableHttpResponse response = httpClient.execute(httpGet); try { int statusCode = response.getStatusLine().getStatusCode(); if (statusCode == 200) { Header locatorHeader = response.getFirstHeader("Sforce-Locator"); Header numRecordsHeader = response.getFirstHeader("Sforce-NumberOfRecords"); String newLocator = (locatorHeader != null) ? locatorHeader.getValue() : null; int numberOfRecords = (numRecordsHeader != null) ? Integer.parseInt(numRecordsHeader.getValue()) : 0; return EntityUtils.toString(response.getEntity()); } else { String errorBody = EntityUtils.toString(response.getEntity()); throw new RuntimeException("Failed to get query job results, status code: " + statusCode + ", error: " + errorBody); } } finally { response.close(); httpClient.close(); } } /** * Create the BulkConnection used to call Bulk API operations. */ public static BulkConnection getBulkConnection(String userName, String password,String url) throws ConnectionException, AsyncApiException { ConnectorConfig partnerConfig = new ConnectorConfig(); partnerConfig.setUsername(userName); partnerConfig.setPassword(password); partnerConfig.setAuthEndpoint(url); // Creating the connection automatically handles login and stores // the session in partnerConfig new PartnerConnection(partnerConfig); // When PartnerConnection is instantiated, a login is implicitly // executed and, if successful, // a valid session is stored in the ConnectorConfig instance. // Use this key to initialize a BulkConnection: ConnectorConfig config = new ConnectorConfig(); config.setSessionId(partnerConfig.getSessionId()); // The endpoint for the Bulk API service is the same as for the normal // SOAP uri until the /Soap/ part. From here it's '/async/versionNumber' String soapEndpoint = partnerConfig.getServiceEndpoint(); String apiVersion = "56.0"; String restEndpoint = soapEndpoint.substring(0, soapEndpoint.indexOf("Soap/")) + "async/" + apiVersion; config.setRestEndpoint(restEndpoint); // This should only be false when doing debugging. config.setCompression(true); // Set this to true to see HTTP requests and responses on stdout config.setTraceMessage(false); BulkConnection connection = new BulkConnection(config); return connection; } public static JobInfo createQueryJob(String api, String soqlQuery,BulkConnection connection) throws Exception { JobInfo job = new JobInfo(); job.setOperation(OperationEnum.queryAll); ContentType jobContentType = ContentType.CSV; job.setContentType(jobContentType); ConcurrencyMode jobConcurrencyMode = ConcurrencyMode.Parallel; job.setConcurrencyMode(jobConcurrencyMode); job.setObject(soqlQuery); String endpoint = connection.getConfig().getRestEndpoint() + "job/"; HashMap headers = getHeaders(JSON_CONTENT_TYPE, JSON_CONTENT_TYPE, connection.getConfig().getSessionId()); HashMap requestBodyMap = new HashMap(); requestBodyMap.put("operation", job.getOperation().toString()); requestBodyMap.put("object", job.getObject()); requestBodyMap.put("contentType", job.getContentType().toString()); requestBodyMap.put("lineEnding", "LF"); requestBodyMap.put("concurrencyMode", job.getConcurrencyMode().toString()); HttpTransportImpl transport = HttpTransportImpl.getInstance(); OutputStream out = transport.connect(endpoint, headers, true, HttpTransportInterface.SupportedHttpMethodType.POST); ObjectMapper mapper = new ObjectMapper(); mapper.setDateFormat(CalendarCodec.getDateFormat()); String requestContent = mapper.writeValueAsString(requestBodyMap); out.write(requestContent.getBytes(UTF_8)); out.close(); InputStream in = transport.getContent(); boolean successfulRequest = transport.isSuccessful(); if (!successfulRequest) { throw new AsyncApiException("Failed to create query job", AsyncExceptionCode.Unknown); } ObjectMapper responseMapper = new ObjectMapper(); responseMapper.configure(DeserializationFeature.ACCEPT_SINGLE_VALUE_AS_ARRAY, true); responseMapper.setTimeZone(TimeZone.getTimeZone("GMT")); return responseMapper.readValue(in, JobInfo.class); } public static JobInfo createQueryJob1(String api, String soqlQuery,BulkConnection connection) throws Exception { JobInfo job = new JobInfo(); InputStream in = null; job.setOperation(OperationEnum.queryAll); ContentType jobContentType = ContentType.CSV; job.setContentType(jobContentType); ConcurrencyMode jobConcurrencyMode = ConcurrencyMode.Parallel; job.setConcurrencyMode(jobConcurrencyMode); job.setObject(soqlQuery); String endpoint =connection.getConfig().getRestEndpoint()+ "job/"; String requestURL = constructRequestURL(endpoint, job.getId()); HashMap headers = null; headers = getHeaders(JSON_CONTENT_TYPE, JSON_CONTENT_TYPE,connection.getConfig().getSessionId()); HashMap requestBodyMap = new HashMap(); ContentType type = job.getContentType(); requestBodyMap.put("object", job.getObject()); requestBodyMap.put("contentType", type.toString()); requestBodyMap.put("lineEnding", "LF"); OutputStream out; HttpTransportImpl transport = HttpTransportImpl.getInstance(); out = transport.connect(requestURL, headers, true, HttpTransportInterface.SupportedHttpMethodType.POST); ObjectMapper mapper = new ObjectMapper(); mapper.setDateFormat(CalendarCodec.getDateFormat()); String requestContent = mapper.writeValueAsString(requestBodyMap); out.write(requestContent.getBytes(UTF_8)); out.close(); in = transport.getContent(); boolean successfulRequest = true; successfulRequest = transport.isSuccessful(); JobInfo result = null; ObjectMapper mapper1 = new ObjectMapper(); mapper1.configure(DeserializationFeature.ACCEPT_SINGLE_VALUE_AS_ARRAY, true); // By default, ObjectMapper generates Calendar instances with UTC TimeZone. // Here, override that to "GMT" to better match the behavior of the WSC XML parser. mapper1.setTimeZone(TimeZone.getTimeZone("GMT")); return mapper1.readValue(in, JobInfo.class); } public static HashMap getHeaders(String requestContentType, String acceptContentType,String seesionId) { HashMap newMap = new HashMap(); String authHeaderValue = AUTH_HEADER_VALUE_PREFIX + seesionId; newMap.put(REQUEST_CONTENT_TYPE_HEADER, requestContentType); newMap.put(ACCEPT_CONTENT_TYPES_HEADER, acceptContentType); newMap.put(AUTH_HEADER, authHeaderValue); newMap.put(SFORCE_CALL_OPTIONS_HEADER, "client=data-dump, batchSize=5000"); return newMap; } private static String constructRequestURL(String urlString,String jobId) { urlString += "query/" + jobId + "/"; return urlString; } /** * Create and upload batches using a CSV file. * The file into the appropriate size batch files. * * @param connection * Connection to use for creating batches * @param jobInfo * Job associated with new batches * @param csvFileName * The source file for batch data */ public static List createBatchesFromCSVFile(BulkConnection connection, JobInfo jobInfo, String csvFileName) throws IOException, AsyncApiException { List batchInfos = new ArrayList(); BufferedReader rdr = new BufferedReader( new InputStreamReader(Files.newInputStream(Paths.get(csvFileName)),"GBK") ); // read the CSV header row byte[] headerBytes = (rdr.readLine() + "\n").getBytes("UTF-8"); int headerBytesLength = headerBytes.length; File tmpFile = File.createTempFile("bulkAPIInsert", ".csv"); // Split the CSV file into multiple batches try { FileOutputStream tmpOut = new FileOutputStream(tmpFile); int maxBytesPerBatch = 10000000; // 10 million bytes per batch int maxRowsPerBatch = 10000; // 10 thousand rows per batch int currentBytes = 0; int currentLines = 0; String nextLine; while ((nextLine = rdr.readLine()) != null) { byte[] bytes = (nextLine + "\n").getBytes("UTF-8"); // Create a new batch when our batch size limit is reached if (currentBytes + bytes.length > maxBytesPerBatch || currentLines > maxRowsPerBatch) { createBatch(tmpOut, tmpFile, batchInfos, connection, jobInfo); currentBytes = 0; currentLines = 0; } if (currentBytes == 0) { tmpOut = new FileOutputStream(tmpFile); tmpOut.write(headerBytes); currentBytes = headerBytesLength; currentLines = 1; } tmpOut.write(bytes); currentBytes += bytes.length; currentLines++; } // Finished processing all rows // Create a final batch for any remaining data if (currentLines > 1) { createBatch(tmpOut, tmpFile, batchInfos, connection, jobInfo); } } finally { tmpFile.delete(); } return batchInfos; } /** * Create a batch by uploading the contents of the file. * This closes the output stream. * * @param tmpOut * The output stream used to write the CSV data for a single batch. * @param tmpFile * The file associated with the above stream. * @param batchInfos * The batch info for the newly created batch is added to this list. * @param connection * The BulkConnection used to create the new batch. * @param jobInfo * The JobInfo associated with the new batch. */ public static void createBatch(FileOutputStream tmpOut, File tmpFile, List batchInfos, BulkConnection connection, JobInfo jobInfo) throws IOException, AsyncApiException { tmpOut.flush(); tmpOut.close(); FileInputStream tmpInputStream = new FileInputStream(tmpFile); try { BatchInfo batchInfo = connection.createBatchFromStream(jobInfo, tmpInputStream); System.out.println(batchInfo); batchInfos.add(batchInfo); } finally { tmpInputStream.close(); } } }