TestMorphlineSolrSink.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.flume.sink.solr.morphline;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.flume.Channel;
import org.apache.flume.ChannelSelector;
import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDeliveryException;
import org.apache.flume.channel.ChannelProcessor;
import org.apache.flume.channel.MemoryChannel;
import org.apache.flume.channel.ReplicatingChannelSelector;
import org.apache.flume.conf.Configurables;
import org.apache.flume.event.EventBuilder;
import org.apache.solr.SolrTestCaseJ4;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.kitesdk.morphline.api.MorphlineContext;
import org.kitesdk.morphline.api.Record;
import org.kitesdk.morphline.base.FaultTolerance;
import org.kitesdk.morphline.base.Fields;
import org.kitesdk.morphline.solr.DocumentLoader;
import org.kitesdk.morphline.solr.SolrLocator;
import org.kitesdk.morphline.solr.SolrMorphlineContext;
import org.kitesdk.morphline.solr.SolrServerDocumentLoader;
import org.kitesdk.morphline.solr.TestEmbeddedSolrServer;
import com.codahale.metrics.MetricRegistry;
import com.google.common.base.Charsets;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.io.Files;

public class TestMorphlineSolrSink extends SolrTestCaseJ4 {

  private EmbeddedSource source;
  private SolrServer solrServer;
  private MorphlineSink sink;
  private Map<String,Integer> expectedRecords;

  private File tmpFile;
  private static final boolean TEST_WITH_EMBEDDED_SOLR_SERVER = true;
  private static final String EXTERNAL_SOLR_SERVER_URL = System.getProperty("externalSolrServer");
//private static final String EXTERNAL_SOLR_SERVER_URL = "http://127.0.0.1:8983/solr";
  private static final String RESOURCES_DIR = "target/test-classes";
//private static final String RESOURCES_DIR = "src/test/resources";
  private static final AtomicInteger SEQ_NUM = new AtomicInteger();
  private static final AtomicInteger SEQ_NUM2 = new AtomicInteger();
  private static final Logger LOGGER = LoggerFactory.getLogger(TestMorphlineSolrSink.class);

  @BeforeClass
  public static void beforeClass() throws Exception {
    initCore(
        RESOURCES_DIR + "/solr/collection1/conf/solrconfig.xml", 
        RESOURCES_DIR + "/solr/collection1/conf/schema.xml",
        RESOURCES_DIR + "/solr");
  }

