refactor common code in DatabaseWriter and CSVFileWriter

Refactor common code in DatabaseWriter and CSVFileWriter by introducing a common interface for Row and TableRow classes.
forcedotcom · Dec 5, 2024 · c8204d1 · c8204d1
1 parent ceb1fd8
commit c8204d1
Show file tree

Hide file tree

Showing 16 changed files with 105 additions and 128 deletions.
diff --git a/src/main/java/com/salesforce/dataloader/action/visitor/AbstractVisitor.java b/src/main/java/com/salesforce/dataloader/action/visitor/AbstractVisitor.java
@@ -139,7 +139,7 @@ protected void writeSuccess(TableRow row, String id, String message) throws Data
             if (message != null && message.length() > 0) {
                 row.put(AppConfig.STATUS_COLUMN_NAME, message);
             }
-            this.successWriter.writeTableRow(row);
+            this.successWriter.writeRow(row);
         }
         addSuccess();
     }
@@ -151,7 +151,7 @@ protected void writeError(TableRow row, String errorMessage) throws DataAccessOb
             } else {
                 row.put(AppConfig.ERROR_COLUMN_NAME, errorMessage);
             }
-            this.errorWriter.writeTableRow(row);
+            this.errorWriter.writeRow(row);
         }
         addErrors();
     }

