data-dump/src/main/java/com/celnet/datadump/util/BulkUtil.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);
                    }
                }
            }
        }
    }

    /**
     * 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();
        }
    }
}