core/org.eclipse.smila.solr/code/src/org/eclipse/smila/solr/update/SolrDocumentConverter.java

/**
 *
 */
package org.eclipse.smila.solr.update;

import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map.Entry;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.solr.common.SolrInputDocument;
import org.eclipse.smila.datamodel.Any;
import org.eclipse.smila.datamodel.AnyMap;
import org.eclipse.smila.datamodel.Record;
import org.eclipse.smila.datamodel.util.AnyUtil;
import org.eclipse.smila.solr.params.SolrParams;
import org.eclipse.smila.solr.params.UpdateParams;

/**
 * @author pwissel
 *
 */
public class SolrDocumentConverter {

  public static final String FIELD_NAME = "fieldName";

  public static final String FIELD_BOOST = "fieldBoost";

  public static final String TYPE = "type";

  private final Log _log = LogFactory.getLog(getClass());

  public enum Type {
    ATTACHMENT, ATTRIBUTE;
    public static Type get(final String label) {
      return valueOf(label.toUpperCase());
    }
  }

  public SolrInputDocument toSolrDocument(final Record record) throws UnsupportedEncodingException {
    return toSolrDocument(record, false);
  }

  public SolrInputDocument toSolrDocument(final Record record, final List<Record> children)
    throws UnsupportedEncodingException {
    return toSolrDocument(record, false, children);
  }

  public List<SolrInputDocument> toSolrDocuments(final List<Record> records) throws UnsupportedEncodingException {
    return toSolrDocuments(records, false);
  }

  public SolrInputDocument toSolrDocument(final Record record, final boolean attachments)
    throws UnsupportedEncodingException {
    final SolrInputDocument document = prepareSolrDocument(record);
    // add all metadata fields if name is not a system key (begins with underscore)
    for (final Entry<String, Any> entry : record.getMetadata().entrySet()) {
      final String name = entry.getKey();
      if (isSystemKey(name)) {
        continue;
      }
      final Any any = entry.getValue();
      if (any.isValue() || any.isSeq()) {
        final Object value = AnyUtil.anyToNative(any);
        document.addField(name, value);
      }
    }
    // add all attachments if name is not a system key (begins with underscore)
    if (attachments && record.hasAttachments()) {
      final Iterator<String> attachmentNames = record.getAttachmentNames();
      while (attachmentNames.hasNext()) {
        final String name = attachmentNames.next();
        if (isSystemKey(name)) {
          continue;
        }
        final byte[] bytes = record.getAttachmentAsBytes(name);
        document.addField(name, new String(bytes, "UTF-8"));
      }
    }
    return document;
  }

  public SolrInputDocument toSolrDocument(final Record record, final boolean attachments,
    final List<Record> children) throws UnsupportedEncodingException {
    final SolrInputDocument document = toSolrDocument(record, attachments);
    for (final Record child : children) {
      final SolrInputDocument childDocument = toSolrDocument(child, attachments);
      document.addChildDocument(childDocument);
    }
    return document;
  }

  public List<SolrInputDocument> toSolrDocuments(final List<Record> records, final boolean attachments)
    throws UnsupportedEncodingException {
    final List<SolrInputDocument> documents = new ArrayList<SolrInputDocument>(records.size());
    for (final Record record : records) {
      final SolrInputDocument document = toSolrDocument(record, attachments);
      documents.add(document);
    }
    return documents;
  }

  public SolrInputDocument toSolrDocument(final Record record, final AnyMap mapping)
    throws UnsupportedEncodingException {
    final SolrInputDocument document = prepareSolrDocument(record);
    final String id = record.getId();
    final AnyMap metadata = record.getMetadata();
    // add mapping fields
    for (final Entry<String, Any> entry : mapping.entrySet()) {
      String name = entry.getKey();
      final Any any = entry.getValue();
      // read mapping from value
      if (any.isValue()) {
        String fieldName = any.asValue().asString();
        if (StringUtils.isEmpty(fieldName)) {
          fieldName = name;
        }
        addObjectToDocument(id, metadata, name, document, fieldName);
      } else if (any.isSeq()) {
        final String[] fieldNames = any.asSeq().asStrings().toArray(new String[any.asSeq().size()]);
        addObjectToDocument(id, metadata, name, document, fieldNames);
      } else if (any.isMap()) {
        // read mapping from map
        final AnyMap map = any.asMap();
        final String fieldName = map.containsKey(FIELD_NAME) ? map.getStringValue(FIELD_NAME) : name;
        float fieldBoost = 0;
        boolean boost = false;
        if (map.containsKey(FIELD_BOOST)) {
          fieldBoost = map.getDoubleValue(FIELD_BOOST).floatValue();
          boost = true;
        }
        final Type type = map.containsKey(TYPE) ? Type.get(map.getStringValue(TYPE)) : Type.ATTRIBUTE;
        Object value = null;
        switch (type) {
          case ATTRIBUTE:
            addObjectToDocument(id, metadata, name, document, fieldName);
            break;
          case ATTACHMENT:
            // check if attachment is available
            if (record.hasAttachment(name)) {
              final byte[] bytes = record.getAttachmentAsBytes(name);
              value = new String(bytes, "UTF-8");
            } else {
              if (_log.isWarnEnabled()) {
                final String message =
                  String.format("Record does not contain attachment for given key. _recordid=%s, key=%s", id, name);
                _log.warn(message);
              }
            }
            break;
        }
        // add field to document
        if (boost) {
          document.addField(fieldName, value, fieldBoost);
        } else {
          document.addField(fieldName, value);
        }
      }
    }
    return document;
  }

  public SolrInputDocument toSolrDocument(final Record record, final AnyMap mapping, final List<Record> children)
    throws UnsupportedEncodingException {
    final SolrInputDocument document = toSolrDocument(record, mapping);
    for (final Record child : children) {
      final SolrInputDocument childDocument = toSolrDocument(child, mapping);
      document.addChildDocument(childDocument);
    }
    return document;
  }

  private void addObjectToDocument(final String id, final AnyMap metadata, final String name,
    final SolrInputDocument document, final String... fieldNames) {
    if (metadata.containsKey(name)) {
      Any val;
      // check if attribute is not a map and get value
      if (!(val = metadata.get(name)).isMap()) {
        final Object value = AnyUtil.anyToNative(val);
        for (final String fieldName : fieldNames) {
          document.addField(fieldName, value);
        }
      } else {
        if (_log.isWarnEnabled()) {
          final String message =
            String.format("Can not convert AnyMap to SolrDocument field. _recordid=%s, key=%s", id, name);
          _log.warn(message);
        }
      }
    } else {
      if (_log.isDebugEnabled()) {
        final String message =
          String.format("Record does not contain mapping for given key. _recordid=%s, key=%s", id, name);
        _log.warn(message);
      }
    }
  }

  private SolrInputDocument prepareSolrDocument(final Record record) {
    // create solr input document
    final SolrInputDocument document = new SolrInputDocument();
    final String id = record.getId();
    document.addField(Record.RECORD_ID, id);
    // set document boost if available
    final UpdateParams params = new UpdateParams(SolrParams.getSolrParams(record));
    final Float documentBoost = params.getDocumentBoost(false);
    if (documentBoost != null) {
      document.setDocumentBoost(documentBoost);
    }
    return document;
  }

  private boolean isSystemKey(final String key) {
    return StringUtils.startsWith(key, "_");
  }

}