diff --git a/src/main/java/com/salesforce/dataloader/action/visitor/DAOLoadVisitor.java b/src/main/java/com/salesforce/dataloader/action/visitor/DAOLoadVisitor.java
@@ -205,7 +205,7 @@ public boolean visit(TableRow row) throws OperationException, DataAccessObjectEx
         }
         try {
             convertBulkAPINulls(sforceDataRow);
-            DynaBean dynaBean = SforceDynaBean.convertToDynaBean(dynaClass, sforceDataRow.convertToRow());
+            DynaBean dynaBean = SforceDynaBean.convertToDynaBean(dynaClass, sforceDataRow);
             Map<String, String> fieldMap = BeanUtils.describe(dynaBean);
             for (String fName : fieldMap.keySet()) {
                 if (fieldMap.get(fName) != null) {

diff --git a/src/main/java/com/salesforce/dataloader/action/visitor/partner/PartnerQueryVisitor.java b/src/main/java/com/salesforce/dataloader/action/visitor/partner/PartnerQueryVisitor.java
@@ -105,7 +105,7 @@ private Row getDaoRow(SObject sob, boolean firstRowInBatch) {
             Row row = getMapper().mapPartnerSObjectSfdcToLocal(sob);
             try {
                 List<String> queryResultFieldsList;
-                queryResultFieldsList = ((DataWriter)controller.getDao()).getColumnNamesFromRow(row);
+                queryResultFieldsList =  row.getColumnNames();
                 SOQLMapper mapper = (SOQLMapper)this.controller.getMapper();
                 mapper.initSoqlMappingFromResultFields(queryResultFieldsList);
                 final List<String> daoColumns = mapper.getDaoColumnsForSoql();

diff --git a/src/main/java/com/salesforce/dataloader/config/ConfigPropertyMetadata.java b/src/main/java/com/salesforce/dataloader/config/ConfigPropertyMetadata.java
@@ -38,6 +38,7 @@
 import com.salesforce.dataloader.dao.csv.CSVFileWriter;
 import com.salesforce.dataloader.exception.DataAccessObjectException;
 import com.salesforce.dataloader.exception.DataAccessObjectInitializationException;
+import com.salesforce.dataloader.model.RowInterface;
 import com.salesforce.dataloader.model.TableHeader;
 import com.salesforce.dataloader.model.TableRow;
 import com.salesforce.dataloader.ui.Labels;
@@ -246,7 +247,7 @@ public static void generateCSV(AppConfig appConfig) {
             headerLabelList.add(COL_IS_READ_ONLY);
             headerLabelList.add(COL_IS_COMMAND_LINE_OPTION);
             headerLabelList.add(COL_IS_ENCRYPTED);
-            ArrayList<TableRow> rowList = new ArrayList<TableRow>(propertiesMap.size());
+            ArrayList<RowInterface> rowList = new ArrayList<RowInterface>(propertiesMap.size());
             TableHeader header = new TableHeader(headerLabelList);
 
             for (ConfigPropertyMetadata propMD : propertiesMap.values()) {
@@ -268,7 +269,7 @@ public static void generateCSV(AppConfig appConfig) {
                 rowList.add(row);
             }
             try {
-                csvWriter.writeTableRowList(rowList);
+                csvWriter.writeRowList(rowList);
             } catch (DataAccessObjectException e) {
                 logger.warn(e.getStackTrace());
             }

diff --git a/src/main/java/com/salesforce/dataloader/dao/DataWriter.java b/src/main/java/com/salesforce/dataloader/dao/DataWriter.java
@@ -30,8 +30,7 @@
 
 import com.salesforce.dataloader.exception.DataAccessObjectException;
 import com.salesforce.dataloader.exception.DataAccessObjectInitializationException;
-import com.salesforce.dataloader.model.Row;
-import com.salesforce.dataloader.model.TableRow;
+import com.salesforce.dataloader.model.RowInterface;
 
 /**
  * Interface to be implemented for data writers -- data access objects that are used for writing rows of data.
@@ -54,17 +53,12 @@ public interface DataWriter extends DataAccessObject {
      * @return Any data columns generated as a result of writing
      * @throws DataAccessObjectException
      */
-    boolean writeRow(Row inputRow) throws DataAccessObjectException;
-    boolean writeTableRow(TableRow inputRow) throws DataAccessObjectException;
+    boolean writeRow(RowInterface inputRow) throws DataAccessObjectException;
 
     /**
      * @param inputRowList
      * @return List of data rows with generated data columns as a result of writing
      * @throws DataAccessObjectException
      */
-    boolean writeRowList(List<Row> inputRowList) throws DataAccessObjectException;
-
-    public List<String> getColumnNamesFromRow(Row row) throws DataAccessObjectInitializationException;
-
-    public boolean writeTableRowList(List<TableRow> rows) throws DataAccessObjectException;
+    boolean writeRowList(List<? extends RowInterface> inputRowList) throws DataAccessObjectException;
 }
diff --git a/src/main/java/com/salesforce/dataloader/dao/csv/CSVFileWriter.java b/src/main/java/com/salesforce/dataloader/dao/csv/CSVFileWriter.java
@@ -33,7 +33,6 @@
 import java.io.OutputStreamWriter;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Set;
 
 import org.apache.logging.log4j.Logger;
 import com.salesforce.dataloader.util.DLLogManager;
@@ -43,8 +42,7 @@
 import com.salesforce.dataloader.dao.DataWriter;
 import com.salesforce.dataloader.exception.DataAccessObjectException;
 import com.salesforce.dataloader.exception.DataAccessObjectInitializationException;
-import com.salesforce.dataloader.model.Row;
-import com.salesforce.dataloader.model.TableRow;
+import com.salesforce.dataloader.model.RowInterface;
 import com.salesforce.dataloader.util.AppUtil;
 
 /**
@@ -177,9 +175,9 @@ private void writeHeaderRow() throws DataAccessObjectInitializationException {
      * @see com.salesforce.dataloader.dao.csv.Writer#writeRow(java.util.Map)
      */
     @Override
-    public boolean writeRow(Row row) throws DataAccessObjectException {
+    public boolean writeRow(RowInterface row) throws DataAccessObjectException {
         if (this.columnNames == null || this.columnNames.isEmpty()) {
-           List<String>colNames = getColumnNamesFromRow(row);
+           List<String>colNames = row.getColumnNames();
            this.setColumnNames(colNames);
         }
         CSVColumnVisitor visitor = new CSVColumnVisitor(fileOut, false, this.columnDelimiter);
@@ -194,66 +192,21 @@ public boolean writeRow(Row row) throws DataAccessObjectException {
             throw new DataAccessObjectException(Messages.getString("CSVWriter.errorWriting"), e); //$NON-NLS-1$
         }
     }
-
-    /*
-     * (non-Javadoc)
-     * @see com.salesforce.dataloader.dao.csv.Writer#writeRow(java.util.Map)
-     */
-    public boolean writeTableRow(TableRow row) throws DataAccessObjectException {
-        if (this.columnNames == null || this.columnNames.isEmpty()) {
-           List<String>colNames = getColumnNamesFromTableRow(row);
-           this.setColumnNames(colNames);
-        }
-        CSVColumnVisitor visitor = new CSVColumnVisitor(fileOut, false, this.columnDelimiter);
-        try {
-            visitTableRowColumns(columnNames, row, visitor);
-            fileOut.newLine();
-            visitor.newRow();
-            currentRowNumber++;
-            return true; // success unless there's an exception
-        } catch (IOException e) {
-            logger.error(Messages.getString("CSVWriter.errorWriting"), e); //$NON-NLS-1$
-            throw new DataAccessObjectException(Messages.getString("CSVWriter.errorWriting"), e); //$NON-NLS-1$
-        }
-    }
-
-    public List<String> getColumnNamesFromRow(Row row) throws DataAccessObjectInitializationException {
-        Set<String> fieldNameSet = row.keySet();
-        return new ArrayList<String>(fieldNameSet);    
-    }
-
-    public List<String> getColumnNamesFromTableRow(TableRow row) {
-        return row.getHeader().getColumns();    
-    }
 
     /*
      * (non-Javadoc)
      * @see com.salesforce.dataloader.dao.csv.Writer#writeRowList(java.util.List)
      */
     @Override
-    public boolean writeRowList(List<Row> rows) throws DataAccessObjectException {
+    public boolean writeRowList(List<? extends RowInterface> rows) throws DataAccessObjectException {
         boolean success = true;
         // return the last result, should be same as others
-        for (Row row : rows) {
+        for (RowInterface row : rows) {
             success = writeRow(row);
         }
         return success;
     }
 
-    /*
-     * (non-Javadoc)
-     * @see com.salesforce.dataloader.dao.csv.Writer#writeRowList(java.util.List)
-     */
-    @Override
-    public boolean writeTableRowList(List<TableRow> rows) throws DataAccessObjectException {
-        boolean success = true;
-        // return the last result, should be same as others
-        for (TableRow trow : rows) {
-            success = writeTableRow(trow);
-        }
-        return success;
-    }
-
     private void visitHeaderColumns(List<String> columnNames, CSVColumnVisitor visitor) throws IOException {
         for (String colName : columnNames) {
             String outColName;
@@ -270,7 +223,7 @@ private void visitHeaderColumns(List<String> columnNames, CSVColumnVisitor visit
         }
     }
 
-    static private void visitColumns(List<String> columnNames, Row row, CSVColumnVisitor visitor) throws IOException {
+    static private void visitColumns(List<String> columnNames, RowInterface row, CSVColumnVisitor visitor) throws IOException {
         for (String colName : columnNames) {
             Object colVal = row.get(colName);
             if (colVal == null && colName.contains("(")) {
@@ -283,19 +236,6 @@ static private void visitColumns(List<String> columnNames, Row row, CSVColumnVis
         }
     }
 
-
-    static private void visitTableRowColumns(List<String> columnNames, TableRow row, CSVColumnVisitor visitor) throws IOException {
-        for (String colName : columnNames) {
-            Object colVal = row.get(colName);
-            if (colVal == null && colName.contains("(")) {
-                int lparenIdx = colName.indexOf('(');
-                int rparenIdx = colName.indexOf(')');
-                colName = colName.substring(lparenIdx + 1, rparenIdx);
-                colVal = row.get(colName);
-            }
-            visitor.visit(colVal != null ? colVal.toString() : "");
-        }
-    }
     @Override
     public List<String> getColumnNames() {
         return columnNames;
@@ -334,5 +274,4 @@ public void setOpen(boolean open) {
     public int getCurrentRowNumber() {
         return currentRowNumber;
     }
-
 }
diff --git a/src/main/java/com/salesforce/dataloader/dao/database/DatabaseWriter.java b/src/main/java/com/salesforce/dataloader/dao/database/DatabaseWriter.java
@@ -30,6 +30,7 @@
 import java.util.*;
 
 import com.salesforce.dataloader.model.Row;
+import com.salesforce.dataloader.model.RowInterface;
 import com.salesforce.dataloader.model.TableRow;
 import org.apache.commons.dbcp2.BasicDataSource;
 import org.apache.logging.log4j.Logger;
@@ -126,7 +127,7 @@ private void setupUpdate() throws DataAccessObjectInitializationException {
      * @see com.salesforce.dataloader.dao.DataWriter#writeRowList(java.util.List)
      */
     @Override
-    public boolean writeRowList(List<Row> inputRowList) throws DataAccessObjectException {
+    public boolean writeRowList(List<? extends RowInterface> inputRowList) throws DataAccessObjectException {
 
         // make sure that the update is setup and ready to go, otherwise stop
         if (!dbContext.isOpen()) { throw new DataAccessObjectInitializationException(Messages
@@ -138,13 +139,26 @@ public boolean writeRowList(List<Row> inputRowList) throws DataAccessObjectExcep
         try {
             //for batchsize = 1, don't do batching, this provides much better error output
             if(inputRowList.size() == 1) {
-                dbContext.setSqlParamValues(sqlConfig, appConfig, inputRowList.get(0));
+                if (inputRowList.get(0) instanceof Row) {
+                    dbContext.setSqlParamValues(sqlConfig, appConfig, (Row)inputRowList.get(0));
+                } else if (inputRowList.get(0) == null) {
+                    dbContext.setSqlParamValues(sqlConfig, appConfig, null);
+                } else {
+                    dbContext.setSqlParamValues(sqlConfig, appConfig, 
+                            ((TableRow)inputRowList.get(0)).convertToRow());
+                }
                 currentRowNumber++;
             } else {
                 // for each row set the Sql params in the prepared statement
                 dbContext.getDataStatement().clearBatch();
-                for (Row inputRow : inputRowList) {
-                    dbContext.setSqlParamValues(sqlConfig, appConfig, inputRow);
+                for (RowInterface inputRow : inputRowList) {
+                    if (inputRow instanceof Row) {
+                        dbContext.setSqlParamValues(sqlConfig, appConfig, (Row)inputRow);
+                    } else if (inputRow == null) {
+                        dbContext.setSqlParamValues(sqlConfig, appConfig, null);
+                    } else {
+                        dbContext.setSqlParamValues(sqlConfig, appConfig, ((TableRow)inputRow).convertToRow());
+                    }
                     dbContext.getDataStatement().addBatch();
                     currentRowNumber++;
                 }
@@ -233,19 +247,12 @@ public boolean writeRowList(List<Row> inputRowList) throws DataAccessObjectExcep
      * @throws DataAccessObjectException
      */
     @Override
-    public boolean writeRow(Row inputRow) throws DataAccessObjectException {
+    public boolean writeRow(RowInterface inputRow) throws DataAccessObjectException {
         // FIXME: Think about refactoring this for the caller to writeRow() and here take care of batching internally
-        List<Row> inputRowList = new ArrayList<Row>();
+        List<RowInterface> inputRowList = new ArrayList<RowInterface>();
         inputRowList.add(inputRow);
         return writeRowList(inputRowList);
     }
-
-    public boolean writeTableRow(TableRow inputRow) throws DataAccessObjectException {
-        // FIXME: Think about refactoring this for the caller to writeRow() and here take care of batching internally
-        List<TableRow> inputRowList = new ArrayList<TableRow>();
-        inputRowList.add(inputRow);
-        return writeTableRowList(inputRowList);
-    }
 
     /**
      * @param sqe
@@ -294,15 +301,4 @@ public List<String> getColumnNamesFromRow(Row row) throws DataAccessObjectInitia
         return null;
     }
 
-    @Override
-    public boolean writeTableRowList(List<TableRow> trowList) throws DataAccessObjectException {
-        if (trowList == null) {
-            return false;
-        }
-        ArrayList<Row> rowList = new ArrayList<Row>(trowList.size());
-        for (TableRow trow : trowList) {
-            rowList.add(trow.convertToRow());
-        }
-        return writeRowList(rowList);
-    }
 }
diff --git a/src/main/java/com/salesforce/dataloader/dyna/SforceDynaBean.java b/src/main/java/com/salesforce/dataloader/dyna/SforceDynaBean.java
@@ -29,6 +29,7 @@
 import java.util.*;
 
 import com.salesforce.dataloader.model.Row;
+import com.salesforce.dataloader.model.TableRow;
 import com.salesforce.dataloader.util.DateOnlyCalendar;
 
 import org.apache.commons.beanutils.*;
@@ -256,6 +257,11 @@ static public DynaBean convertToDynaBean(BasicDynaClass dynaClass, Row sforceDat
             throw new LoadException(e);
         }
     }
+
+    static public DynaBean convertToDynaBean(BasicDynaClass dynaClass, TableRow sforceDataRow)
+            throws ConversionException, LoadException {
+        return convertToDynaBean(dynaClass, sforceDataRow.convertToRow());
+    }
 
     /**
      * Set all the fields specified in the dynaBean to null on the sObj

diff --git a/src/main/java/com/salesforce/dataloader/model/Row.java b/src/main/java/com/salesforce/dataloader/model/Row.java
@@ -27,8 +27,8 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -42,26 +42,14 @@
  * methods and probably stop implementing Map interface. All Row behavior should be moved into this
  * class and not be spread in multiple class.
  */
-public class Row implements Map<String, Object> {
+public class Row implements Map<String, Object>, RowInterface {
     private final Map<String, Object> internalMap;
     private final Map<String, String> keyMap = new HashMap<String, String>();
 
     public Row() {
         this.internalMap = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
     }
 
-    public Row(Map<String, Object> internalMap) {
-        this.internalMap = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
-        this.internalMap.putAll(internalMap);
-        for (String key : internalMap.keySet()) {
-            this.keyMap.put(key.toLowerCase(), key);
-        }
-    }
-
-    public static Row singleEntryImmutableRow(String key, Object value) {
-        return new Row(Collections.singletonMap(key, value));
-    }
-
     @Override
     public int size() {
         return internalMap.size();
@@ -151,4 +139,15 @@ public TableRow convertToTableRow(TableHeader header) {
         }
         return trow;
     }
+
+    @Override
+    public Object get(String key) {
+        return get((Object)key);
+    }
+
+    @Override
+    public List<String> getColumnNames() {
+        Set<String> fieldNameSet = this.keySet();
+        return new ArrayList<String>(fieldNameSet);
+    }
 }