/*
* 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.solr.handler.dataimport;
import org.apache.commons.io.IOUtils;
import org.apache.tika.config.TikaConfig;
import org.apache.tika.metadata.Metadata;
import org.apache.tika.parser.AutoDetectParser;
import org.apache.tika.parser.ParseContext;
import org.apache.tika.parser.Parser;
import org.apache.tika.parser.html.HtmlMapper;
import org.apache.tika.parser.html.IdentityHtmlMapper;
import org.apache.tika.sax.BodyContentHandler;
import org.apache.tika.sax.ContentHandlerDecorator;
import org.apache.tika.sax.XHTMLContentHandler;
import org.xml.sax.Attributes;
import org.xml.sax.ContentHandler;
import org.xml.sax.SAXException;
import org.xml.sax.helpers.DefaultHandler;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.sax.SAXTransformerFactory;
import javax.xml.transform.sax.TransformerHandler;
import javax.xml.transform.stream.StreamResult;
import java.io.File;
import java.io.InputStream;
import java.io.StringWriter;
import java.io.Writer;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import static org.apache.solr.handler.dataimport.DataImportHandlerException.SEVERE;
import static org.apache.solr.handler.dataimport.DataImportHandlerException.wrapAndThrow;
import static org.apache.solr.handler.dataimport.DataImporter.COLUMN;
import static org.apache.solr.handler.dataimport.XPathEntityProcessor.URL;
/**
* <p>An implementation of {@link EntityProcessor} which reads data from rich docs
* using <a href="http://tika.apache.org/">Apache Tika</a>
*
* <p>To index latitude/longitude data that might
* be extracted from a file's metadata, identify
* the geo field for this information with this attribute:
* <code>spatialMetadataField</code>
*
* @since solr 3.1
*/
public class TikaEntityProcessor extends EntityProcessorBase {
private TikaConfig tikaConfig;
private String format = "text";
private boolean done = false;
private boolean extractEmbedded = false;
private String parser;
static final String AUTO_PARSER = "org.apache.tika.parser.AutoDetectParser";
private String htmlMapper;
private String spatialMetadataField;
@Override
public void init(Context context) {
super.init(context);
done = false;
}
@Override
protected void firstInit(Context context) {
super.firstInit(context);
try {
String tikaConfigFile = context.getResolvedEntityAttribute("tikaConfig");
if (tikaConfigFile == null) {
ClassLoader classLoader = context.getSolrCore().getResourceLoader().getClassLoader();
tikaConfig = new TikaConfig(classLoader);
} else {
File configFile = new File(tikaConfigFile);
if (!configFile.isAbsolute()) {
configFile = new File(context.getSolrCore().getResourceLoader().getConfigDir(), tikaConfigFile);
}
tikaConfig = new TikaConfig(configFile);
}
} catch (Exception e) {
wrapAndThrow (SEVERE, e,"Unable to load Tika Config");
}
String extractEmbeddedString = context.getResolvedEntityAttribute("extractEmbedded");
if ("true".equals(extractEmbeddedString)) {
extractEmbedded = true;
}
format = context.getResolvedEntityAttribute("format");
if(format == null)
format = "text";
if (!"html".equals(format) && !"xml".equals(format) && !"text".equals(format)&& !"none".equals(format) )
throw new DataImportHandlerException(SEVERE, "'format' can be one of text|html|xml|none");
htmlMapper = context.getResolvedEntityAttribute("htmlMapper");
if (htmlMapper == null)
htmlMapper = "default";
if (!"default".equals(htmlMapper) && !"identity".equals(htmlMapper))
throw new DataImportHandlerException(SEVERE, "'htmlMapper', if present, must be 'default' or 'identity'");
parser = context.getResolvedEntityAttribute("parser");
if(parser == null) {
parser = AUTO_PARSER;
}
spatialMetadataField = context.getResolvedEntityAttribute("spatialMetadataField");
}
@Override
public Map<String, Object> nextRow() {
if(done) return null;
Map<String, Object> row = new HashMap<>();
DataSource<InputStream> dataSource = context.getDataSource();
InputStream is = dataSource.getData(context.getResolvedEntityAttribute(URL));
ContentHandler contentHandler = null;
Metadata metadata = new Metadata();
StringWriter sw = new StringWriter();
try {
if ("html".equals(format)) {
contentHandler = getHtmlHandler(sw);
} else if ("xml".equals(format)) {
contentHandler = getXmlContentHandler(sw);
} else if ("text".equals(format)) {
contentHandler = getTextContentHandler(sw);
} else if("none".equals(format)){
contentHandler = new DefaultHandler();
}
} catch (TransformerConfigurationException e) {
wrapAndThrow(SEVERE, e, "Unable to create content handler");
}
Parser tikaParser = null;
if(parser.equals(AUTO_PARSER)){
tikaParser = new AutoDetectParser(tikaConfig);
} else {
tikaParser = context.getSolrCore().getResourceLoader().newInstance(parser, Parser.class);
}
try {
ParseContext context = new ParseContext();
if ("identity".equals(htmlMapper)){
context.set(HtmlMapper.class, IdentityHtmlMapper.INSTANCE);
}
if (extractEmbedded) {
context.set(Parser.class, tikaParser);
}
tikaParser.parse(is, contentHandler, metadata , context);
} catch (Exception e) {
if(SKIP.equals(onError)) {
throw new DataImportHandlerException(DataImportHandlerException.SKIP_ROW,
"Document skipped :" + e.getMessage());
}
wrapAndThrow(SEVERE, e, "Unable to read content");
}
IOUtils.closeQuietly(is);
for (Map<String, String> field : context.getAllEntityFields()) {
if (!"true".equals(field.get("meta"))) continue;
String col = field.get(COLUMN);
String s = metadata.get(col);
if (s != null) row.put(col, s);
}
if(!"none".equals(format) ) row.put("text", sw.toString());
tryToAddLatLon(metadata, row);
done = true;
return row;
}
private void tryToAddLatLon(Metadata metadata, Map<String, Object> row) {
if (spatialMetadataField == null) return;
String latString = metadata.get(Metadata.LATITUDE);
String lonString = metadata.get(Metadata.LONGITUDE);
if (latString != null && lonString != null) {
row.put(spatialMetadataField, String.format(Locale.ROOT, "%s,%s", latString, lonString));
}
}
private static ContentHandler getHtmlHandler(Writer writer)
throws TransformerConfigurationException {
SAXTransformerFactory factory = (SAXTransformerFactory)
TransformerFactory.newInstance();
TransformerHandler handler = factory.newTransformerHandler();
handler.getTransformer().setOutputProperty(OutputKeys.METHOD, "html");
handler.setResult(new StreamResult(writer));
return new ContentHandlerDecorator(handler) {
@Override
public void startElement(
String uri, String localName, String name, Attributes atts)
throws SAXException {
if (XHTMLContentHandler.XHTML.equals(uri)) {
uri = null;
}
if (!"head".equals(localName)) {
super.startElement(uri, localName, name, atts);
}
}
@Override
public void endElement(String uri, String localName, String name)
throws SAXException {
if (XHTMLContentHandler.XHTML.equals(uri)) {
uri = null;
}
if (!"head".equals(localName)) {
super.endElement(uri, localName, name);
}
}
@Override
public void startPrefixMapping(String prefix, String uri) {/*no op*/ }
@Override
public void endPrefixMapping(String prefix) {/*no op*/ }
};
}
private static ContentHandler getTextContentHandler(Writer writer) {
return new BodyContentHandler(writer);
}
private static ContentHandler getXmlContentHandler(Writer writer)
throws TransformerConfigurationException {
SAXTransformerFactory factory = (SAXTransformerFactory)
TransformerFactory.newInstance();
TransformerHandler handler = factory.newTransformerHandler();
handler.getTransformer().setOutputProperty(OutputKeys.METHOD, "xml");
handler.setResult(new StreamResult(writer));
return handler;
}
}