metron-platform/metron-indexing/metron-indexing-common/src/main/java/org/apache/metron/indexing/dao/MultiIndexDao.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.indexing.dao;

 import com.google.common.base.Joiner;
 import com.google.common.collect.Iterables;
 import org.apache.commons.lang.ClassUtils;
 import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.apache.metron.indexing.dao.search.FieldType;
 import org.apache.metron.indexing.dao.search.GetRequest;
 import org.apache.metron.indexing.dao.search.GroupRequest;
 import org.apache.metron.indexing.dao.search.GroupResponse;
 import org.apache.metron.indexing.dao.search.InvalidSearchException;
 import org.apache.metron.indexing.dao.search.SearchRequest;
 import org.apache.metron.indexing.dao.search.SearchResponse;
 import org.apache.metron.indexing.dao.update.CommentAddRemoveRequest;
 import org.apache.metron.indexing.dao.update.Document;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.IOException;
 import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;

 public class MultiIndexDao implements IndexDao {

   private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
   private List<IndexDao> indices;

   public MultiIndexDao(IndexDao... composedDao) {
     indices = new ArrayList<>();
     Collections.addAll(indices, composedDao);
   }

   public MultiIndexDao(Iterable<IndexDao> composedDao) {
     this.indices = new ArrayList<>();
     Iterables.addAll(indices, composedDao);
   }

   public MultiIndexDao(Iterable<IndexDao> composedDao, Function<IndexDao, IndexDao> decoratorTransformation) {
     this(Iterables.transform(composedDao, x -> decoratorTransformation.apply(x)));
   }

   @Override
   public Document update(final Document update, Optional<String> index) throws IOException {
     List<String> exceptions =
     indices.parallelStream().map(dao -> {
       try {
         dao.update(update, index);
         return null;
       } catch (Throwable e) {
         return dao.getClass() + ": " + e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e);
       }
     }).filter(e -> e != null).collect(Collectors.toList());
     if(exceptions.size() > 0) {
       throw new IOException(Joiner.on("\n").join(exceptions));
     }
     return update;
   }

   @Override
   public Map<Document, Optional<String>> batchUpdate(Map<Document, Optional<String>> updates) throws IOException {
     List<String> exceptions =
         indices.parallelStream().map(dao -> {
           try {
             dao.batchUpdate(updates);
             return null;
           } catch (Throwable e) {
             return dao.getClass() + ": " + e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e);
           }
         }).filter(e -> e != null).collect(Collectors.toList());
     if (exceptions.size() > 0) {
       throw new IOException(Joiner.on("\n").join(exceptions));
     }
     return updates;
   }

   @Override
   public Map<String, FieldType> getColumnMetadata(List<String> in) throws IOException {
     for(IndexDao dao : indices) {
       Map<String, FieldType> r = dao.getColumnMetadata(in);
       if(r != null) {
         return r;
       }
     }
     return null;
   }

   @Override
   public Document addCommentToAlert(CommentAddRemoveRequest request) throws IOException {
     Document latest = getLatest(request.getGuid(), request.getSensorType());
     return addCommentToAlert(request, latest);
   }

   /**
    * Adds comments to an alert.  Updates are written to each Dao in parallel with the assumption that all updates
    * are identical.  The first update to be applied is returned as the current version of the alert with comments added.
    * @param request Request to add comments
    * @param latest The latest version of the alert the comments will be added to.
    * @return The complete alert document with comments added.
    * @throws IOException
    */
   @Override
   public Document addCommentToAlert(CommentAddRemoveRequest request, Document latest) throws IOException {
     List<DocumentContainer> output = indices
             .parallelStream()
             .map(dao -> addCommentToAlert(dao, request, latest))
             .collect(Collectors.toList());
     return getLatestDocument(output);
   }

   private DocumentContainer addCommentToAlert(IndexDao indexDao, CommentAddRemoveRequest request, Document latest) {
     DocumentContainer container;
     try {
       Document document = indexDao.addCommentToAlert(request, latest);
       container = new DocumentContainer(document);
       LOG.debug("Added comment to alert; indexDao={}, guid={}, sensorType={}, document={}",
               ClassUtils.getShortClassName(indexDao.getClass()), document.getGuid(), document.getSensorType(), document);

     } catch (Throwable e) {
       container = new DocumentContainer(e);
       LOG.error("Unable to add comment to alert; indexDao={}, error={}",
               ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
     }

     return container;
   }

   @Override
   public Document removeCommentFromAlert(CommentAddRemoveRequest request) throws IOException {
     Document latest = getLatest(request.getGuid(), request.getSensorType());
     return removeCommentFromAlert(request, latest);
   }

   /**
    * Removes comments from an alert.  Updates are written to each Dao in parallel with the assumption that all updates
    * are identical.  The first update to be applied is returned as the current version of the alert with comments removed.
    * @param request Request to remove comments
    * @param latest The latest version of the alert the comments will be removed from.
    * @return The complete alert document with comments removed.
    * @throws IOException
    */
   @Override
   public Document removeCommentFromAlert(CommentAddRemoveRequest request, Document latest) throws IOException {
     List<DocumentContainer> output = indices
             .parallelStream()
             .map(dao -> removeCommentFromAlert(dao, request, latest))
             .collect(Collectors.toList());
     return getLatestDocument(output);
   }

   private DocumentContainer removeCommentFromAlert(IndexDao indexDao, CommentAddRemoveRequest request, Document latest) {
     DocumentContainer container;
     try {
       Document document = indexDao.removeCommentFromAlert(request, latest);
       container = new DocumentContainer(document);
       LOG.debug("Removed comment from alert; indexDao={}, guid={}, sensorType={}, document={}",
               ClassUtils.getShortClassName(indexDao.getClass()), document.getGuid(), document.getSensorType(), document);

     } catch (Throwable e) {
       container = new DocumentContainer(e);
       LOG.error("Unable to remove comment from alert; indexDao={}, error={}",
               ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
     }

     return container;
   }

   protected static class DocumentContainer {
     private Optional<Document> d = Optional.empty();
     private Optional<Throwable> t = Optional.empty();
     public DocumentContainer(Document d) {
       this.d = Optional.ofNullable(d);
     }
     public DocumentContainer(Throwable t) {
       this.t = Optional.ofNullable(t);
     }

     public Optional<Document> getDocument() {
       return d;
     }
     public Optional<Throwable> getException() {
       return t;
     }

   }

   private static class DocumentIterableContainer {
     private Optional<Iterable<Document>> d = Optional.empty();
     private Optional<Throwable> t = Optional.empty();
     public DocumentIterableContainer(Iterable<Document> d) {
       this.d = Optional.ofNullable(d);
     }
     public DocumentIterableContainer(Throwable t) {
       this.t = Optional.ofNullable(t);
     }

     public Optional<Iterable<Document>> getDocumentIterable() {
       return d;
     }
     public Optional<Throwable> getException() {
       return t;
     }

   }

   @Override
   public SearchResponse search(SearchRequest searchRequest) throws InvalidSearchException {
     for(IndexDao dao : indices) {
       SearchResponse s = dao.search(searchRequest);
       if(s != null) {
         return s;
       }
     }
     return null;
   }

   @Override
   public GroupResponse group(GroupRequest groupRequest) throws InvalidSearchException {
     for(IndexDao dao : indices) {
       GroupResponse s = dao.group(groupRequest);
       if(s != null) {
         return s;
       }
     }
     return null;
   }

   @Override
   public void init(AccessConfig config) {
     for(IndexDao dao : indices) {
       dao.init(config);
     }
   }

   @Override
   public Document getLatest(final String guid, String sensorType) throws IOException {
     List<DocumentContainer> output = indices
             .parallelStream()
             .map(dao -> getLatest(dao, guid, sensorType))
             .collect(Collectors.toList());
     return getLatestDocument(output);
   }

   private DocumentContainer getLatest(IndexDao indexDao, String guid, String sensorType) {
     DocumentContainer container;
     try {
       Document document = indexDao.getLatest(guid, sensorType);
       container = new DocumentContainer(document);
       LOG.debug("Found latest document; indexDao={}, guid={}, sensorType={}, document={}",
               ClassUtils.getShortClassName(indexDao.getClass()), guid, sensorType, document);

     } catch (Throwable e) {
       container = new DocumentContainer(e);
       LOG.error("Unable to find latest document; indexDao={}, error={}",
               ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
     }

     return container;
   }

   @Override
   public Iterable<Document> getAllLatest(
       List<GetRequest> getRequests) throws IOException {
     Iterable<Document> ret = null;
     List<DocumentIterableContainer> output =
         indices.parallelStream().map(dao -> {
           try {
             return new DocumentIterableContainer(dao.getAllLatest(getRequests));
           } catch (Throwable e) {
             return new DocumentIterableContainer(e);
           }
         }).collect(Collectors.toList());

     List<String> error = new ArrayList<>();
     for(DocumentIterableContainer dc : output) {
       if(dc.getException().isPresent()) {
         Throwable e = dc.getException().get();
         error.add(e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e));
       }
       else {
         if(dc.getDocumentIterable().isPresent()) {
           Iterable<Document> documents = dc.getDocumentIterable().get();
           if(ret == null) {
             ret = documents;
           }
         }
       }
     }
     if(error.size() > 0) {
       throw new IOException(Joiner.on("\n").join(error));
     }
     return ret;
   }

   public List<IndexDao> getIndices() {
     return indices;
   }

   /**
    * Returns the most recent {@link Document} from a list of {@link DocumentContainer}s.
    *
    * @param documentContainers A list of containers; each retrieved from a separate index.
    * @return The latest {@link Document} found.
    * @throws IOException If any of the {@link DocumentContainer}s contain an exception.
    */
   private Document getLatestDocument(List<DocumentContainer> documentContainers) throws IOException {
     Document latestDocument = null;
     List<String> error = new ArrayList<>();

     for(DocumentContainer dc : documentContainers) {
       if(dc.getException().isPresent()) {
         // collect each exception; multiple can occur, one in each index
         Throwable e = dc.getException().get();
         error.add(e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e));

       } else if(dc.getDocument().isPresent()) {
         Document d = dc.getDocument().get();
         // is this the latest document so far?
         if(latestDocument == null || latestDocument.getTimestamp() < d.getTimestamp()) {
           latestDocument = d;
         }

       } else {
         // no document was found in the index
       }
     }
     if(error.size() > 0) {
       // report all of the errors encountered
       throw new IOException(Joiner.on("\n").join(error));
     }
     return latestDocument;
   }
 }
	/**
	* 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.indexing.dao;

	import com.google.common.base.Joiner;
	import com.google.common.collect.Iterables;
	import org.apache.commons.lang.ClassUtils;
	import org.apache.commons.lang3.exception.ExceptionUtils;
	import org.apache.metron.indexing.dao.search.FieldType;
	import org.apache.metron.indexing.dao.search.GetRequest;
	import org.apache.metron.indexing.dao.search.GroupRequest;
	import org.apache.metron.indexing.dao.search.GroupResponse;
	import org.apache.metron.indexing.dao.search.InvalidSearchException;
	import org.apache.metron.indexing.dao.search.SearchRequest;
	import org.apache.metron.indexing.dao.search.SearchResponse;
	import org.apache.metron.indexing.dao.update.CommentAddRemoveRequest;
	import org.apache.metron.indexing.dao.update.Document;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.IOException;
	import java.lang.invoke.MethodHandles;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.List;
	import java.util.Map;
	import java.util.Optional;
	import java.util.function.Function;
	import java.util.stream.Collectors;

	public class MultiIndexDao implements IndexDao {

	private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
	private List<IndexDao> indices;

	public MultiIndexDao(IndexDao... composedDao) {
	indices = new ArrayList<>();
	Collections.addAll(indices, composedDao);
	}

	public MultiIndexDao(Iterable<IndexDao> composedDao) {
	this.indices = new ArrayList<>();
	Iterables.addAll(indices, composedDao);
	}

	public MultiIndexDao(Iterable<IndexDao> composedDao, Function<IndexDao, IndexDao> decoratorTransformation) {
	this(Iterables.transform(composedDao, x -> decoratorTransformation.apply(x)));
	}

	@Override
	public Document update(final Document update, Optional<String> index) throws IOException {
	List<String> exceptions =
	indices.parallelStream().map(dao -> {
	try {
	dao.update(update, index);
	return null;
	} catch (Throwable e) {
	return dao.getClass() + ": " + e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e);
	}
	}).filter(e -> e != null).collect(Collectors.toList());
	if(exceptions.size() > 0) {
	throw new IOException(Joiner.on("\n").join(exceptions));
	}
	return update;
	}

	@Override
	public Map<Document, Optional<String>> batchUpdate(Map<Document, Optional<String>> updates) throws IOException {
	List<String> exceptions =
	indices.parallelStream().map(dao -> {
	try {
	dao.batchUpdate(updates);
	return null;
	} catch (Throwable e) {
	return dao.getClass() + ": " + e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e);
	}
	}).filter(e -> e != null).collect(Collectors.toList());
	if (exceptions.size() > 0) {
	throw new IOException(Joiner.on("\n").join(exceptions));
	}
	return updates;
	}

	@Override
	public Map<String, FieldType> getColumnMetadata(List<String> in) throws IOException {
	for(IndexDao dao : indices) {
	Map<String, FieldType> r = dao.getColumnMetadata(in);
	if(r != null) {
	return r;
	}
	}
	return null;
	}

	@Override
	public Document addCommentToAlert(CommentAddRemoveRequest request) throws IOException {
	Document latest = getLatest(request.getGuid(), request.getSensorType());
	return addCommentToAlert(request, latest);
	}

	/**
	* Adds comments to an alert. Updates are written to each Dao in parallel with the assumption that all updates
	* are identical. The first update to be applied is returned as the current version of the alert with comments added.
	* @param request Request to add comments
	* @param latest The latest version of the alert the comments will be added to.
	* @return The complete alert document with comments added.
	* @throws IOException
	*/
	@Override
	public Document addCommentToAlert(CommentAddRemoveRequest request, Document latest) throws IOException {
	List<DocumentContainer> output = indices
	.parallelStream()
	.map(dao -> addCommentToAlert(dao, request, latest))
	.collect(Collectors.toList());
	return getLatestDocument(output);
	}

	private DocumentContainer addCommentToAlert(IndexDao indexDao, CommentAddRemoveRequest request, Document latest) {
	DocumentContainer container;
	try {
	Document document = indexDao.addCommentToAlert(request, latest);
	container = new DocumentContainer(document);
	LOG.debug("Added comment to alert; indexDao={}, guid={}, sensorType={}, document={}",
	ClassUtils.getShortClassName(indexDao.getClass()), document.getGuid(), document.getSensorType(), document);

	} catch (Throwable e) {
	container = new DocumentContainer(e);
	LOG.error("Unable to add comment to alert; indexDao={}, error={}",
	ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
	}

	return container;
	}

	@Override
	public Document removeCommentFromAlert(CommentAddRemoveRequest request) throws IOException {
	Document latest = getLatest(request.getGuid(), request.getSensorType());
	return removeCommentFromAlert(request, latest);
	}

	/**
	* Removes comments from an alert. Updates are written to each Dao in parallel with the assumption that all updates
	* are identical. The first update to be applied is returned as the current version of the alert with comments removed.
	* @param request Request to remove comments
	* @param latest The latest version of the alert the comments will be removed from.
	* @return The complete alert document with comments removed.
	* @throws IOException
	*/
	@Override
	public Document removeCommentFromAlert(CommentAddRemoveRequest request, Document latest) throws IOException {
	List<DocumentContainer> output = indices
	.parallelStream()
	.map(dao -> removeCommentFromAlert(dao, request, latest))
	.collect(Collectors.toList());
	return getLatestDocument(output);
	}

	private DocumentContainer removeCommentFromAlert(IndexDao indexDao, CommentAddRemoveRequest request, Document latest) {
	DocumentContainer container;
	try {
	Document document = indexDao.removeCommentFromAlert(request, latest);
	container = new DocumentContainer(document);
	LOG.debug("Removed comment from alert; indexDao={}, guid={}, sensorType={}, document={}",
	ClassUtils.getShortClassName(indexDao.getClass()), document.getGuid(), document.getSensorType(), document);

	} catch (Throwable e) {
	container = new DocumentContainer(e);
	LOG.error("Unable to remove comment from alert; indexDao={}, error={}",
	ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
	}

	return container;
	}

	protected static class DocumentContainer {
	private Optional<Document> d = Optional.empty();
	private Optional<Throwable> t = Optional.empty();
	public DocumentContainer(Document d) {
	this.d = Optional.ofNullable(d);
	}
	public DocumentContainer(Throwable t) {
	this.t = Optional.ofNullable(t);
	}

	public Optional<Document> getDocument() {
	return d;
	}
	public Optional<Throwable> getException() {
	return t;
	}

	}

	private static class DocumentIterableContainer {
	private Optional<Iterable<Document>> d = Optional.empty();
	private Optional<Throwable> t = Optional.empty();
	public DocumentIterableContainer(Iterable<Document> d) {
	this.d = Optional.ofNullable(d);
	}
	public DocumentIterableContainer(Throwable t) {
	this.t = Optional.ofNullable(t);
	}

	public Optional<Iterable<Document>> getDocumentIterable() {
	return d;
	}
	public Optional<Throwable> getException() {
	return t;
	}

	}

	@Override
	public SearchResponse search(SearchRequest searchRequest) throws InvalidSearchException {
	for(IndexDao dao : indices) {
	SearchResponse s = dao.search(searchRequest);
	if(s != null) {
	return s;
	}
	}
	return null;
	}

	@Override
	public GroupResponse group(GroupRequest groupRequest) throws InvalidSearchException {
	for(IndexDao dao : indices) {
	GroupResponse s = dao.group(groupRequest);
	if(s != null) {
	return s;
	}
	}
	return null;
	}

	@Override
	public void init(AccessConfig config) {
	for(IndexDao dao : indices) {
	dao.init(config);
	}
	}

	@Override
	public Document getLatest(final String guid, String sensorType) throws IOException {
	List<DocumentContainer> output = indices
	.parallelStream()
	.map(dao -> getLatest(dao, guid, sensorType))
	.collect(Collectors.toList());
	return getLatestDocument(output);
	}

	private DocumentContainer getLatest(IndexDao indexDao, String guid, String sensorType) {
	DocumentContainer container;
	try {
	Document document = indexDao.getLatest(guid, sensorType);
	container = new DocumentContainer(document);
	LOG.debug("Found latest document; indexDao={}, guid={}, sensorType={}, document={}",
	ClassUtils.getShortClassName(indexDao.getClass()), guid, sensorType, document);

	} catch (Throwable e) {
	container = new DocumentContainer(e);
	LOG.error("Unable to find latest document; indexDao={}, error={}",
	ClassUtils.getShortClassName(indexDao.getClass()), ExceptionUtils.getRootCauseMessage(e));
	}

	return container;
	}

	@Override
	public Iterable<Document> getAllLatest(
	List<GetRequest> getRequests) throws IOException {
	Iterable<Document> ret = null;
	List<DocumentIterableContainer> output =
	indices.parallelStream().map(dao -> {
	try {
	return new DocumentIterableContainer(dao.getAllLatest(getRequests));
	} catch (Throwable e) {
	return new DocumentIterableContainer(e);
	}
	}).collect(Collectors.toList());

	List<String> error = new ArrayList<>();
	for(DocumentIterableContainer dc : output) {
	if(dc.getException().isPresent()) {
	Throwable e = dc.getException().get();
	error.add(e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e));
	}
	else {
	if(dc.getDocumentIterable().isPresent()) {
	Iterable<Document> documents = dc.getDocumentIterable().get();
	if(ret == null) {
	ret = documents;
	}
	}
	}
	}
	if(error.size() > 0) {
	throw new IOException(Joiner.on("\n").join(error));
	}
	return ret;
	}

	public List<IndexDao> getIndices() {
	return indices;
	}

	/**
	* Returns the most recent {@link Document} from a list of {@link DocumentContainer}s.
	*
	* @param documentContainers A list of containers; each retrieved from a separate index.
	* @return The latest {@link Document} found.
	* @throws IOException If any of the {@link DocumentContainer}s contain an exception.
	*/
	private Document getLatestDocument(List<DocumentContainer> documentContainers) throws IOException {
	Document latestDocument = null;
	List<String> error = new ArrayList<>();

	for(DocumentContainer dc : documentContainers) {
	if(dc.getException().isPresent()) {
	// collect each exception; multiple can occur, one in each index
	Throwable e = dc.getException().get();
	error.add(e.getMessage() + "\n" + ExceptionUtils.getStackTrace(e));

	} else if(dc.getDocument().isPresent()) {
	Document d = dc.getDocument().get();
	// is this the latest document so far?
	if(latestDocument == null \|\| latestDocument.getTimestamp() < d.getTimestamp()) {
	latestDocument = d;
	}

	} else {
	// no document was found in the index
	}
	}
	if(error.size() > 0) {
	// report all of the errors encountered
	throw new IOException(Joiner.on("\n").join(error));
	}
	return latestDocument;
	}
	}