  @Before
  @Override
  public void setUp() throws Exception {
    super.setUp();
    String path = RESOURCES_DIR + "/test-documents";
    expectedRecords = new HashMap();
    expectedRecords.put(path + "/sample-statuses-20120906-141433.avro", 2);
    expectedRecords.put(path + "/sample-statuses-20120906-141433", 2);
    expectedRecords.put(path + "/sample-statuses-20120906-141433.gz", 2);
    expectedRecords.put(path + "/sample-statuses-20120906-141433.bz2", 2);
    expectedRecords.put(path + "/cars.csv", 5);
    expectedRecords.put(path + "/cars.csv.gz", 5);
    expectedRecords.put(path + "/cars.tar.gz", 4);
    expectedRecords.put(path + "/cars.tsv", 5);
    expectedRecords.put(path + "/cars.ssv", 5);

    final Map<String, String> context = new HashMap();
    
    if (EXTERNAL_SOLR_SERVER_URL != null) {
      throw new UnsupportedOperationException();
      //solrServer = new ConcurrentUpdateSolrServer(EXTERNAL_SOLR_SERVER_URL, 2, 2);
      //solrServer = new SafeConcurrentUpdateSolrServer(EXTERNAL_SOLR_SERVER_URL, 2, 2);
      //solrServer = new HttpSolrServer(EXTERNAL_SOLR_SERVER_URL);
    } else {
      if (TEST_WITH_EMBEDDED_SOLR_SERVER) {
        solrServer = new TestEmbeddedSolrServer(h.getCoreContainer(), "");
      } else {
        throw new RuntimeException("Not yet implemented");
        //solrServer = new TestSolrServer(getSolrServer());
      }
    }

    Map<String, String> channelContext = new HashMap();
    channelContext.put("capacity", "1000000");
    channelContext.put("keep-alive", "0"); // for faster tests
    Channel channel = new MemoryChannel();
    channel.setName(channel.getClass().getName() + SEQ_NUM.getAndIncrement());
    Configurables.configure(channel, new Context(channelContext));
 
    class MySolrSink extends MorphlineSolrSink {
      public MySolrSink(MorphlineHandlerImpl indexer) {
        super(indexer);
      }
    }
    
    int batchSize = SEQ_NUM2.incrementAndGet() % 2 == 0 ? 100 : 1;
    DocumentLoader testServer = new SolrServerDocumentLoader(solrServer, batchSize);
    MorphlineContext solrMorphlineContext = new SolrMorphlineContext.Builder()
        .setDocumentLoader(testServer)
        .setExceptionHandler(new FaultTolerance(false, false, SolrServerException.class.getName()))
        .setMetricRegistry(new MetricRegistry()).build();
    
    MorphlineHandlerImpl impl = new MorphlineHandlerImpl();
    impl.setMorphlineContext(solrMorphlineContext);
    
    class MySolrLocator extends SolrLocator { // trick to access protected ctor
      public MySolrLocator(MorphlineContext indexer) {
        super(indexer);
      }
    }

    SolrLocator locator = new MySolrLocator(solrMorphlineContext);
    locator.setSolrHomeDir(testSolrHome + "/collection1");
    String str1 = "SOLR_LOCATOR : " + locator.toString();
    //File solrLocatorFile = new File("target/test-classes/test-morphlines/solrLocator.conf");
    //String str1 = Files.toString(solrLocatorFile, Charsets.UTF_8);
    File morphlineFile = new File("target/test-classes/test-morphlines/solrCellDocumentTypes.conf");
    String str2 = Files.toString(morphlineFile, Charsets.UTF_8);
    tmpFile = File.createTempFile("morphline", ".conf");
    tmpFile.deleteOnExit();
    Files.write(str1 + "\n" + str2, tmpFile, Charsets.UTF_8);    
    context.put("morphlineFile", tmpFile.getPath());

    impl.configure(new Context(context));
    sink = new MySolrSink(impl);
    sink.setName(sink.getClass().getName() + SEQ_NUM.getAndIncrement());
    sink.configure(new Context(context));
    sink.setChannel(channel);
    sink.start();
    
    source = new EmbeddedSource(sink);    
    ChannelSelector rcs = new ReplicatingChannelSelector();
    rcs.setChannels(Collections.singletonList(channel));
    ChannelProcessor chp = new ChannelProcessor(rcs);
    Context chpContext = new Context();
    chpContext.put("interceptors", "uuidinterceptor");
    chpContext.put("interceptors.uuidinterceptor.type", UUIDInterceptor.Builder.class.getName());
    chp.configure(chpContext);
    source.setChannelProcessor(chp);
    
    deleteAllDocuments();
  }
  
  private void deleteAllDocuments() throws SolrServerException, IOException {
    SolrServer s = solrServer;
    s.deleteByQuery("*:*"); // delete everything!
    s.commit();
  }

  @After
  @Override
  public void tearDown() throws Exception {
    try {
      if (source != null) {
        source.stop();
        source = null;
      }
      if (sink != null) {
        sink.stop();
        sink = null;
      }
      if (tmpFile != null) {
        tmpFile.delete();
      }
    } finally {
      solrServer = null;
      expectedRecords = null;
      super.tearDown();
    }
  }

