metron-platform/metron-solr/metron-solr-common/src/main/java/org/apache/metron/solr/writer/SolrWriter.java - metron - Git at Google

 /**
  * 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.metron.solr.writer;

 import static org.apache.metron.solr.SolrConstants.SOLR_WRITER_NAME;

 import com.google.common.base.Joiner;
 import java.io.IOException;
 import java.io.Serializable;
 import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;

 import org.apache.commons.lang3.StringUtils;
 import org.apache.metron.common.Constants;
 import org.apache.metron.common.configuration.writer.WriterConfiguration;
 import org.apache.metron.common.writer.BulkMessageWriter;
 import org.apache.metron.common.writer.BulkMessage;
 import org.apache.metron.common.writer.BulkWriterResponse;
 import org.apache.metron.common.writer.MessageId;
 import org.apache.metron.solr.SolrConstants;
 import org.apache.metron.stellar.common.utils.ConversionUtils;
 import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.client.solrj.impl.HttpClientUtil;
 import org.apache.solr.client.solrj.impl.HttpSolrClient;
 import org.apache.solr.client.solrj.impl.Krb5HttpClientConfigurer;
 import org.apache.solr.client.solrj.response.UpdateResponse;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.SolrInputDocument;
 import org.json.simple.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class SolrWriter implements BulkMessageWriter<JSONObject>, Serializable {

   public static final String JAVA_SECURITY_CONFIG_PROPERTY = "java.security.auth.login.config";

   public enum SolrProperties {
     ZOOKEEPER_QUORUM(SolrConstants.SOLR_ZOOKEEPER),
     COMMIT_PER_BATCH("solr.commitPerBatch", Optional.of(true)),
     COMMIT_WAIT_SEARCHER("solr.commit.waitSearcher", Optional.of(true)),
     COMMIT_WAIT_FLUSH("solr.commit.waitFlush", Optional.of(true)),
     COMMIT_SOFT("solr.commit.soft", Optional.of(false)),
     DEFAULT_COLLECTION("solr.collection", Optional.of("metron")),
     HTTP_CONFIG("solr.http.config", Optional.of(new HashMap<>()))
     ;
     String name;
     Optional<Object> defaultValue;

     SolrProperties(String name) {
       this(name, Optional.empty());
     }
     SolrProperties(String name, Optional<Object> defaultValue) {
       this.name = name;
       this.defaultValue = defaultValue;
     }

     public <T> Optional<T> coerceOrDefault(Map<String, Object> globalConfig, Class<T> clazz) {
       Object val = globalConfig.get(name);
       if(val != null) {
         T ret = null;
         try {
           ret = ConversionUtils.convert(val, clazz);
         }
         catch(ClassCastException cce) {
           ret = null;
         }
         if(ret == null) {
           //unable to convert value
           LOG.warn("Unable to convert {} to {}, was {}", name, clazz.getName(), "" + val);
           if(defaultValue.isPresent()) {
             return Optional.ofNullable(ConversionUtils.convert(defaultValue.get(), clazz));
           }
           else {
             return Optional.empty();
           }
         }
         else {
           return Optional.ofNullable(ret);
         }
       }
       else {
         if(defaultValue.isPresent()) {
           return Optional.ofNullable(ConversionUtils.convert(defaultValue.get(), clazz));
         }
         else {
           return Optional.empty();
         }
       }
     }

     public Supplier<IllegalArgumentException> errorOut(Map<String, Object> globalConfig) {
       String message = "Unable to retrieve " + name + " from global config, value associated is " + globalConfig.get(name);
       return () -> new IllegalArgumentException(message);
     }

     public <T> T coerceOrDefaultOrExcept(Map<String, Object> globalConfig, Class<T> clazz) {
          return this.coerceOrDefault(globalConfig, clazz).orElseThrow(this.errorOut(globalConfig));
     }

   }


   private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

   private Boolean shouldCommit;
   private Boolean softCommit;
   private Boolean waitSearcher;
   private Boolean waitFlush;
   private String zookeeperUrl;
   private String defaultCollection;
   private Map<String, Object> solrHttpConfig;

   private MetronSolrClient solr;

   public SolrWriter withMetronSolrClient(MetronSolrClient solr) {
     this.solr = solr;
     return this;
   }

   public void initializeFromGlobalConfig(Map<String, Object> globalConfiguration) {
     zookeeperUrl = SolrProperties.ZOOKEEPER_QUORUM.coerceOrDefaultOrExcept(globalConfiguration, String.class);
     defaultCollection = SolrProperties.DEFAULT_COLLECTION.coerceOrDefaultOrExcept(globalConfiguration, String.class);
     solrHttpConfig = SolrProperties.HTTP_CONFIG.coerceOrDefaultOrExcept(globalConfiguration, Map.class);
     shouldCommit = SolrProperties.COMMIT_PER_BATCH.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
     softCommit = SolrProperties.COMMIT_SOFT.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
     waitSearcher = SolrProperties.COMMIT_WAIT_SEARCHER.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
     waitFlush = SolrProperties.COMMIT_WAIT_FLUSH.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
   }

   @Override
   public void init(Map stormConf, WriterConfiguration configurations) throws IOException, SolrServerException {
     Map<String, Object> globalConfiguration = configurations.getGlobalConfig();
     initializeFromGlobalConfig(globalConfiguration);
     LOG.info("Initializing SOLR writer: {}", zookeeperUrl);
     LOG.info("Forcing commit per batch: {}", shouldCommit);
     LOG.info("Soft commit: {}", softCommit);
     LOG.info("Commit Wait Searcher: {}", waitSearcher);
     LOG.info("Commit Wait Flush: {}", waitFlush);
     LOG.info("Default Collection: {}", "" + defaultCollection );
     if(solr == null) {
       if (isKerberosEnabled(stormConf)) {
         HttpClientUtil.addConfigurer(new Krb5HttpClientConfigurer());
       }
       solr = new MetronSolrClient(zookeeperUrl, solrHttpConfig);
     }
     solr.setDefaultCollection(defaultCollection);

   }

   public Collection<SolrInputDocument> toDocs(Iterable<BulkMessage<JSONObject>> messages) {
     Collection<SolrInputDocument> ret = new ArrayList<>();
     for(BulkMessage<JSONObject> bulkWriterMessage: messages) {
       SolrInputDocument document = new SolrInputDocument();
       JSONObject message = bulkWriterMessage.getMessage();
       for (Object key : message.keySet()) {
         Object value = message.get(key);
         if (value instanceof Iterable) {
           for (Object v : (Iterable) value) {
             document.addField("" + key, v);
           }
         } else {
           document.addField("" + key, value);
         }
       }
       if (!document.containsKey(Constants.GUID)) {
         document.addField(Constants.GUID, UUID.randomUUID().toString());
       }
       ret.add(document);
     }
     return ret;
   }

   protected String getCollection(String sourceType, WriterConfiguration configurations) {
     String collection = configurations.getIndex(sourceType);
     if(StringUtils.isEmpty(collection)) {
       return solr.getDefaultCollection();
     }
     return collection;
   }

   @Override
   public BulkWriterResponse write(String sourceType, WriterConfiguration configurations, List<BulkMessage<JSONObject>> messages) throws Exception {
     String collection = getCollection(sourceType, configurations);
     BulkWriterResponse bulkResponse = new BulkWriterResponse();
     Collection<SolrInputDocument> docs = toDocs(messages);
     Set<MessageId> ids = messages.stream().map(BulkMessage::getId).collect(Collectors.toSet());
     try {
       Optional<SolrException> exceptionOptional = fromUpdateResponse(solr.add(collection, docs));
       // Solr commits the entire batch or throws an exception for it.  There's no way to get partial failures.
       if(exceptionOptional.isPresent()) {
         bulkResponse.addAllErrors(exceptionOptional.get(), ids);
       }
       else {
         if (shouldCommit) {
           exceptionOptional = fromUpdateResponse(solr.commit(collection, waitFlush, waitSearcher, softCommit));
           if(exceptionOptional.isPresent()) {
             bulkResponse.addAllErrors(exceptionOptional.get(), ids);
           }
         }
         if(!exceptionOptional.isPresent()) {
           bulkResponse.addAllSuccesses(ids);
         }
       }
     }
     catch(HttpSolrClient.RemoteSolrException sse) {
       bulkResponse.addAllErrors(sse, ids);
     }

     return bulkResponse;
   }

   protected Optional<SolrException> fromUpdateResponse(UpdateResponse response) {
     if(response != null && response.getStatus() > 0) {
       String message = "Solr Update response: " + Joiner.on(",").join(response.getResponse());
       return Optional.of(new SolrException(SolrException.ErrorCode.BAD_REQUEST, message));
     }
     return Optional.empty();
   }

   @Override
   public String getName() {
     return SOLR_WRITER_NAME;
   }

   @Override
   public void close() throws Exception {
     if(solr != null) {
       solr.close();
     }
   }

   private boolean isKerberosEnabled(Map stormConfig) {
     if (stormConfig == null) {
       return false;
     }
     String value = (String) stormConfig.get(JAVA_SECURITY_CONFIG_PROPERTY);
     return value != null && !value.isEmpty();
   }
 }
	/**
	* 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.metron.solr.writer;

	import static org.apache.metron.solr.SolrConstants.SOLR_WRITER_NAME;

	import com.google.common.base.Joiner;
	import java.io.IOException;
	import java.io.Serializable;
	import java.lang.invoke.MethodHandles;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Optional;
	import java.util.Set;
	import java.util.UUID;
	import java.util.function.Supplier;
	import java.util.stream.Collectors;

	import org.apache.commons.lang3.StringUtils;
	import org.apache.metron.common.Constants;
	import org.apache.metron.common.configuration.writer.WriterConfiguration;
	import org.apache.metron.common.writer.BulkMessageWriter;
	import org.apache.metron.common.writer.BulkMessage;
	import org.apache.metron.common.writer.BulkWriterResponse;
	import org.apache.metron.common.writer.MessageId;
	import org.apache.metron.solr.SolrConstants;
	import org.apache.metron.stellar.common.utils.ConversionUtils;
	import org.apache.solr.client.solrj.SolrServerException;
	import org.apache.solr.client.solrj.impl.HttpClientUtil;
	import org.apache.solr.client.solrj.impl.HttpSolrClient;
	import org.apache.solr.client.solrj.impl.Krb5HttpClientConfigurer;
	import org.apache.solr.client.solrj.response.UpdateResponse;
	import org.apache.solr.common.SolrException;
	import org.apache.solr.common.SolrInputDocument;
	import org.json.simple.JSONObject;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class SolrWriter implements BulkMessageWriter<JSONObject>, Serializable {

	public static final String JAVA_SECURITY_CONFIG_PROPERTY = "java.security.auth.login.config";

	public enum SolrProperties {
	ZOOKEEPER_QUORUM(SolrConstants.SOLR_ZOOKEEPER),
	COMMIT_PER_BATCH("solr.commitPerBatch", Optional.of(true)),
	COMMIT_WAIT_SEARCHER("solr.commit.waitSearcher", Optional.of(true)),
	COMMIT_WAIT_FLUSH("solr.commit.waitFlush", Optional.of(true)),
	COMMIT_SOFT("solr.commit.soft", Optional.of(false)),
	DEFAULT_COLLECTION("solr.collection", Optional.of("metron")),
	HTTP_CONFIG("solr.http.config", Optional.of(new HashMap<>()))
	;
	String name;
	Optional<Object> defaultValue;

	SolrProperties(String name) {
	this(name, Optional.empty());
	}
	SolrProperties(String name, Optional<Object> defaultValue) {
	this.name = name;
	this.defaultValue = defaultValue;
	}

	public <T> Optional<T> coerceOrDefault(Map<String, Object> globalConfig, Class<T> clazz) {
	Object val = globalConfig.get(name);
	if(val != null) {
	T ret = null;
	try {
	ret = ConversionUtils.convert(val, clazz);
	}
	catch(ClassCastException cce) {
	ret = null;
	}
	if(ret == null) {
	//unable to convert value
	LOG.warn("Unable to convert {} to {}, was {}", name, clazz.getName(), "" + val);
	if(defaultValue.isPresent()) {
	return Optional.ofNullable(ConversionUtils.convert(defaultValue.get(), clazz));
	}
	else {
	return Optional.empty();
	}
	}
	else {
	return Optional.ofNullable(ret);
	}
	}
	else {
	if(defaultValue.isPresent()) {
	return Optional.ofNullable(ConversionUtils.convert(defaultValue.get(), clazz));
	}
	else {
	return Optional.empty();
	}
	}
	}

	public Supplier<IllegalArgumentException> errorOut(Map<String, Object> globalConfig) {
	String message = "Unable to retrieve " + name + " from global config, value associated is " + globalConfig.get(name);
	return () -> new IllegalArgumentException(message);
	}

	public <T> T coerceOrDefaultOrExcept(Map<String, Object> globalConfig, Class<T> clazz) {
	return this.coerceOrDefault(globalConfig, clazz).orElseThrow(this.errorOut(globalConfig));
	}

	}


	private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

	private Boolean shouldCommit;
	private Boolean softCommit;
	private Boolean waitSearcher;
	private Boolean waitFlush;
	private String zookeeperUrl;
	private String defaultCollection;
	private Map<String, Object> solrHttpConfig;

	private MetronSolrClient solr;

	public SolrWriter withMetronSolrClient(MetronSolrClient solr) {
	this.solr = solr;
	return this;
	}

	public void initializeFromGlobalConfig(Map<String, Object> globalConfiguration) {
	zookeeperUrl = SolrProperties.ZOOKEEPER_QUORUM.coerceOrDefaultOrExcept(globalConfiguration, String.class);
	defaultCollection = SolrProperties.DEFAULT_COLLECTION.coerceOrDefaultOrExcept(globalConfiguration, String.class);
	solrHttpConfig = SolrProperties.HTTP_CONFIG.coerceOrDefaultOrExcept(globalConfiguration, Map.class);
	shouldCommit = SolrProperties.COMMIT_PER_BATCH.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
	softCommit = SolrProperties.COMMIT_SOFT.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
	waitSearcher = SolrProperties.COMMIT_WAIT_SEARCHER.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
	waitFlush = SolrProperties.COMMIT_WAIT_FLUSH.coerceOrDefaultOrExcept(globalConfiguration, Boolean.class);
	}

	@Override
	public void init(Map stormConf, WriterConfiguration configurations) throws IOException, SolrServerException {
	Map<String, Object> globalConfiguration = configurations.getGlobalConfig();
	initializeFromGlobalConfig(globalConfiguration);
	LOG.info("Initializing SOLR writer: {}", zookeeperUrl);
	LOG.info("Forcing commit per batch: {}", shouldCommit);
	LOG.info("Soft commit: {}", softCommit);
	LOG.info("Commit Wait Searcher: {}", waitSearcher);
	LOG.info("Commit Wait Flush: {}", waitFlush);
	LOG.info("Default Collection: {}", "" + defaultCollection );
	if(solr == null) {
	if (isKerberosEnabled(stormConf)) {
	HttpClientUtil.addConfigurer(new Krb5HttpClientConfigurer());
	}
	solr = new MetronSolrClient(zookeeperUrl, solrHttpConfig);
	}
	solr.setDefaultCollection(defaultCollection);

	}

	public Collection<SolrInputDocument> toDocs(Iterable<BulkMessage<JSONObject>> messages) {
	Collection<SolrInputDocument> ret = new ArrayList<>();
	for(BulkMessage<JSONObject> bulkWriterMessage: messages) {
	SolrInputDocument document = new SolrInputDocument();
	JSONObject message = bulkWriterMessage.getMessage();
	for (Object key : message.keySet()) {
	Object value = message.get(key);
	if (value instanceof Iterable) {
	for (Object v : (Iterable) value) {
	document.addField("" + key, v);
	}
	} else {
	document.addField("" + key, value);
	}
	}
	if (!document.containsKey(Constants.GUID)) {
	document.addField(Constants.GUID, UUID.randomUUID().toString());
	}
	ret.add(document);
	}
	return ret;
	}

	protected String getCollection(String sourceType, WriterConfiguration configurations) {
	String collection = configurations.getIndex(sourceType);
	if(StringUtils.isEmpty(collection)) {
	return solr.getDefaultCollection();
	}
	return collection;
	}

	@Override
	public BulkWriterResponse write(String sourceType, WriterConfiguration configurations, List<BulkMessage<JSONObject>> messages) throws Exception {
	String collection = getCollection(sourceType, configurations);
	BulkWriterResponse bulkResponse = new BulkWriterResponse();
	Collection<SolrInputDocument> docs = toDocs(messages);
	Set<MessageId> ids = messages.stream().map(BulkMessage::getId).collect(Collectors.toSet());
	try {
	Optional<SolrException> exceptionOptional = fromUpdateResponse(solr.add(collection, docs));
	// Solr commits the entire batch or throws an exception for it. There's no way to get partial failures.
	if(exceptionOptional.isPresent()) {
	bulkResponse.addAllErrors(exceptionOptional.get(), ids);
	}
	else {
	if (shouldCommit) {
	exceptionOptional = fromUpdateResponse(solr.commit(collection, waitFlush, waitSearcher, softCommit));
	if(exceptionOptional.isPresent()) {
	bulkResponse.addAllErrors(exceptionOptional.get(), ids);
	}
	}
	if(!exceptionOptional.isPresent()) {
	bulkResponse.addAllSuccesses(ids);
	}
	}
	}
	catch(HttpSolrClient.RemoteSolrException sse) {
	bulkResponse.addAllErrors(sse, ids);
	}

	return bulkResponse;
	}

	protected Optional<SolrException> fromUpdateResponse(UpdateResponse response) {
	if(response != null && response.getStatus() > 0) {
	String message = "Solr Update response: " + Joiner.on(",").join(response.getResponse());
	return Optional.of(new SolrException(SolrException.ErrorCode.BAD_REQUEST, message));
	}
	return Optional.empty();
	}

	@Override
	public String getName() {
	return SOLR_WRITER_NAME;
	}

	@Override
	public void close() throws Exception {
	if(solr != null) {
	solr.close();
	}
	}

	private boolean isKerberosEnabled(Map stormConfig) {
	if (stormConfig == null) {
	return false;
	}
	String value = (String) stormConfig.get(JAVA_SECURITY_CONFIG_PROPERTY);
	return value != null && !value.isEmpty();
	}
	}