Added generic method for retry logic

Author: Krish-cloudsufi

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

@@ -228,7 +228,7 @@ public String getOrgId(OAuthInfo oAuthInfo) throws ConnectionException {
     );
 
     PartnerConnection partnerConnection = SalesforceConnectionUtil.getPartnerConnection(credentials);
-    GetUserInfoResult userInfo = PartnerConnectionRetryWrapper.getUserInfoWithRetry(partnerConnection, this);
+    GetUserInfoResult userInfo = PartnerConnectionRetryWrapper.getUserInfoWithRetry(partnerConnection);
     return userInfo.getOrganizationId();
   }
 

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

@@ -186,7 +186,7 @@ public void validateSObjects(FailureCollector collector, @Nullable OAuthInfo oAu
                                                                           this.getConnection().getProxyUrl());
       PartnerConnection partnerConnection =
         SalesforceConnectionUtil.getPartnerConnection(credentials);
-        describeGlobalResult = PartnerConnectionRetryWrapper.describeGlobalWithRetry(partnerConnection, this);
+        describeGlobalResult = PartnerConnectionRetryWrapper.describeGlobalWithRetry(partnerConnection);
     } catch (ConnectionException e) {
       String message = SalesforceConnectionUtil.getSalesforceErrorMessageFromException(e);
       throw new IllegalArgumentException(String.format("Unable to connect to Salesforce due to error: %s", message), e);
@@ -229,7 +229,7 @@ private List<String> getSObjects() {
     try {
       PartnerConnection partnerConnection =
         SalesforceConnectionUtil.getPartnerConnection(getConnection().getAuthenticatorCredentials());
-        describeGlobalResult = PartnerConnectionRetryWrapper.describeGlobalWithRetry(partnerConnection, this);
+        describeGlobalResult = PartnerConnectionRetryWrapper.describeGlobalWithRetry(partnerConnection);
     } catch (ConnectionException e) {
       String message = SalesforceConnectionUtil.getSalesforceErrorMessageFromException(e);
       throw new IllegalArgumentException(String.format("Unable to connect to Salesforce due to error: %s", message), e);

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

@@ -66,11 +66,10 @@ public SalesforceSoapRecordReader(Schema schema, String query, SoapRecordToMapTr
    * @param taskAttemptContext task context
    */
 
-  private Configuration configuration;
   @Override
   public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext) {
-    this.configuration = taskAttemptContext.getConfiguration();
-    AuthenticatorCredentials credentials = SalesforceConnectionUtil.getAuthenticatorCredentials(configuration);
+    Configuration conf = taskAttemptContext.getConfiguration();
+    AuthenticatorCredentials credentials = SalesforceConnectionUtil.getAuthenticatorCredentials(conf);
     initialize(credentials);
   }
 

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

@@ -23,16 +23,64 @@
 import com.sforce.soap.partner.sobject.SObject;
 import com.sforce.ws.ConnectionException;
 import dev.failsafe.Failsafe;
-import io.cdap.plugin.salesforce.plugin.source.batch.SalesforceBaseSourceConfig;
+import dev.failsafe.FailsafeException;
+import dev.failsafe.RetryPolicy;
 
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.Callable;
+import java.util.function.Predicate;
+
 
 
 /**
  * A utility class that wraps Salesforce PartnerConnection calls with retry logic using Failsafe.
  */
 public class PartnerConnectionRetryWrapper {
 
+  /**
+   * Executes a Salesforce PartnerConnection operation with retry logic using the provided {@link RetryPolicy}.
+   * <p>
+   * This method wraps the execution in {@link Failsafe} to handle transient {@link ConnectionException}s
+   * and retries the operation if the exception is considered retryable.
+   * It also performs a null or custom result validation check, throwing an {@link IllegalArgumentException}
+   * if the result is invalid.
+   *
+   * @param retryPolicy   the retry policy to use for retries
+   * @param operation     the Salesforce PartnerConnection operation to execute
+   * @param errorContext  a descriptive message used when the result is null or fails the validation check
+   * @param nullCheck     an optional predicate to validate the result; can be {@code null}
+   * @param <T>           the type of the result returned by the operation
+   * @return the result of the successful operation
+   * @throws ConnectionException if the operation fails with a non-retryable {@link ConnectionException}
+   * @throws FailsafeException if the operation fails with a non-ConnectionException or all retries are exhausted
+   * @throws IllegalArgumentException if the result is null or fails the provided validation check
+   */
+  private static <T> T executeWithRetry(RetryPolicy<Object> retryPolicy, Callable<T> operation,
+                                        String errorContext, Predicate<T> nullCheck) throws ConnectionException {
+    try {
+      return Failsafe.with(retryPolicy).get(() -> {
+        try {
+          T result = operation.call();
+          if (nullCheck != null && (result == null || !nullCheck.test(result))) {
+            throw new IllegalArgumentException(errorContext);
+          }
+          return result;
+        } catch (ConnectionException e) {
+          if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
+            throw new SalesforceQueryExecutionException(e);
+          }
+          throw e;
+        }
+      });
+    } catch (FailsafeException ex) {
+      if (ex.getCause() instanceof ConnectionException) {
+        throw (ConnectionException) ex.getCause();
+      }
+      throw ex;
+    }
+  }
+
   /**
    * Executes a SOQL query using PartnerConnection with retry logic.
    *
@@ -42,20 +90,12 @@ public class PartnerConnectionRetryWrapper {
    * @throws ConnectionException if the query fails and is not retryable
    */
   public static QueryResult queryWithRetry(PartnerConnection connection, String query) throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        QueryResult result = connection.query(query);
-        if (result == null) {
-          throw new IllegalArgumentException("The query returned a null result: " + query);
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.query(query),
+        "The query returned a null result: " + query,
+        Objects::nonNull
+    );
   }
 
   /**
@@ -69,21 +109,12 @@ public static QueryResult queryWithRetry(PartnerConnection connection, String qu
    */
   public static QueryResult queryMoreWithRetry(PartnerConnection connection, String locator,
                                                String originalQueryForContext) throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        QueryResult result = connection.queryMore(locator);
-        if (result == null) {
-          throw new IllegalArgumentException("The QueryMore returned a null result for locator used in query: "
-              + originalQueryForContext);
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.queryMore(locator),
+        "The QueryMore returned a null result for locator used in query: " + originalQueryForContext,
+        Objects::nonNull
+    );
   }
 
   /**
@@ -97,20 +128,12 @@ public static QueryResult queryMoreWithRetry(PartnerConnection connection, Strin
   public static DescribeSObjectResult[] describeSObjectsWithRetry(PartnerConnection connection,
                                                                   List<String> sObjectNames) throws
       ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        DescribeSObjectResult[] result = connection.describeSObjects(sObjectNames.toArray(new String[0]));
-        if (result == null) {
-          throw new IllegalArgumentException("The describeSObjects returned a null result: " + sObjectNames);
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.describeSObjects(sObjectNames.toArray(new String[0])),
+        "The describeSObjects returned a null result: " + sObjectNames,
+        Objects::nonNull
+    );
   }
 
   /**
@@ -123,20 +146,12 @@ public static DescribeSObjectResult[] describeSObjectsWithRetry(PartnerConnectio
    */
   public static DescribeSObjectResult describeSObjectWithRetry(PartnerConnection connection, String sObjectName)
       throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        DescribeSObjectResult result = connection.describeSObject(sObjectName);
-        if (result == null) {
-          throw new IllegalArgumentException("The describeSObject returned a null result: " + sObjectName);
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.describeSObject(sObjectName),
+        "The describeSObject returned a null result: " + sObjectName,
+        Objects::nonNull
+    );
   }
 
   /**
@@ -148,76 +163,12 @@ public static DescribeSObjectResult describeSObjectWithRetry(PartnerConnection c
    */
   public static DescribeGlobalResult describeGlobalWithRetry(PartnerConnection connection)
       throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        DescribeGlobalResult result = connection.describeGlobal();
