Merge pull request #336 from data-integrations/pk-chunk-check-dev

Moved the pk chunking check to SplitUtil class

Author: vikasrathee-cs

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

@@ -57,6 +57,8 @@ public class SalesforceSourceConstants {
   public static final int MIN_PK_CHUNK_SIZE = 1;
   // 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.
+  // Update this list with each API version upgrade.
   public static final List<String> SUPPORTED_OBJECTS_WITH_PK_CHUNK = Arrays.asList("account",
                                                                                    "accountcontactrelation",
                                                                                    "accountteammember",
@@ -108,7 +110,7 @@ public class SalesforceSourceConstants {
                                                                                    "loginhistory",
                                                                                    "loyaltyaggrpointexprledger",
                                                                                    "loyaltyledger",
-                                                                                   "LoyaltyLedgerTraceability",
+                                                                                   "loyaltyledgertraceability",
                                                                                    "loyaltymembercurrency",
                                                                                    "loyaltymembertier",
                                                                                    "loyaltypartnerproduct",

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

@@ -15,6 +15,7 @@
  */
 package io.cdap.plugin.salesforce.plugin.source.batch.util;
 
+import com.google.common.base.Strings;
 import com.sforce.async.AsyncApiException;
 import com.sforce.async.AsyncExceptionCode;
 import com.sforce.async.BatchInfo;
@@ -252,4 +253,18 @@ public static RetryPolicy<Object> getRetryPolicy(Long initialRetryDuration, Long
       return RetryPolicy.builder().withMaxRetries(0).build();
     }
   }
+
+  // This is added for UCS use case only to identify the objects where PK chunking needs to be enabled by default.
+  public static boolean isPkChunkingSupported(String sobjectName) {
+    if (!Strings.isNullOrEmpty(sobjectName)) {
+      return SalesforceSourceConstants.SUPPORTED_OBJECTS_WITH_PK_CHUNK.contains(sobjectName.toLowerCase())
+        || isCustomObject(sobjectName);
+    }
+    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/SalesforceBulkUtilTest.java

@@ -23,6 +23,9 @@
 import com.sforce.async.BulkConnection;
 import com.sforce.async.ConcurrencyMode;
 import com.sforce.async.JobInfo;
+import io.cdap.plugin.salesforce.plugin.source.batch.util.SalesforceSourceConstants;
+import io.cdap.plugin.salesforce.plugin.source.batch.util.SalesforceSplitUtil;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -83,4 +86,25 @@ public void testAwaitCompletionBatchCompleted() throws Exception {
     // Call the awaitCompletion method and verify that it completes successfully
     SalesforceBulkUtil.awaitCompletion(bulkConnection, job, Collections.singletonList(batchInfo), true);
   }
+
+  @Test
+  public void testPKChunkingEnabledForSobjectsWithAllCases() {
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("Account"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("account"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("aCcOuNt"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("Case"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("CONTACT"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("test_custom__c"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("Test_custom__c"));
+    Assert.assertTrue(SalesforceSplitUtil.isPkChunkingSupported("TEST_CUSTOM__C"));
+    Assert.assertFalse(SalesforceSplitUtil.isPkChunkingSupported("Not_Supported"));
+    Assert.assertFalse(SalesforceSplitUtil.isPkChunkingSupported("notsupported"));
+  }
+
+  @Test
+  public void testPKChunkingSupportedListContainsNoUpperCaseValues() {
+    // SUPPORTED_OBJECTS_WITH_PK_CHUNK list must have only lower case values.
+    Assert.assertFalse(SalesforceSourceConstants.SUPPORTED_OBJECTS_WITH_PK_CHUNK.stream()
+      .anyMatch(s -> s.matches(".*[A-Z].*")));
+  }
 }