  @Test
  public void testDocumentTypes() throws Exception {
    String path = RESOURCES_DIR + "/test-documents";
    String[] files = new String[] {
        path + "/testBMPfp.txt",
        path + "/boilerplate.html",
        path + "/NullHeader.docx",
        path + "/testWORD_various.doc",          
        path + "/testPDF.pdf",
        path + "/testJPEG_EXIF.jpg",
        path + "/testXML.xml",          
//        path + "/cars.csv",
//        path + "/cars.tsv",
//        path + "/cars.ssv",
//        path + "/cars.csv.gz",
//        path + "/cars.tar.gz",
        path + "/sample-statuses-20120906-141433.avro",
        path + "/sample-statuses-20120906-141433",
        path + "/sample-statuses-20120906-141433.gz",
        path + "/sample-statuses-20120906-141433.bz2",
    };
    testDocumentTypesInternal(files);
  }

  @Test
  public void testDocumentTypes2() throws Exception {
    String path = RESOURCES_DIR + "/test-documents";
    String[] files = new String[] {
        path + "/testPPT_various.ppt",
        path + "/testPPT_various.pptx",        
        path + "/testEXCEL.xlsx",
        path + "/testEXCEL.xls", 
        path + "/testPages.pages", 
        path + "/testNumbers.numbers", 
        path + "/testKeynote.key",
        
        path + "/testRTFVarious.rtf", 
        path + "/complex.mbox", 
        path + "/test-outlook.msg", 
        path + "/testEMLX.emlx",
//        path + "/testRFC822",  
        path + "/rsstest.rss", 
//        path + "/testDITA.dita", 
        
        path + "/testMP3i18n.mp3", 
        path + "/testAIFF.aif", 
        path + "/testFLAC.flac", 
//        path + "/testFLAC.oga", 
//        path + "/testVORBIS.ogg",  
        path + "/testMP4.m4a", 
        path + "/testWAV.wav", 
//        path + "/testWMA.wma", 
        
        path + "/testFLV.flv", 
//        path + "/testWMV.wmv", 
        
        path + "/testBMP.bmp", 
        path + "/testPNG.png", 
        path + "/testPSD.psd",        
        path + "/testSVG.svg",  
        path + "/testTIFF.tif",     

//        path + "/test-documents.7z", 
//        path + "/test-documents.cpio",
//        path + "/test-documents.tar", 
//        path + "/test-documents.tbz2", 
//        path + "/test-documents.tgz",
//        path + "/test-documents.zip",
//        path + "/test-zip-of-zip.zip",
//        path + "/testJAR.jar",
        
//        path + "/testKML.kml", 
//        path + "/testRDF.rdf", 
        path + "/testTrueType.ttf", 
        path + "/testVISIO.vsd",
//        path + "/testWAR.war", 
//        path + "/testWindows-x86-32.exe",
//        path + "/testWINMAIL.dat", 
//        path + "/testWMF.wmf", 
    };   
    testDocumentTypesInternal(files);
  }

  @Test
  public void testAvroRoundTrip() throws Exception {
    String file = RESOURCES_DIR + "/test-documents" + "/sample-statuses-20120906-141433.avro";
    testDocumentTypesInternal(file);
    QueryResponse rsp = query("*:*");
    Iterator<SolrDocument> iter = rsp.getResults().iterator();
    ListMultimap<String, String> expectedFieldValues;
    expectedFieldValues = ImmutableListMultimap.of("id", "1234567890", "text", "sample tweet one",
                                                   "user_screen_name", "fake_user1");
    assertEquals(expectedFieldValues, next(iter));
    expectedFieldValues = ImmutableListMultimap.of("id", "2345678901", "text", "sample tweet two",
                                                   "user_screen_name", "fake_user2");
    assertEquals(expectedFieldValues, next(iter));
    assertFalse(iter.hasNext());
  }
  
  private ListMultimap<String, Object> next(Iterator<SolrDocument> iter) {
    SolrDocument doc = iter.next();
    Record record = toRecord(doc);
    record.removeAll("_version_"); // the values of this field are unknown and internal to solr
    return record.getFields();    
  }
  
  private Record toRecord(SolrDocument doc) {
    Record record = new Record();
    for (String key : doc.keySet()) {
      record.getFields().replaceValues(key, doc.getFieldValues(key));        
    }
    return record;
  }
  
