1   package com.atlassian.bonnie.index;
2   
3   import com.atlassian.bonnie.ILuceneConnection;
4   import org.apache.lucene.document.Document;
5   import org.apache.lucene.index.IndexWriter;
6   import org.slf4j.LoggerFactory;
7   import org.slf4j.Logger;
8   
9   import java.io.IOException;
10  
11  /**
12   * Each thread has its own IndexWriter which writes to a temp directory.
13   */
14  public class TempDirectoryDocumentWritingScheme extends BaseDocumentWritingScheme
15  {
16      private static final Logger log = LoggerFactory.getLogger(TempDirectoryDocumentWritingScheme.class);
17      protected final TempIndexWriter tempIndexWriter;
18  
19      public TempDirectoryDocumentWritingScheme(BatchOpIndexer indexer)
20      {
21          this(indexer, null);
22      }
23  
24      public TempDirectoryDocumentWritingScheme(BatchOpIndexer indexer, String tmpDir)
25      {
26          this.tempIndexWriter = new TempIndexWriter(indexer.getAnalyzer(), tmpDir);
27      }
28  
29      public void write(final Document doc)
30      {
31          if (doc == null)
32          {
33              progress.incrementCounter();
34              return;
35          }
36  
37          String key = getWriterKey(doc);
38          try
39          {
40              tempIndexWriter.addDocument(key, doc);
41              progress.incrementCounter("Indexed: " + getDocumentTitle(doc) + " - " + progress.progressAsString());
42          }
43          catch (IOException e)
44          {
45              progress.incrementCounter("Error indexing: " + getDocumentTitle(doc) + " (" + e.toString() + ") - " + progress.progressAsString());
46              log.error("Error encountered", e);
47          }
48      }
49  
50      protected String getWriterKey(Document doc)
51      {
52          return Thread.currentThread().getName();
53      }
54  
55      protected String getDocumentTitle(Document doc)
56      {
57          return doc.get("title");
58      }
59  
60      public void runComplete()
61      {
62          // noop.
63      }
64  
65      /**
66       * Merge the temp indices into the final index, deleting the temp directories.
67       */
68      public void close(final ILuceneConnection luceneConnection) throws IOException
69      {
70          luceneConnection.withBatchUpdate(new ILuceneConnection.BatchUpdateAction() {
71              public void perform() throws Exception {
72                  luceneConnection.withWriter(new ILuceneConnection.WriterAction()
73                  {
74                      public void perform(IndexWriter writer) throws IOException
75                      {
76                          tempIndexWriter.merge(writer);
77                      }
78                  });
79              }
80          });
81          // remove the temporary indexes
82          tempIndexWriter.closeAll();
83      }
84  }