Implement PK Chunking

Author: ardian4

docs/Salesforce-batchsource.md

@@ -113,3 +113,99 @@ The Salesforce types will be automatically mapped to schema types as shown below
 |             | textarea, phone, id, url, email, encryptedstring,                          |              |
 |             | datacategorygroupreference, location, address, anyType, json, complexvalue |              |
 
+**Enable PK Chunking:**
+Primary key (PK) Chunking splits query on large tables into chunks based on the record IDs, or primary keys, of the queried records.
+
+Salesforce recommends that you enable PK chunking when querying tables with more than 10 million records or when a bulk query consistently times out. 
+However, the effectiveness of PK chunking depends on the specifics of the query and the queried data.
+
+For example, let’s say you enable PK chunking for the following query on an Account table with 10,000,000 records.
+
+`SELECT Name FROM Account`
+
+Assuming a chunk size of 250,000 the query is split into the following 40 queries. Each query is processed parallely
+
+| Queries |         
+| ------- | 
+| SELECT Name FROM Account WHERE Id >= 001300000000000 AND Id < 00130000000132G |
+| SELECT Name FROM Account WHERE Id >= 00130000000132G AND Id < 00130000000264W |
+| SELECT Name FROM Account WHERE Id >= 00130000000264W AND Id < 00130000000396m |
+| ... |
+| SELECT Name FROM Account WHERE Id >= 00130000000euQ4 AND Id < 00130000000fxSK |
+
+PK chunking works only with queries that don’t include `SELECT` clauses or conditions other than `WHERE`.
+
+PK chunking only works with the following objects:
+
+| Objects |         
+| ------- | 
+|Account|
+|AccountContactRelation|
+|AccountTeamMember|
+|AiVisitSummary|
+|Asset|
+|B2BMktActivity|
+|B2BMktProspect|
+|Campaign|
+|CampaignMember|
+|CandidateAnswer|
+|Case|
+|CaseArticle|
+|CaseComment|
+|Claim|
+|ClaimParticipant|
+|Contact|
+|ContractLineItem|
+|ConversationEntry|
+|CustomerProperty|
+|EinsteinAnswerFeedback|
+|EmailMessage|
+|EngagementScore|
+|Event|
+|EventRelation|
+|FeedItem|
+|Individual|
+|InsurancePolicy|
+|InsurancePolicyAsset|
+|InsurancePolicyParticipant|
+|Lead|
+|LeadInsight|
+|LiveChatTranscript|
+|LoginHistory|
+|LoyaltyLedger|
+|LoyaltyMemberCurrency|
+|LoyaltyMemberTier|
+|LoyaltyPartnerProduct|
+|LoyaltyProgramMember|
+|LoyaltyProgramPartner|
+|Note|
+|ObjectTerritory2Association|
+|Opportunity|
+|OpportunityContactRole|
+|OpportunityHistory|
+|OpportunityLineItem|
+|OpportunitySplit|
+|OpportunityTeamMember|
+|Pricebook2|
+|PricebookEntry|
+|Product2|
+|ProductConsumed|
+|ProductRequired|
+|QuickText|
+|Quote|
+|QuoteLineItem|
+|ReplyText|
+|ScoreIntelligence|
+|ServiceContract|
+|Task|
+|TermDocumentFrequency|
+|TransactionJournal|
+|User|
+|UserRole|
+|VoiceCall|
+|WorkOrder|
+|WorkOrderLineItem|
+
+Support also includes custom objects, and any Sharing and History tables that support standard objects.
+
+**Chunk Size:** Specify size of chunk. Maximum Size is 250,000. Default Size is 100,000. 

pom.xml

@@ -31,7 +31,7 @@
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <!-- version properties -->
     <avro.version>1.7.7</avro.version>
-    <cdap.version>6.2.0</cdap.version>
+    <cdap.version>6.4.0-SNAPSHOT</cdap.version>
     <junit.version>4.12</junit.version>
     <commons.io.version>2.6</commons.io.version>
     <commons.lang.version>2.6</commons.lang.version>

src/main/java/io/cdap/plugin/salesforce/SObjectsDescribeResult.java

@@ -221,4 +221,8 @@ private static DescribeSObjectResult describe(PartnerConnection connection, Stri
     cache.put(name.toLowerCase(), describe);
     return describe;
   }
+
+  public static boolean isCustomObject(PartnerConnection connection, String name) throws ConnectionException {
+    return describe(connection, name, new HashMap<>()).isCustom();
+  }
 }

src/main/java/io/cdap/plugin/salesforce/SalesforceBulkUtil.java

@@ -37,6 +37,7 @@
 import java.io.InputStream;
 import java.io.SequenceInputStream;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -109,21 +110,59 @@ public static void closeJob(BulkConnection bulkConnection, String jobId) throws
    *
    * @param bulkConnection bulk connection instance
    * @param query a SOQL query
+   * @param enablePKChunk enable PK Chunk
    * @return an array of batches
    * @throws AsyncApiException  if there is an issue creating the job
    * @throws IOException failed to close the query
    */