  private void testDocumentTypesInternal(String... files) throws Exception {
    int numDocs = 0;
    long startTime = System.currentTimeMillis();
    
    assertEquals(numDocs, queryResultSetSize("*:*"));      
//  assertQ(req("*:*"), "//*[@numFound='0']");
    for (int i = 0; i < 1; i++) {      
      for (String file : files) {
        File f = new File(file);
        byte[] body = Files.toByteArray(f);
        Event event = EventBuilder.withBody(body);
        event.getHeaders().put(Fields.ATTACHMENT_NAME, f.getName());
        load(event);
        Integer count = expectedRecords.get(file);
        if (count != null) {
          numDocs += count;
        } else {
          numDocs++;
        }
        assertEquals(numDocs, queryResultSetSize("*:*"));
      }
      LOGGER.trace("iter: {}", i);
    }
    LOGGER.trace("all done with put at {}", System.currentTimeMillis() - startTime);
    assertEquals(numDocs, queryResultSetSize("*:*"));
    LOGGER.trace("sink: ", sink);
  }

//  @Test
  public void benchmarkDocumentTypes() throws Exception {
    int iters = 200;
    
//    LogManager.getLogger(getClass().getPackage().getName()).setLevel(Level.INFO);
    
    assertEquals(0, queryResultSetSize("*:*"));      
    String path = RESOURCES_DIR + "/test-documents";
    String[] files = new String[] {
//        path + "/testBMPfp.txt",
//        path + "/boilerplate.html",
//        path + "/NullHeader.docx",
//        path + "/testWORD_various.doc",          
//        path + "/testPDF.pdf",
//        path + "/testJPEG_EXIF.jpg",
//        path + "/testXML.xml",          
//        path + "/cars.csv",
//        path + "/cars.csv.gz",
//        path + "/cars.tar.gz",
//        path + "/sample-statuses-20120906-141433.avro",
        path + "/sample-statuses-20120906-141433-medium.avro",
    };
    
    List<Event> events = new ArrayList();
    for (String file : files) {
      File f = new File(file);
      byte[] body = Files.toByteArray(f);
      Event event = EventBuilder.withBody(body);
//      event.getHeaders().put(Metadata.RESOURCE_NAME_KEY, f.getName());
      events.add(event);
    }
    
    long startTime = System.currentTimeMillis();
    for (int i = 0; i < iters; i++) {
      if (i % 10000 == 0) {
        LOGGER.info("iter: {}", i);
      }
      for (Event event : events) {
        event = EventBuilder.withBody(event.getBody(), new HashMap(event.getHeaders()));
        event.getHeaders().put("id", UUID.randomUUID().toString());
        load(event);
      }
    }
    
    float secs = (System.currentTimeMillis() - startTime) / 1000.0f;
    long numDocs = queryResultSetSize("*:*");
    LOGGER.info("Took secs: " + secs + ", iters/sec: " + (iters / secs));
    LOGGER.info("Took secs: " + secs + ", docs/sec: " + (numDocs / secs));
    LOGGER.info("Iterations: " + iters + ", numDocs: " + numDocs);
    LOGGER.info("sink: ", sink);
  }

  private void load(Event event) throws EventDeliveryException {
    source.load(event);
  }

  private void commit() throws SolrServerException, IOException {
    solrServer.commit(false, true, true);
  }
  
  private int queryResultSetSize(String query) throws SolrServerException, IOException {
    commit();
    QueryResponse rsp = query(query);
    LOGGER.debug("rsp: {}", rsp);
    int size = rsp.getResults().size();
    return size;
  }
  
  private QueryResponse query(String query) throws SolrServerException, IOException {
    commit();
    QueryResponse rsp = solrServer.query(new SolrQuery(query).setRows(Integer.MAX_VALUE));
    LOGGER.debug("rsp: {}", rsp);
    return rsp;
  }
  
}