validate PK Chunking for large queries

Author: harishhk107

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

@@ -83,6 +83,18 @@ public static String createSObjectIdQuery(String query) {
     return SELECT + FIELD_ID + " " + fromStatement;
   }
 
+  /**
+   * Creates a COUNT query from an existing SOQL query.
+   * Replaces the SELECT fields with COUNT() while preserving the FROM and WHERE clauses.
+   *
+   * @param query the original SOQL query
+   * @return a COUNT SOQL query string
+   */
+  public static String createCountQuery(String query) {
+    String fromStatement = SalesforceQueryParser.getFromStatement(query);
+    return SELECT + "COUNT() " + fromStatement;
+  }
+
   /**
    * Generates SObject query filter based on provided values.
    *

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

@@ -48,6 +48,8 @@
 import io.cdap.plugin.salesforce.plugin.source.batch.util.BulkConnectionRetryWrapper;
 import io.cdap.plugin.salesforce.plugin.source.batch.util.SalesforceSourceConstants;
 import io.cdap.plugin.salesforce.plugin.source.batch.util.SalesforceSplitUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -71,6 +73,7 @@ public class SalesforceBatchSource extends
   public static final String NAME = "Salesforce";
 
   private final SalesforceSourceConfig config;
+  private static final Logger LOG = LoggerFactory.getLogger(SalesforceBatchSource.class);
   private Schema schema;
   private MapToRecordTransformer transformer;
   private Set<String> jobIds = new HashSet<>();
@@ -139,7 +142,7 @@ public void prepareRun(BatchSourceContext context) {
 
     authenticatorCredentials = config.getConnection().getAuthenticatorCredentials();
     List<SalesforceSplit> querySplits =
-        getSplits(config, authenticatorCredentials, context.getLogicalStartTime(), oAuthInfo);
+        getSplits(config, authenticatorCredentials, context.getLogicalStartTime(), oAuthInfo, false);
     querySplits.stream().forEach(salesforceSplit -> jobIds.add(salesforceSplit.getJobId()));
     context.setInput(Input.of(config.getReferenceNameOrNormalizedFQN(orgId, sObjectName),
         new SalesforceInputFormatProvider(
@@ -148,10 +151,22 @@ public void prepareRun(BatchSourceContext context) {
 
   public static List<SalesforceSplit> getSplits(
       SalesforceSourceConfig config, AuthenticatorCredentials authenticatorCredentials,
-      long logicStartTime, OAuthInfo oAuthInfo) {
+      long logicStartTime, OAuthInfo oAuthInfo, boolean autoEnablePKChunk) {
     String query = config.getQuery(logicStartTime, oAuthInfo);
     BulkConnection bulkConnection = SalesforceSplitUtil.getBulkConnection(authenticatorCredentials);
+
     boolean enablePKChunk = config.getEnablePKChunk();
+
+    if (enablePKChunk && autoEnablePKChunk) {
+      enablePKChunk = SalesforceSplitUtil.shouldAutoPKChunk(
+          query, authenticatorCredentials, SalesforceSourceConstants.AUTO_PK_CHUNK_THRESHOLD);
+      if (!enablePKChunk) {
+        LOG.info("PK chunking skipped: record count is below threshold {} or "
+                + "object/query does not support it.",
+                SalesforceSourceConstants.AUTO_PK_CHUNK_THRESHOLD);
+      }
+    }
+
     if (enablePKChunk) {
       String parent = config.getParent();
       int chunkSize = config.getChunkSize();

src/main/java/io/cdap/plugin/salesforce/plugin/source/batch/util/SalesforceSourceConstants.java

@@ -55,6 +55,7 @@ public class SalesforceSourceConstants {
   public static final int MAX_PK_CHUNK_SIZE = 250000;
   public static final int DEFAULT_PK_CHUNK_SIZE = 100000;
   public static final int MIN_PK_CHUNK_SIZE = 1;
+  public static final long AUTO_PK_CHUNK_THRESHOLD = 1_000_000;
   // https://developer.salesforce.com/docs/atlas.en-us.252.0.api_asynch.meta/api_asynch/
   // async_api_headers_enable_pk_chunking.htm
   // **Always use lowercase names** to ensure consistency, especially if the sObject name is manually provided.

src/main/java/io/cdap/plugin/salesforce/plugin/source/batch/util/SalesforceSplitUtil.java

@@ -26,16 +26,21 @@
 import com.sforce.async.JobInfo;
 import com.sforce.async.JobStateEnum;
 import com.sforce.async.OperationEnum;
+import com.sforce.soap.partner.PartnerConnection;
+import com.sforce.soap.partner.QueryResult;
+import com.sforce.ws.ConnectionException;
 import dev.failsafe.FailsafeException;
 import dev.failsafe.RetryPolicy;
 import dev.failsafe.TimeoutExceededException;
 import io.cdap.plugin.salesforce.BulkAPIBatchException;
 import io.cdap.plugin.salesforce.InvalidConfigException;
 import io.cdap.plugin.salesforce.SObjectDescriptor;
 import io.cdap.plugin.salesforce.SalesforceBulkUtil;
+import io.cdap.plugin.salesforce.SalesforceConnectionUtil;
 import io.cdap.plugin.salesforce.SalesforceQueryUtil;
 import io.cdap.plugin.salesforce.authenticator.Authenticator;
 import io.cdap.plugin.salesforce.authenticator.AuthenticatorCredentials;
+import io.cdap.plugin.salesforce.parser.SalesforceQueryParser;
 import io.cdap.plugin.salesforce.plugin.source.batch.SalesforceSplit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -263,6 +268,43 @@ public static boolean isPkChunkingSupported(String sobjectName) {
     return false;
   }
 
+  /**
+   * Determines whether PK chunking should be enabled automatically based on record count,
+   * object support, and query compatibility.
+   *
+   * @param query       the SOQL query
+   * @param credentials authenticator credentials for SOAP API
+   * @param threshold   the record count threshold above which PK chunking is enabled
+   * @return true if PK chunking should be auto-enabled, false otherwise
+   */
+  public static boolean shouldAutoPKChunk(String query, AuthenticatorCredentials credentials, long threshold) {
+    try {
+      String sObjectName = SObjectDescriptor.fromQuery(query).getName();
+      if (!isPkChunkingSupported(sObjectName)) {
+        LOG.debug("PK Chunking auto-decision: object '{}' is not supported for PK chunking", sObjectName);
+        return false;
+      }
+      if (SalesforceQueryParser.isRestrictedPKQuery(query)) {
+        LOG.debug("PK Chunking auto-decision: query contains restricted clauses, skipping PK chunking");
+        return false;
+      }
+      String countQuery = SalesforceQueryUtil.createCountQuery(query);
+      PartnerConnection partnerConnection = SalesforceConnectionUtil.getPartnerConnection(credentials);
+      QueryResult result = partnerConnection.query(countQuery);
+      int recordCount = result.getSize();
+      LOG.debug("PK Chunking auto-decision: object '{}' has {} records, threshold is {}",
+              sObjectName, recordCount, threshold);
+      return recordCount >= threshold;
+    } catch (ConnectionException e) {
+      LOG.warn("PK Chunking auto-decision: failed to execute count query, falling back to no PK chunking", e);
+      return false;
+    } catch (Exception e) {
+      LOG.warn("PK Chunking auto-decision: unexpected error, falling back to no PK chunking", e);
+      return false;
+    }
+  }
+
+
   // This is added only for UCS use case.
   private static boolean isCustomObject(String sobjectName) {
     return sobjectName.toLowerCase().endsWith("__c");

src/test/java/io/cdap/plugin/salesforce/SalesforceQueryUtilTest.java

@@ -227,4 +227,18 @@ public void testCreateSObjectIdQuery() {
 
     Assert.assertEquals("SELECT Id " + fromClause, sObjectIdQuery);
   }
+
+  @Test
+  public void testCreateCountQuery() {
+    String query = "SELECT Id,Name,SomeField FROM sObjectName WHERE LastModifiedDate>=2019-04-12T23:23:23Z";
+    String countQuery = SalesforceQueryUtil.createCountQuery(query);
+    Assert.assertEquals("SELECT COUNT() FROM sObjectName WHERE LastModifiedDate>=2019-04-12T23:23:23Z", countQuery);
+  }
+
+  @Test
+  public void testCreateCountQueryWithoutWhere() {
+    String query = "SELECT Id, Name FROM Account";
+    String countQuery = SalesforceQueryUtil.createCountQuery(query);
+    Assert.assertEquals("SELECT COUNT() FROM Account", countQuery);
+  }
 }

