Extending the default Dataload search index mediator

package com.mycompany.commerce.foundation.dataimport.dataload.mediator;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.StringTokenizer;

import org.apache.solr.common.SolrInputDocument;

import com.ibm.commerce.foundation.common.util.logging.LoggingHelper;
import com.ibm.commerce.foundation.dataimport.dataload.mediator.AbstractSolrInputDocumentMediator;
import com.ibm.commerce.foundation.dataimport.dataload.object.SolrDocumentDataObject;
import com.ibm.commerce.foundation.dataload.DataLoadConstants;
import com.ibm.commerce.foundation.dataload.exception.DataLoadException;

/**
 * This class handles the mapping between the inventory map object to SolrInputDocument object and adds it to SolrInputDocDataObjects.
 * It calculates the product quantity value based on the its SKUs, and creates SolrDoc for the product and its SKUs.
 */
public class InventorySolrInputDocumentMediator extends
AbstractSolrInputDocumentMediator{
  /**
   * IBM Copyright notice field.
   */
  private static final String COPYRIGHT = com.ibm.commerce.copyright.IBMCopyright.SHORT_COPYRIGHT;

  /**
   * Class name used for trace and logging.
   */
  private final static String CLASSNAME = InventorySolrInputDocumentMediator.class.getName();

  /**
   * Logger used for logging.
   */
  private static final java.util.logging.Logger LOGGER = LoggingHelper
      .getLogger(InventorySolrInputDocumentMediator.class);
  
  private String CATENTRY_ID_SOLR_FIELD_NAME="catentry_id";  
  private String QUANTITY_SOLR_FIELD_NAME="inv_strlocqty_";
  
  
  /**
   * Transform the specified map object to SolrInputDocument object.
   * @param dataObject, the actual map object as passed from the builder.
   * @param deleteFlag if it is delete operation.
   * @throws DataLoadException
   */
  protected void transform(Object dataObject, boolean deleteFlag)throws DataLoadException {
      final String METHODNAME = "transform(Object, boolean)";
    if(LoggingHelper.isTraceEnabled(LOGGER)) {
      LOGGER.entering(CLASSNAME, METHODNAME, new Object[]{dataObject, deleteFlag});
    }  
    
    String PRODUCTCATENTRYID_DATA_NAME="productCatentryId";
    String ITEMCATENTRYID_DATA_NAME="itemCatentryId";
    String ITEMSTOREQUANTITY_DATA_NAME="store_location_quantity";
    
    if (dataObject instanceof Map) {
      Map<String, Object> data = (Map<String, Object>) dataObject;
      ArrayList productCatentryIdValue= null;  
      ArrayList itemsCatEntryIdValue=null;
      ArrayList itemsStoreQuantityValue=null;      
      String catentryId=null;
          
      try{              
        //get data from data file, Assuming that the loadItem 'valueIsList' config property is set to true, 
        //therefore all Objects are returned as an ArrayList
        if (data != null && !data.isEmpty()) {            
          productCatentryIdValue = ((ArrayList)data.get(PRODUCTCATENTRYID_DATA_NAME));
          itemsCatEntryIdValue = (ArrayList)data.get(ITEMCATENTRYID_DATA_NAME);
          itemsStoreQuantityValue = (ArrayList)data.get(ITEMSTOREQUANTITY_DATA_NAME);        
        }
        
        Map<String, String> productStoreQuantityMap= new HashMap();
        // Ensure not null, and the number of items equals the number of quantities/stores
        if(itemsCatEntryIdValue!=null && itemsStoreQuantityValue!=null  
            && itemsCatEntryIdValue.size()==itemsStoreQuantityValue.size()){            
          
          for(int i=0;i<itemsCatEntryIdValue.size();i++){
            //create the item Solr document
            catentryId=itemsCatEntryIdValue.get(i).toString();
            StringTokenizer st = new StringTokenizer(itemsStoreQuantityValue.get(i).toString()," ");
            String storeQuantity = null;
            String storeId_id=null;
            String quantityValue=null;
            Map<String, String> itemStoreQuantityMap= new HashMap();
            Map<String, String> tempProductStoreQuantityMap= new HashMap();
            while(st.hasMoreElements()){
              storeQuantity = st.nextToken();
              StringTokenizer innerST = new StringTokenizer(storeQuantity,":");
              storeId_id = innerST.nextToken();
              quantityValue = innerST.nextToken();
              String storeQuantityFieldName = QUANTITY_SOLR_FIELD_NAME+storeId_id;
              itemStoreQuantityMap.put(storeQuantityFieldName, quantityValue);  
              tempProductStoreQuantityMap.put(storeQuantityFieldName, quantityValue);
            }  
            SolrDocumentDataObject solrDocDO = createSolrDocumentDataObject(catentryId, itemStoreQuantityMap, deleteFlag);
            if(solrDocDO != null){
            	addSolrInputDocDataObjects(solrDocDO);
            }
            
            //handles product level quantity
            Iterator<String> iter = tempProductStoreQuantityMap.keySet().iterator();
            while (iter.hasNext()) {
              String fieldName = iter.next();
              Float quantity = Float.valueOf(tempProductStoreQuantityMap.get(fieldName));  
              if(productStoreQuantityMap.containsKey(fieldName))  {                
                quantity +=Float.valueOf(productStoreQuantityMap.get(fieldName));
                productStoreQuantityMap.put(fieldName, quantity.toString());
              }
              else
                productStoreQuantityMap.put(fieldName, quantity.toString());
            }        
          }          
        }
        
        //Ensure not null, and there is at lease one productPartNumber returned in the ArrayList
        if(productCatentryIdValue !=null && productCatentryIdValue.size()>0)
        {
          //create the product Solr document
          catentryId=productCatentryIdValue.get(0).toString();
          SolrDocumentDataObject solrDocDO = createSolrDocumentDataObject(catentryId, productStoreQuantityMap, deleteFlag);
          if(solrDocDO != null){
          	addSolrInputDocDataObjects(solrDocDO);
          }
        }
      }catch (Exception ex)
      {
        throw new DataLoadException();
      }    
    }
    else{
    	LOGGER.info("The dataObject is not an instance of Map");
    }
    
    if(LoggingHelper.isEntryExitTraceEnabled(LOGGER)) {
      LOGGER.exiting(CLASSNAME, METHODNAME);
    }    
  }
    
  /**
   * This method creates new SolrInputDocument and returns it to the calling method
   * @param catentryIdVal
   *       the catalog entry internal id value
   * @param storeQuantityMap
   *       the store-quantity name-value-pair
   * @param deleteFlag
   *       the solr delete flag
   */
  private SolrDocumentDataObject createSolrDocumentDataObject(String catentryIdVal,Map<String, String> storeQuantityMap,boolean deleteFlag){
    final String METHODNAME = "createSolrDocumentDataObject";
    if (LoggingHelper.isEntryExitTraceEnabled(LOGGER)) {
      LOGGER.entering(CLASSNAME, METHODNAME, new Object[] { catentryIdVal, storeQuantityMap,deleteFlag});
    }
    SolrDocumentDataObject solrDocDO = null;
    if(catentryIdVal !=null && storeQuantityMap!=null){
      SolrInputDocument doc= new SolrInputDocument();
              
      String idFieldName = getSolrIdFieldName();    
      if (idFieldName != null) {        
        doc.addField(CATENTRY_ID_SOLR_FIELD_NAME, catentryIdVal);
        Iterator<String> iter = storeQuantityMap.keySet().iterator();
        while (iter.hasNext()) {
          String fieldName = iter.next();
          String fieldValue = storeQuantityMap.get(fieldName);
          doc.addField(fieldName, fieldValue);
        }
      }  
              
      solrDocDO = new SolrDocumentDataObject(idFieldName, doc);    
      if (deleteFlag) {
        solrDocDO.setOpMode(DataLoadConstants.DL_OP_DELETE);
      } else {
        solrDocDO.setOpMode(DataLoadConstants.DL_OP_UPDATE);
      }
    }
    else{ 
      LOGGER.info("The value of catenryId is missing!");
    }
    
    if (LoggingHelper.isEntryExitTraceEnabled(LOGGER)) {
      LOGGER.exiting(CLASSNAME, METHODNAME,  new Object[] { solrDocDO});
    }
    return solrDocDO;
  }
}