step.core.collections.mongodb.MongoDBCollection Maven / Gradle / Ivy

Go to download
/*******************************************************************************
 * Copyright (C) 2020, exense GmbH
 *  
 * This file is part of STEP
 *  
 * STEP is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *  
 * STEP is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *  
 * You should have received a copy of the GNU Affero General Public License
 * along with STEP.  If not, see .
 ******************************************************************************/
package step.core.collections.mongodb;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import org.bson.Document;
import org.bson.UuidRepresentation;
import org.bson.conversions.Bson;
import org.bson.types.ObjectId;
import org.mongojack.JacksonMongoCollection;
import org.mongojack.ObjectMapperConfigurer;
import org.mongojack.SerializationOptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.mongodb.MongoExecutionTimeoutException;
import com.mongodb.MongoNamespace;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.model.CountOptions;
import com.mongodb.client.model.Projections;

import step.core.accessors.AbstractIdentifiableObject;
import step.core.collections.Collection;
import step.core.collections.Filter;
import step.core.collections.SearchOrder;
import step.core.collections.AbstractCollection;

public class MongoDBCollection extends AbstractCollection implements Collection {
	
	private static final Logger logger = LoggerFactory.getLogger(MongoDBCollection.class);

	protected static final String CSV_DELIMITER = ";";
	
	private final MongoClientSession mongoClientSession;
	private final Class entityClass;

	private JacksonMongoCollection collection;
	
	/**
	 * @param collectionName the name of the mongo collection
	 * @param entityClass the
	 */
	public MongoDBCollection(MongoClientSession mongoClientSession, String collectionName, Class entityClass) {
		this.mongoClientSession = mongoClientSession;
		this.entityClass = entityClass;
		collection = JacksonMongoCollection.builder()
				.withObjectMapper(ObjectMapperConfigurer.configureObjectMapper(MongoDBCollectionJacksonMapperProvider.getObjectMapper()))
				.withSerializationOptions(SerializationOptions.builder().withSimpleFilterSerialization(true).build())
				.build(mongoClientSession.getMongoDatabase(), collectionName, entityClass, UuidRepresentation.JAVA_LEGACY);
	}
	
	@Override
	public long count(Filter filter, Integer limit) {
		Bson query = filterToQuery(filter);
		CountOptions countOptions = new CountOptions();
		if(limit != null) {
			countOptions.limit(limit);
		}
		return collection.countDocuments(query, countOptions);
	}

	@Override
	public long estimatedCount() {
		return collection.estimatedDocumentCount();
	}

	/**
	 * @param columnName the name of the column (field)
	 * @return the distinct values of the column 
	 */
	@Override
	public List distinct(String columnName, Filter filter) {
		Bson query = filterToQuery(filter);
		
		if (columnName.equals(AbstractIdentifiableObject.ID)) {
			return collection.distinct("_id", query, ObjectId.class).map(ObjectId::toString).into(new ArrayList());
		} else {
			return collection.distinct(columnName, query, String.class).into(new ArrayList());
		}
	}
	
	private Bson filterToQuery(Filter filter) {
		return new MongoDBFilterFactory().buildFilter(filter);
	}
	
	@Override
	public Stream find(Filter filter, SearchOrder order, Integer skip, Integer limit, int maxTime) {
		return find(filter, order, skip, limit, maxTime, null);
	}

	@Override
	public Stream findReduced(Filter filter, SearchOrder order, Integer skip, Integer limit, int maxTime, List reduceFields ) {
		return find(filter, order, skip, limit, maxTime, reduceFields);
	}

	private Stream find(Filter filter, SearchOrder order, Integer skip, Integer limit, int maxTime, List reduceFields ) {
		Bson query = filterToQuery(filter);
		//long count = collection.estimatedDocumentCount();
		
		//CountOptions option = new CountOptions();
		//option.skip(0).limit(DEFAULT_LIMIT);
		//long countResults = collection.countDocuments(query, option);
		
		FindIterable find = (reduceFields != null && !reduceFields.isEmpty()) ?
				collection.find(query).projection(Projections.include(reduceFields)).maxTime(maxTime, TimeUnit.SECONDS).batchSize(mongoClientSession.getBatchSize()) :
				collection.find(query).maxTime(maxTime, TimeUnit.SECONDS).batchSize(mongoClientSession.getBatchSize());

		if (order != null) {
			String attributeName = fixAttributeName(order.getAttributeName());
			Document sortDoc = new Document(attributeName, order.getOrder());
			find.sort(sortDoc);
		}
		if(skip!=null) {
			find.skip(skip);
		}
		if(limit!=null) {
			find.limit(limit);
		}
		MongoCursor iterator;
		try {
			iterator = find.iterator();
		} catch (MongoExecutionTimeoutException e) {
			logger.error("Query execution exceeded timeout of " + maxTime + " " + TimeUnit.SECONDS);
			throw e;
		}
		
		Iterator enrichedIterator = new Iterator() {
			@Override
			public boolean hasNext() {
				return iterator.hasNext();
			}

			@Override
			public T next() {
				T next = iterator.next();
				fixIdAfterRead(next);
				return next;
			}
		};
		
		return StreamSupport.stream(Spliterators.spliteratorUnknownSize(enrichedIterator, Spliterator.ORDERED), false);
	}

	private String fixAttributeName(String attributeName) {
		if(attributeName.equals(AbstractIdentifiableObject.ID)) {
			attributeName = "_id";
		} else if (attributeName.contains("."+AbstractIdentifiableObject.ID)) {
			attributeName = attributeName.replace("."+AbstractIdentifiableObject.ID, "._id");
		}
		return attributeName;
	}