-        if (result == null) {
-          throw new IllegalArgumentException("The describeGlobal() returned a null result");
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
-  }
-
-  /**
-   * Describes all available Salesforce objects using describeGlobal with a custom retry policy from config.
-   *
-   * @param connection the Salesforce PartnerConnection
-   * @param config the base source config containing retry parameters
-   * @return DescribeGlobalResult containing metadata for all objects
-   * @throws ConnectionException if the describeGlobal call fails and is not retryable
-   */
-  public static DescribeGlobalResult describeGlobalWithRetry(PartnerConnection connection,
-                                                             SalesforceBaseSourceConfig config)
-      throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getRetryPolicy(config.getInitialRetryDuration(),
-        config.getMaxRetryDuration(), config.getMaxRetryCount())).get(() -> {
-      try {
-        DescribeGlobalResult result = connection.describeGlobal();
-        if (result == null) {
-          throw new IllegalArgumentException("The describeGlobal() returned a null result");
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
-  }
-
-  /**
-   * Retrieves user info from the Salesforce connection with retry logic using config.
-   *
-   * @param connection the Salesforce PartnerConnection
-   * @param config the base source config containing retry parameters
-   * @return GetUserInfoResult containing user and organization info
-   * @throws ConnectionException if the getUserInfo call fails and is not retryable
-   */
-  public static GetUserInfoResult getUserInfoWithRetry(PartnerConnection connection,
-                                                       SalesforceBaseSourceConfig config)
-      throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getRetryPolicy(config.getInitialRetryDuration(),
-        config.getMaxRetryDuration(), config.getMaxRetryCount())).get(() -> {
-      try {
-        GetUserInfoResult userInfo = connection.getUserInfo();
-        if (userInfo == null || userInfo.getOrganizationId() == null) {
-          throw new IllegalArgumentException("UserInfo or OrganizationId is null");
-        }
-        return userInfo;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        connection::describeGlobal,
+        "The describeGlobal() returned a null result",
+        Objects::nonNull
+    );
   }
 
   /**
@@ -234,20 +185,29 @@ public static SObject[] retrieveWithRetry(PartnerConnection connection,
                                             String fields,
                                             String sObjectName,
                                             String[] sObjectIds) throws ConnectionException {
-    return Failsafe.with(SalesforceSplitUtil.getDefaultRetryPolicy()).get(() -> {
-      try {
-        SObject[] result = connection.retrieve(fields, sObjectName, sObjectIds);
-        if (result == null) {
-          throw new IllegalArgumentException("The Retrieve returned a null result: " + sObjectName);
-        }
-        return result;
-      } catch (ConnectionException e) {
-        if (SalesforceSplitUtil.isRetryableConnectionError(e)) {
-          throw new SalesforceQueryExecutionException(e);
-        }
-        throw e;
-      }
-    });
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.retrieve(fields, sObjectName, sObjectIds),
+        "The Retrieve returned a null result: " + sObjectName,
+        Objects::nonNull
+    );
+  }
+
+  /**
+   * Retrieves user info from the Salesforce connection with retry logic.
+   *
+   * @param connection the Salesforce PartnerConnection
+   * @return GetUserInfoResult containing user and organization info
+   * @throws ConnectionException if the getUserInfo call fails and is not retryable
+   */
+  public static GetUserInfoResult getUserInfoWithRetry(PartnerConnection connection)
+      throws ConnectionException {
+    return executeWithRetry(
+        SalesforceSplitUtil.getDefaultRetryPolicy(),
+        () -> connection.getUserInfo(),
+        "UserInfo or OrganizationId is null",
+        result -> result != null && result.getOrganizationId() != null
+    );
   }
 
 }

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

@@ -259,13 +259,6 @@ private static OperationEnum getOperationEnum(String operation) {
     }
   }
 
-  public static RetryPolicy<Object> getRetryPolicy(SalesforceSourceConfig config) {
-    return getRetryPolicy(
-            config.getInitialRetryDuration(),
-            config.getMaxRetryDuration(),
-            config.getMaxRetryCount());
-  }
-
   public static RetryPolicy<Object> getRetryPolicy(Long initialRetryDuration, Long maxRetryDuration,
                                                     Integer maxRetryCount) {
     // Exponential backoff with initial retry of 5 seconds and max retry of 80 seconds.