-  public static BatchInfo[] runBulkQuery(BulkConnection bulkConnection, String query)
+  public static BatchInfo[] runBulkQuery(BulkConnection bulkConnection, String query, boolean enablePKChunk)
     throws AsyncApiException, IOException {
 
     SObjectDescriptor sObjectDescriptor = SObjectDescriptor.fromQuery(query);
     JobInfo job = createJob(bulkConnection, sObjectDescriptor.getName(), OperationEnum.query, null);
-
+    BatchInfo batchInfo;
     try (ByteArrayInputStream bout = new ByteArrayInputStream(query.getBytes())) {
-      bulkConnection.createBatchFromStream(job, bout);
+      batchInfo = bulkConnection.createBatchFromStream(job, bout);
     }
+    return enablePKChunk ? waitForBatchChunks(bulkConnection, job.getId(), batchInfo.getId()) :
+      bulkConnection.getBatchInfoList(job.getId()).getBatchInfo();
+  }
+
+  /** When PK Chunk is enabled, wait for state of initial batch to be NotProcessed, in this case Salesforce API will
+   * decide how many batches will be created
+   * @param bulkConnection bulk connection instance
+   * @param jobId a job id
+   * @param initialBatchId a batch id
+   * @return Array with Batches created by Salesforce API
+   *
+   * @throws AsyncApiException if there is an issue creating the job
+   */
+  private static BatchInfo[] waitForBatchChunks(BulkConnection bulkConnection, String jobId, String initialBatchId)
+    throws AsyncApiException {
+    BatchInfo initialBatchInfo = null;
+    for (int i = 0; i < GET_BATCH_RESULTS_TRIES; i++) {
+      //check if the job is aborted
+      if (bulkConnection.getJobStatus(jobId).getState() == JobStateEnum.Aborted) {
+        LOG.info(String.format("Job with Id: '%s' is aborted", jobId));
+        return new BatchInfo[0];
+      }
+      initialBatchInfo = bulkConnection.getBatchInfo(jobId, initialBatchId);
 
-    return bulkConnection.getBatchInfoList(job.getId()).getBatchInfo();
+      if (initialBatchInfo.getState() == BatchStateEnum.NotProcessed) {
+        BatchInfo[] result = bulkConnection.getBatchInfoList(jobId).getBatchInfo();
+        return Arrays.stream(result).filter(batchInfo -> batchInfo.getState() != BatchStateEnum.NotProcessed)
+          .toArray(BatchInfo[]::new);
+      } else if (initialBatchInfo.getState() == BatchStateEnum.Failed) {
+        throw new BulkAPIBatchException("Batch failed", initialBatchInfo);
+      } else {
+        try {
+          Thread.sleep(GET_BATCH_RESULTS_SLEEP_MS);
+        } catch (InterruptedException e) {
+          throw new RuntimeException("Job is aborted", e);
+        }
+      }
+    }
+    throw new BulkAPIBatchException("Timeout waiting for batch results", initialBatchInfo);
   }
 
   /**

src/main/java/io/cdap/plugin/salesforce/parser/SalesforceQueryParser.java

@@ -65,6 +65,18 @@ public static boolean isRestrictedQuery(String query) {
     return visitor.visit(parser.statement());
   }
 
+  /**
+   * Checks if query has restricted syntax that cannot be processed by Bulk API with PK Enable.
+   *
+   * @param query SOQL query
+   * @return true if query has restricted syntax, false otherwise
+   */
+  public static boolean isRestrictedPKQuery(String query) {
+    SOQLParser parser = initParser(query);
+    SalesforceQueryVisitor.RestrictedPKQueryVisitor visitor = new SalesforceQueryVisitor.RestrictedPKQueryVisitor();
+    return visitor.visit(parser.statement());
+  }
+
   private static SOQLParser initParser(String query) {
     SOQLLexer lexer = new SOQLLexer(CharStreams.fromString(query));
     lexer.removeErrorListeners();

src/main/java/io/cdap/plugin/salesforce/parser/SalesforceQueryVisitor.java

@@ -304,4 +304,28 @@ protected Boolean defaultResult() {
       return Boolean.FALSE;
     }
   }
+
+  /**
+   * Visits query fields that are restricted by Bulk API with PKChunk Enable.
+   * Returns true if there is condition clause other than 'WHERE', false otherwise.
+   * For example: SELECT Name FROM Opportunity LIMIT 10.
+   */
+  public static class RestrictedPKQueryVisitor extends SOQLBaseVisitor<Boolean> {
+
+    @Override
+    public Boolean visitStatement(SOQLParser.StatementContext ctx) {
+      SOQLParser.FromStatementContext fromStatementContext = ctx.fromStatement();
+      if (
+        fromStatementContext.WITH() != null ||
+          fromStatementContext.GROUP() != null ||
+          fromStatementContext.ORDER() != null ||
+          fromStatementContext.LIMIT() != null ||
+          fromStatementContext.OFFSET() != null ||
+          fromStatementContext.FOR() != null) {
+        return true;
+      }
+
+      return ctx.fieldList().accept(new RestrictedFieldVisitor());
+    }
+  }
 }