	private void fixIdAfterRead(T next) {
		if(next instanceof step.core.collections.Document) {
			step.core.collections.Document document = (step.core.collections.Document) next;
			Object id = document.get("_id");
			if(id instanceof ObjectId) {
				id = id.toString();
			}
			document.put(AbstractIdentifiableObject.ID, id);
			document.remove("_id");
		}
	}
	
	private void fixIdBeforeSave(T next) {
		if(next instanceof step.core.collections.Document) {
			step.core.collections.Document document = (step.core.collections.Document) next;
			Object id = document.get(AbstractIdentifiableObject.ID);
			if(id instanceof String) {
				id = new ObjectId((String) id);
			}
			document.put("_id", id);
			document.remove(AbstractIdentifiableObject.ID);
		}
	}

	@Override
	public void remove(Filter filter) {
		collection.deleteMany(filterToQuery(filter));
	}

	@Override
	public T save(T entity) {
		if(getId(entity) == null) {
			setId(entity, new ObjectId());
		}
		fixIdBeforeSave(entity);
		collection.save(entity);
		fixIdAfterRead(entity);
		return entity;
	}

	@Override
	public void save(Iterable entities) {
		entities.forEach(this::save);
	}

	@Override
	public void createOrUpdateIndex(String field) {
		createOrUpdateIndex(field, 1);
	}

	@Override
	public void createOrUpdateIndex(String field, int order) {
		createOrUpdateIndex(collection, field, order);
	}

	@Override
	public void createOrUpdateCompoundIndex(String... fields) {
		Map fieldsMap = new LinkedHashMap<>();
		Arrays.asList(fields).forEach(s -> fieldsMap.put(s,1));
		createOrUpdateCompoundIndex(fieldsMap);
	}

	@Override
	public void createOrUpdateCompoundIndex(Map fields) {
		createOrUpdateCompoundIndex(collection, fields);
	}


	public static void createOrUpdateIndex(com.mongodb.client.MongoCollection collection, String attribute, int order) {
		Document index = getIndex(collection, Set.of(attribute));
		if(index==null) {
			collection.createIndex(new Document(attribute,order));
		}
	}

	public static void createOrUpdateCompoundIndex(com.mongodb.client.MongoCollection collection, Map fields) {
		Document index = getIndex(collection, fields.keySet());
		
		if(index==null) {
			Document newIndex = new Document();
			
			for(String key : fields.keySet())
				newIndex.append(key, fields.get(key));

			collection.createIndex(newIndex);
		}
	}
	
	private static Document getIndex(com.mongodb.client.MongoCollection collection, Set attributes) {

		for(Document index:collection.listIndexes()) {  // inspect all indexes, looking for a match
			Object o = index.get("key");

			if(o instanceof Document) {
				Document d = ((Document)o);
				
				if(attributes.equals(d.keySet())) {
					return d;
				}
			}
		}
		return null;
	}

	@Override
	public void rename(String newName) {
		collection.renameCollection(new MongoNamespace(mongoClientSession.getMongoDatabase().getName(), newName));
	}

	@Override
	public void drop() {
		collection.drop();
	}

	@Override
	public Class getEntityClass() {
		return entityClass;
	}

//	/**
//	 * Export data to CSV
//	 * @param query
//	 * @param columns
//	 * @param writer
//	 */
//	public void export(Filter filter, Map columns, PrintWriter writer) {
//		Bson query = filterToQuery(filter);
//		FindIterable find = collection.find(query);
//		MongoCursor iterator;
//		iterator = find.iterator();
//		if (!iterator.hasNext()) {
//			return;
//		}
//		BasicDBObject basicDBObject = iterator.next();
//		//if column names not provided by the caller, get them from the collection
//		if (columns == null || columns.size() == 0) {
//			columns = getExportFields();
//		}
//		//if the collection has also no specification, dump all keys found in first object
//		if (columns == null || columns.size() == 0 && iterator.hasNext()) {
//			columns = getExportFields();		
//			for (String key : basicDBObject.keySet()) {
//				columns.put(key, new CollectionField(key,key));			
//			}
//		}
//		//write headers
//		columns.values().forEach((v)->{
//			String title = v.getTitle().replaceAll("^ID", "id");
//			writer.print(title);
//			writer.print(CSV_DELIMITER);
//		});
//		writer.println();
//		//Dump first row (required when writting all keys found in 1st object)
//		dumpRow(basicDBObject,columns,writer);
//
//		int count = 1;
//		while(iterator.hasNext()) {
//			count++;
//			basicDBObject = iterator.next();
//			dumpRow(basicDBObject,columns,writer);
//		}
//	}
//	
//	private void dumpRow(BasicDBObject basicDBObject, Map fields, PrintWriter writer) {
//		fields.forEach((key,field)->{
//			Object value = basicDBObject.get(key);
//			if (value != null) {
//				String valueStr = field.getFormat().format(value);
//				if(valueStr.contains(CSV_DELIMITER)||valueStr.contains("\n")||valueStr.contains("\"")) {
//					valueStr = "\"" + valueStr.replaceAll("\"", "\"\"") + "\"";
//				}
//				writer.print(valueStr);
//			}
//			writer.print(CSV_DELIMITER);
//		});
//		writer.println();
//	}
//
//	protected Map getExportFields() {
//		Map result = new HashMap ();
//		return result;
//	}
	
}