src/test/java/io/cdap/plugin/salesforce/etl/SalesforceBatchSourceETLTest.java

@@ -624,4 +624,131 @@ private void testPKChunk(ImmutableMap<String, String> properties) throws Excepti
     Assert.assertEquals("record3", results.get(2).get("Name"));
     Assert.assertEquals("record4", results.get(3).get("Name"));
   }
+
+  /**
+   * Test: UI toggle TRUE + records > 10M threshold → PK Chunking should be enabled
+   */
+  @Test
+  public void testPKChunkWithUITrueAndRecordsAboveThreshold() throws Exception {
+    // Setup: User explicitly enables PK Chunking and records exceed 10M threshold
+    ImmutableMap<String, String> properties = new ImmutableMap.Builder<String, String>()
+        .put("enablePKChunk", "true")
+        .put("chunkSize", "100000")
+        .build();
+
+    // Verify behavior: PK Chunking should be enabled
+    testPKChunk(properties);
+  }
+
+  /**
+   * Test: UI toggle FALSE → PK Chunking should never be enabled
+   * regardless of record count (respects explicit user setting)
+   */
+  @Test
+  public void testPKChunkWithUIFalse() throws Exception {
+    String sObjectName = createCustomObject("IT_PKChunkDisabled", null);
+
+    List<SObject> sObjects = new ImmutableList.Builder<SObject>()
+        .add(new SObjectBuilder()
+            .setType(sObjectName)
+            .put("Name", "record1")
+            .build())
+        .add(new SObjectBuilder()
+            .setType(sObjectName)
+            .put("Name", "record2")
+            .build())
+        .build();
+
+    addSObjects(sObjects, true);
+
+    // User explicitly disables PK Chunking
+    ImmutableMap<String, String> properties = new ImmutableMap.Builder<String, String>()
+        .put("enablePKChunk", "false")
+        .build();
+
+    ImmutableMap.Builder<String, String> baseProperties = getBaseProperties("SalesforceReaderPKChunkDisabled");
+    baseProperties.putAll(properties);
+    baseProperties.put(SalesforceSourceConstants.PROPERTY_QUERY, "Select Name from " + sObjectName);
+
+    List<StructuredRecord> results = getPipelineResults(baseProperties.build(), SalesforceBatchSource.NAME,
+        "SalesforceBatch");
+
+    // Verify: Records are returned but PK Chunking should not be applied
+    results.sort(Comparator.comparing(record -> record.get("Name")));
+    Assert.assertEquals(2, results.size());
+    Assert.assertEquals("record1", results.get(0).get("Name"));
+    Assert.assertEquals("record2", results.get(1).get("Name"));
+  }
+
+  /**
+   * Test: Default (no UI setting) + small filtered query (< 10M records)
+   * → PK Chunking should be disabled (prevents empty chunks)
+   */
+  @Test
+  public void testPKChunkWithDefaultSettingAndSmallFilteredQuery() throws Exception {
+    String sObjectName = createCustomObject("IT_PKChunkSmallQuery", null);
+
+    // Create a small dataset
+    List<SObject> sObjects = new ImmutableList.Builder<SObject>()
+        .add(new SObjectBuilder()
+            .setType(sObjectName)
+            .put("Name", "filtered_record")
+            .put("Status", "Active")
+            .build())
+        .build();
+
+    addSObjects(sObjects, true);
+
+    // No explicit PK Chunking setting - use default
+    ImmutableMap.Builder<String, String> baseProperties = getBaseProperties("SalesforceReaderSmallQuery");
+    baseProperties.put(SalesforceSourceConstants.PROPERTY_QUERY,
+        "Select Name from " + sObjectName + " WHERE Status = 'Active'");
+
+    List<StructuredRecord> results = getPipelineResults(baseProperties.build(), SalesforceBatchSource.NAME,
+        "SalesforceBatch");
+
+    // Verify: Small filtered query should NOT use PK Chunking
+    Assert.assertEquals(1, results.size());
+    Assert.assertEquals("filtered_record", results.get(0).get("Name"));
+  }
+
+  /**
+   * Test: Custom Object with PK Chunking enabled
+   * Verifies that custom objects (__c suffix) support PK Chunking
+   */
+  @Test
+  public void testPKChunkWithCustomObject() throws Exception {
+    String customObjectName = createCustomObject("CustomPKChunk__c", null);
+
+    List<SObject> sObjects = new ImmutableList.Builder<SObject>()
+        .add(new SObjectBuilder()
+            .setType(customObjectName)
+            .put("Name", "custom_record1")
+            .build())
+        .add(new SObjectBuilder()
+            .setType(customObjectName)
+            .put("Name", "custom_record2")
+            .build())
+        .build();
+
+    addSObjects(sObjects, true);
+
+    ImmutableMap<String, String> properties = new ImmutableMap.Builder<String, String>()
+        .put("enablePKChunk", "true")
+        .put("chunkSize", "100000")
+        .build();
+
+    ImmutableMap.Builder<String, String> baseProperties = getBaseProperties("SalesforceReaderCustomPKChunk");
+    baseProperties.putAll(properties);
+    baseProperties.put(SalesforceSourceConstants.PROPERTY_QUERY, "Select Name from " + customObjectName);
+
+    List<StructuredRecord> results = getPipelineResults(baseProperties.build(), SalesforceBatchSource.NAME,
+        "SalesforceBatch");
+
+    // Verify: Custom object with PK Chunking enabled should work correctly
+    results.sort(Comparator.comparing(record -> record.get("Name")));
+    Assert.assertEquals(2, results.size());
+    Assert.assertEquals("custom_record1", results.get(0).get("Name"));
+    Assert.assertEquals("custom_record2", results.get(1).get("Name"));
+  }
 }