513 lines
23 KiB
Java
513 lines
23 KiB
Java
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<String> incomplete = new HashSet<String>();
|
||
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<BatchInfo> batchInfoList)
|
||
throws AsyncApiException {
|
||
long sleepTime = 0L;
|
||
Set<String> incomplete = new HashSet<String>();
|
||
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<String> incomplete = new HashSet<String>();
|
||
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<String, String> headers = getHeaders(JSON_CONTENT_TYPE, JSON_CONTENT_TYPE, connection.getConfig().getSessionId());
|
||
|
||
HashMap<String, Object> requestBodyMap = new HashMap<String, Object>();
|
||
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<String, String> headers = null;
|
||
headers = getHeaders(JSON_CONTENT_TYPE, JSON_CONTENT_TYPE,connection.getConfig().getSessionId());
|
||
HashMap<String, Object> requestBodyMap = new HashMap<String, Object>();
|
||
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<String, String> getHeaders(String requestContentType, String acceptContentType,String seesionId) {
|
||
HashMap<String, String> newMap = new HashMap<String, String>();
|
||
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<BatchInfo> createBatchesFromCSVFile(BulkConnection connection,
|
||
JobInfo jobInfo, String csvFileName)
|
||
throws IOException, AsyncApiException {
|
||
List<BatchInfo> batchInfos = new ArrayList<BatchInfo>();
|
||
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<BatchInfo> 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();
|
||
}
|
||
}
|
||
}
|