src/main/java/io/cdap/plugin/salesforce/plugin/source/batch/SalesforceInputFormat.java

@@ -64,8 +64,16 @@ public List<InputSplit> getSplits(JobContext context) {
     List<String> queries = GSON.fromJson(configuration.get(SalesforceSourceConstants.CONFIG_QUERIES), QUERIES_TYPE);
     BulkConnection bulkConnection = getBulkConnection(configuration);
 
+    boolean enablePKChunk = configuration.getBoolean(SalesforceSourceConstants.CONFIG_PK_CHUNK_ENABLE, false);
+    if (enablePKChunk) {
+      int chunkSize = configuration.getInt(SalesforceSourceConstants.CONFIG_CHUNK_SIZE,
+                                           SalesforceSourceConstants.DEFAULT_PK_CHUNK_SIZE);
+      bulkConnection.addHeader(SalesforceSourceConstants.HEADER_ENABLE_PK_CHUNK,
+                               String.format(SalesforceSourceConstants.HEADER_VALUE_PK_CHUNK, chunkSize));
+    }
+
     return queries.parallelStream()
-      .map(query -> getQuerySplits(query, bulkConnection))
+      .map(query -> getQuerySplits(query, bulkConnection, enablePKChunk))
       .flatMap(Collection::stream)
       .collect(Collectors.toList());
   }
@@ -87,8 +95,8 @@ public RecordReader createRecordReader(InputSplit split, TaskAttemptContext cont
     return new SalesforceRecordReaderWrapper(sObjectName, sObjectNameField, getDelegateRecordReader(query, schema));
   }
 
-  private List<SalesforceSplit> getQuerySplits(String query, BulkConnection bulkConnection) {
-    return Stream.of(getBatches(query, bulkConnection))
+  private List<SalesforceSplit> getQuerySplits(String query, BulkConnection bulkConnection, boolean enablePKChunk) {
+    return Stream.of(getBatches(query, bulkConnection, enablePKChunk))
       .map(batch -> new SalesforceSplit(batch.getJobId(), batch.getId(), query))
       .collect(Collectors.toList());
   }
@@ -116,15 +124,16 @@ private BulkConnection getBulkConnection(Configuration conf) {
    *
    * @param query SOQL query
    * @param bulkConnection bulk connection
+   * @param enablePKChunk enable PK Chunking
    * @return array of batch info
    */
-  private BatchInfo[] getBatches(String query, BulkConnection bulkConnection) {
+  private BatchInfo[] getBatches(String query, BulkConnection bulkConnection, boolean enablePKChunk) {
     try {
       if (!SalesforceQueryUtil.isQueryUnderLengthLimit(query)) {
         LOG.debug("Wide object query detected. Query length '{}'", query.length());
         query = SalesforceQueryUtil.createSObjectIdQuery(query);
       }
-      BatchInfo[] batches = SalesforceBulkUtil.runBulkQuery(bulkConnection, query);
+       BatchInfo[] batches = SalesforceBulkUtil.runBulkQuery(bulkConnection, query, enablePKChunk);
       LOG.debug("Number of batches received from Salesforce: '{}'", batches.length);
       return batches;
     } catch (AsyncApiException | IOException e) {

src/main/java/io/cdap/plugin/salesforce/plugin/source/batch/SalesforceInputFormatProvider.java

@@ -66,6 +66,28 @@ public SalesforceInputFormatProvider(SalesforceBaseSourceConfig config,
     this.conf = configBuilder.build();
   }
 
+  public SalesforceInputFormatProvider(SalesforceSourceConfig config,
+                                       List<String> queries,
+                                       Map<String, String> schemas,
+                                       @Nullable String sObjectNameField) {
+    ImmutableMap.Builder<String, String> builder = new ImmutableMap.Builder<String, String>()
+      .put(SalesforceConstants.CONFIG_USERNAME, config.getUsername())
+      .put(SalesforceConstants.CONFIG_PASSWORD, config.getPassword())
+      .put(SalesforceConstants.CONFIG_CONSUMER_KEY, config.getConsumerKey())
+      .put(SalesforceConstants.CONFIG_CONSUMER_SECRET, config.getConsumerSecret())
+      .put(SalesforceConstants.CONFIG_LOGIN_URL, config.getLoginUrl())
+      .put(SalesforceSourceConstants.CONFIG_QUERIES, GSON.toJson(queries))
+      .put(SalesforceSourceConstants.CONFIG_SCHEMAS, GSON.toJson(schemas))
+      .put(SalesforceSourceConstants.CONFIG_PK_CHUNK_ENABLE, String.valueOf(config.getEnablePKChunk()))
+      .put(SalesforceSourceConstants.CONFIG_CHUNK_SIZE, String.valueOf(config.getChunkSize()));
+
+    if (sObjectNameField != null) {
+      builder.put(SalesforceSourceConstants.CONFIG_SOBJECT_NAME_FIELD, sObjectNameField);
+    }
+
+    this.conf = builder.build();
+  }
+
   @Override
   public Map<String, String> getInputFormatConfiguration() {
     return conf;