From ff1385b5022b9b24f42dcb9051c35c78c921a345 Mon Sep 17 00:00:00 2001
From: ashitsalesforce <ashit.patel@salesforce.com>
Date: Sat, 13 Jan 2024 22:28:30 -0800
Subject: [PATCH] modularize code

create helper methods for loadBatch() to make it easier to set additional headers in the future.
---
 .../action/visitor/PartnerLoadVisitor.java    | 64 +++++++++++--------
 1 file changed, 37 insertions(+), 27 deletions(-)

diff --git a/src/main/java/com/salesforce/dataloader/action/visitor/PartnerLoadVisitor.java b/src/main/java/com/salesforce/dataloader/action/visitor/PartnerLoadVisitor.java
index a9d6a553..095992c3 100644
--- a/src/main/java/com/salesforce/dataloader/action/visitor/PartnerLoadVisitor.java
+++ b/src/main/java/com/salesforce/dataloader/action/visitor/PartnerLoadVisitor.java
@@ -61,22 +61,7 @@ public PartnerLoadVisitor(Controller controller, ILoaderProgress monitor, DataWr
     @Override
     protected void loadBatch() throws DataAccessObjectException, LoadException {
         Object[] results = null;
-        Config config = this.controller.getConfig();
-        OwnerChangeOption keepAccountTeamOption = new OwnerChangeOption();
-        OwnerChangeOption[] ownerChangeOptionArray;
-        if (config.getBoolean(Config.PROCESS_KEEP_ACCOUNT_TEAM) 
-                && "Account".equalsIgnoreCase(config.getString(Config.ENTITY))) {
-            // Support for Keeping Account keepAccountTeam during Account ownership change
-            // More details at https://developer.salesforce.com/docs/atlas.en-us.api.meta/api/sforce_api_header_ownerchangeoptions.htm
-            keepAccountTeamOption.setExecute(true);
-            keepAccountTeamOption.setType(OwnerChangeOptionType.KeepAccountTeam); // Transfer Open opportunities owned by the account's owner
-            ownerChangeOptionArray = new OwnerChangeOption[] {keepAccountTeamOption};
-        } else {
-            // clear ownerChangeOptions from the existing connection otherwise.
-            ownerChangeOptionArray = new OwnerChangeOption[] {};
-        }
-        this.controller.getPartnerClient().getClient().setOwnerChangeOptions(ownerChangeOptionArray);
-        
+        setHeaders();
         try {
             results = executeClientAction(getController().getPartnerClient(), dynaArray);
         } catch (ApiFault e) {
@@ -85,7 +70,20 @@ protected void loadBatch() throws DataAccessObjectException, LoadException {
             handleException(e);
         }
 
-        // set the current processed
+        writeOutputToWriter(results);
+        setLastRunProperties(results);
+
+        // update Monitor
+        getProgressMonitor().worked(results.length);
+        getProgressMonitor().setSubTask(getRateCalculator().calculateSubTask(getNumberOfRows(), getNumberErrors()));
+
+        // now clear the arrays
+        clearArrays();
+
+    }
+    
+    private void setLastRunProperties(Object[] results) throws LoadException {
+        // set the last processed row number in the config (*_lastRun.properties) file
         int currentProcessed;
         try {
             currentProcessed = getConfig().getInt(LastRun.LAST_LOAD_BATCH_ROW);
@@ -102,16 +100,28 @@ protected void loadBatch() throws DataAccessObjectException, LoadException {
             getLogger().error(errMsg, e);
             handleException(errMsg, e);
         }
-
-        writeOutputToWriter(results);
-
-        // update Monitor
-        getProgressMonitor().worked(results.length);
-        getProgressMonitor().setSubTask(getRateCalculator().calculateSubTask(getNumberOfRows(), getNumberErrors()));
-
-        // now clear the arrays
-        clearArrays();
-
+    }
+    
+    private void setHeaders() {
+        setKeepAccountTeamHeader();
+    }
+    
+    private void setKeepAccountTeamHeader() {
+        Config config = this.controller.getConfig();
+        OwnerChangeOption keepAccountTeamOption = new OwnerChangeOption();
+        OwnerChangeOption[] ownerChangeOptionArray;
+        if (config.getBoolean(Config.PROCESS_KEEP_ACCOUNT_TEAM) 
+                && "Account".equalsIgnoreCase(config.getString(Config.ENTITY))) {
+            // Support for Keeping Account keepAccountTeam during Account ownership change
+            // More details at https://developer.salesforce.com/docs/atlas.en-us.api.meta/api/sforce_api_header_ownerchangeoptions.htm
+            keepAccountTeamOption.setExecute(true);
+            keepAccountTeamOption.setType(OwnerChangeOptionType.KeepAccountTeam); // Transfer Open opportunities owned by the account's owner
+            ownerChangeOptionArray = new OwnerChangeOption[] {keepAccountTeamOption};
+        } else {
+            // clear ownerChangeOptions from the existing connection otherwise.
+            ownerChangeOptionArray = new OwnerChangeOption[] {};
+        }
+        this.controller.getPartnerClient().getClient().setOwnerChangeOptions(ownerChangeOptionArray);
     }
 
     private void writeOutputToWriter(Object[] results)