AnnotationMongoDBAdaptor.java

/*
 * Copyright 2015-2016 OpenCB
 *
 * Licensed 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.opencb.opencga.catalog.db.mongodb;

import com.mongodb.client.model.Aggregates;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.Projections;
import com.mongodb.client.model.Updates;
import com.mongodb.client.result.UpdateResult;
import org.apache.commons.collections.map.LinkedMap;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.opencb.commons.datastore.core.ObjectMap;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.commons.datastore.core.QueryParam;
import org.opencb.commons.datastore.core.QueryResult;
import org.opencb.commons.datastore.mongodb.GenericDocumentComplexConverter;
import org.opencb.commons.datastore.mongodb.MongoDBCollection;
import org.opencb.opencga.catalog.exceptions.CatalogDBException;
import org.opencb.opencga.catalog.models.Annotable;
import org.opencb.opencga.catalog.models.Annotation;
import org.opencb.opencga.catalog.models.AnnotationSet;
import org.opencb.opencga.catalog.models.Variable;
import org.opencb.opencga.catalog.models.summaries.FeatureCount;
import org.opencb.opencga.catalog.models.summaries.VariableSummary;
import org.slf4j.Logger;

import javax.annotation.Nullable;
import java.util.*;

import static org.opencb.commons.datastore.core.QueryParam.Type.*;

/**
 * Created by pfurio on 07/07/16.
 */
abstract class AnnotationMongoDBAdaptor extends MongoDBAdaptor {

    AnnotationMongoDBAdaptor(Logger logger) {
        super(logger);
    }

    protected abstract GenericDocumentComplexConverter<? extends Annotable> getConverter();

    protected abstract MongoDBCollection getCollection();

    enum AnnotationSetParams implements QueryParam {
        ID("id", TEXT, ""),
        VARIABLE_SET_ID("variableSetId", DOUBLE, ""),

        ANNOTATION_SETS("annotationSets", TEXT_ARRAY, ""),
        ANNOTATION_SETS_NAME("annotationSets.name", TEXT, ""),
        ANNOTATION_SETS_VARIABLE_SET_ID("annotationSets.variableSetId", DECIMAL, ""),
        ANNOTATION_SETS_ANNOTATIONS("annotationSets.annotations", TEXT_ARRAY, ""),
        ANNOTATION_SETS_ANNOTATIONS_NAME("annotationSets.annotations.name", TEXT, ""),
        ANNOTATION_SETS_ANNOTATIONS_VALUE("annotationSets.annotations.value", TEXT, ""),

        ANNOTATIONS("annotations", TEXT_ARRAY, ""),
        NAME("name", TEXT, ""),
        VALUE("value", TEXT, "");

        private static Map<String, AnnotationSetParams> map;
        static {
            map = new LinkedMap();
            for (AnnotationSetParams params : AnnotationSetParams.values()) {
                map.put(params.key(), params);
            }
        }

        private final String key;
        private Type type;
        private String description;

        AnnotationSetParams(String key, Type type, String description) {
            this.key = key;
            this.type = type;
            this.description = description;
        }

        @Override
        public String key() {
            return key;
        }

        @Override
        public Type type() {
            return type;
        }

        @Override
        public String description() {
            return description;
        }

        public static Map<String, AnnotationSetParams> getMap() {
            return map;
        }

        public static AnnotationSetParams getParam(String key) {
            return map.get(key);
        }
    }

    public QueryResult<AnnotationSet> createAnnotationSet(long id, AnnotationSet annotationSet) throws CatalogDBException {
        long startTime = startQuery();
        // Check if there already exists an annotation set with the same name
        QueryResult<Long> count = getCollection().count(
                new Document()
                        .append(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSet.getName())
                        .append(PRIVATE_ID, id));

        if (count.first() > 0) {
            throw CatalogDBException.alreadyExists("AnnotationSet", AnnotationSetParams.NAME.key(), annotationSet.getName());
        }

        Document document = MongoDBUtils.getMongoDBDocument(annotationSet, "AnnotationSet");

        // Insert the annotation set in the database
        Bson query = Filters.and(
                Filters.eq(PRIVATE_ID, id),
                Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), new Document("$ne", annotationSet.getName()))
        );
        Bson update = new Document("$push", new Document(AnnotationSetParams.ANNOTATION_SETS.key(), document));
        QueryResult<UpdateResult> queryResult = getCollection().update(query, update, null);

        if (queryResult.first().getModifiedCount() != 1) {
            throw CatalogDBException.alreadyExists("AnnotationSet", AnnotationSetParams.NAME.key(), annotationSet.getName());
        }
        return endQuery("Create annotation set", startTime, getAnnotationSet(id, annotationSet.getName()));
    }

    public QueryResult<AnnotationSet> getAnnotationSet(long id, @Nullable String annotationSetName) throws CatalogDBException {
        long startTime = startQuery();

        QueryResult<? extends Annotable> aggregate = commonGetAnnotationSet(id, annotationSetName);

        List<AnnotationSet> annotationSets = new ArrayList<>(aggregate.getNumResults());
        for (Annotable annotable : aggregate.getResult()) {
            annotationSets.add((AnnotationSet) annotable.getAnnotationSets().get(0));
        }

        return endQuery("Get annotation set", startTime, annotationSets);
    }

    public QueryResult<ObjectMap> getAnnotationSetAsMap(long id, @Nullable String annotationSetName) throws CatalogDBException {
        long startTime = startQuery();

        QueryResult<? extends Annotable> aggregate = commonGetAnnotationSet(id, annotationSetName);

        List<ObjectMap> annotationSets = new ArrayList<>(aggregate.getNumResults());
        for (Annotable annotable : aggregate.getResult()) {
            annotationSets.add((ObjectMap) annotable.getAnnotationSetAsMap().get(0));
        }

        return endQuery("Get annotation set", startTime, annotationSets);
    }

    private QueryResult<? extends Annotable> commonGetAnnotationSet(long id, @Nullable String annotationSetName) {
        List<Bson> aggregation = new ArrayList<>();
        aggregation.add(Aggregates.match(Filters.eq(PRIVATE_ID, id)));
        aggregation.add(Aggregates.project(Projections.include(AnnotationSetParams.ID.key(), AnnotationSetParams.ANNOTATION_SETS.key())));
        aggregation.add(Aggregates.unwind("$" + AnnotationSetParams.ANNOTATION_SETS.key()));

        List<Bson> filters = new ArrayList<>();
        if (annotationSetName != null && !annotationSetName.isEmpty()) {
            filters.add(Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSetName));
        }

        if (filters.size() > 0) {
            Bson filter = filters.size() == 1 ? filters.get(0) : Filters.and(filters);
            aggregation.add(Aggregates.match(filter));
        }

        for (Bson bson : aggregation) {
            logger.debug("Get annotation: {}", bson.toBsonDocument(Document.class, com.mongodb.MongoClient.getDefaultCodecRegistry()));
        }

        return getCollection().aggregate(aggregation, getConverter(), null);
    }



    public QueryResult<AnnotationSet> updateAnnotationSet(long id, AnnotationSet annotationSet) throws CatalogDBException {
        long startTime = startQuery();

        // Check if there already exists an annotation set with the same name
        QueryResult<Long> count = getCollection().count(
                new Document()
                        .append(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSet.getName())
                        .append(PRIVATE_ID, id));

        if (count.first() == 0) {
            throw CatalogDBException.idNotFound("AnnotationSet", annotationSet.getName());
        }

        Document document = MongoDBUtils.getMongoDBDocument(annotationSet, "AnnotationSet");

        // Insert the annotation set in the database
        Bson query = Filters.and(
                Filters.eq(PRIVATE_ID, id),
                Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSet.getName())
        );
        Bson update = new Document("$set", new Document(AnnotationSetParams.ANNOTATION_SETS.key() + ".$", document));
        QueryResult<UpdateResult> queryResult = getCollection().update(query, update, null);

        if (queryResult.first().getModifiedCount() != 1) {
            throw new CatalogDBException("The annotation set could not be updated.");
        }

        return endQuery("Update annotation set", startTime, getAnnotationSet(id, annotationSet.getName()));
    }

    public void deleteAnnotationSet(long id, String annotationSetName) throws CatalogDBException {
        QueryResult<AnnotationSet> annotationSet = getAnnotationSet(id, annotationSetName);
        if (annotationSet == null || annotationSet.getNumResults() == 0) {
            throw CatalogDBException.idNotFound("Annotation set", annotationSetName);
        }

        Bson eq = Filters.eq(PRIVATE_ID, id);
        Bson pull = Updates.pull(AnnotationSetParams.ANNOTATION_SETS.key(),
                new Document(AnnotationSetParams.NAME.key(), annotationSetName));

        QueryResult<UpdateResult> update = getCollection().update(eq, pull, null);
        if (update.first().getModifiedCount() < 1) {
            throw new CatalogDBException("Could not delete the annotation set");
        }
    }

    public QueryResult<Long> addVariableToAnnotations(long variableSetId, Variable variable) throws CatalogDBException {
        long startTime = startQuery();

        Annotation annotation = new Annotation(variable.getName(), variable.getDefaultValue());

        // Obtain the annotation names of the annotations that are using the variableSet variableSetId
        List<Bson> aggregation = new ArrayList<>(4);
        aggregation.add(Aggregates.match(Filters.eq(AnnotationSetParams.ANNOTATION_SETS_VARIABLE_SET_ID.key(), variableSetId)));
        aggregation.add(Aggregates.unwind("$" + AnnotationSetParams.ANNOTATION_SETS.key()));
        aggregation.add(Aggregates.project(Projections.include(
                AnnotationSetParams.ANNOTATION_SETS_NAME.key(), AnnotationSetParams.ANNOTATION_SETS_VARIABLE_SET_ID.key()
        )));
        aggregation.add(Aggregates.match(Filters.eq(AnnotationSetParams.ANNOTATION_SETS_VARIABLE_SET_ID.key(), variableSetId)));
        QueryResult<Document> aggregationResult = getCollection().aggregate(aggregation, null);

        // Store the different annotation names in the set
        Set<String> annotationNames = new HashSet<>(aggregationResult.getNumResults());
        for (Document document : aggregationResult.getResult()) {
            annotationNames.add((String) ((Document) document.get(AnnotationSetParams.ANNOTATION_SETS.key()))
                    .get(AnnotationSetParams.NAME.key()));
        }

        // Prepare the update event
        Bson update = Updates.push(AnnotationSetParams.ANNOTATION_SETS.key() + ".$." + AnnotationSetParams.ANNOTATIONS.key(),
                MongoDBUtils.getMongoDBDocument(annotation, "Annotation"));

        // Construct the query dynamically for each different annotation set and make the update
        long modifiedCount = 0;
        Bson bsonQuery;
        for (String annotationId : annotationNames) {
            bsonQuery = Filters.elemMatch(AnnotationSetParams.ANNOTATION_SETS.key(), Filters.and(
                    Filters.eq(AnnotationSetParams.VARIABLE_SET_ID.key(), variableSetId),
                    Filters.eq(AnnotationSetParams.NAME.key(), annotationId)
            ));

            modifiedCount += getCollection().update(bsonQuery, update, new QueryOptions(MongoDBCollection.MULTI, true)).first()
                    .getModifiedCount();
        }

        return endQuery("Add annotation", startTime, Collections.singletonList(modifiedCount));
    }

    public QueryResult<Long> renameAnnotationField(long variableSetId, String oldName, String newName) throws CatalogDBException {
        long startTime = startQuery();
        long renamedAnnotations = 0;
        List<Document> aggregateResult = getAnnotationDocuments(variableSetId, oldName);

        if (aggregateResult.size() > 0) {
            // Each document will be a cohort, sample or individual
            for (Document entity : aggregateResult) {
                Object entityId = entity.get(AnnotationSetParams.ID.key());
                Document annotationSet = ((Document) entity.get(AnnotationSetParams.ANNOTATION_SETS.key()));

                String annotationSetName = annotationSet.getString(AnnotationSetParams.NAME.key());

                // Build a query to look for the particular annotations
                Bson bsonQuery = Filters.and(
                        Filters.eq(PRIVATE_ID, entityId),
                        Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSetName),
                        Filters.eq(AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_NAME.key(), oldName)
                );

                // And extract those annotations from the annotation set
                Bson update = Updates.pull(AnnotationSetParams.ANNOTATION_SETS.key() + ".$." + AnnotationSetParams.ANNOTATIONS.key(),
                        Filters.eq(AnnotationSetParams.NAME.key(), oldName));

                QueryResult<UpdateResult> queryResult = getCollection().update(bsonQuery, update, null);

                if (queryResult.first().getModifiedCount() != 1) {
                    throw new CatalogDBException("VariableSet {id: " + variableSetId + "} - AnnotationSet {name: "
                            + annotationSet.getString(AnnotationSetParams.NAME.key()) + "} - An unexpected error happened when "
                            + "extracting the annotation " + oldName + ". Please, report this error to the OpenCGA developers.");
                }

                // Obtain the value of the annotation
                Object value = ((Document) annotationSet.get(AnnotationSetParams.ANNOTATIONS.key())).get(AnnotationSetParams.VALUE.key());

                // Create a new annotation with the new id and the former value
                Annotation annotation = new Annotation(newName, value);

                bsonQuery = Filters.and(
                        Filters.eq(PRIVATE_ID, entityId),
                        Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSetName)
                );

                // Push the again the annotation with the new name
                update = Updates.push(AnnotationSetParams.ANNOTATION_SETS.key() + ".$." + AnnotationSetParams.ANNOTATIONS.key(),
                        MongoDBUtils.getMongoDBDocument(annotation, "Annotation"));
                queryResult = getCollection().update(bsonQuery, update, null);

                if (queryResult.first().getModifiedCount() != 1) {
                    throw new CatalogDBException("VariableSet {id: " + variableSetId + "} - AnnotationSet {name: "
                            + annotationSetName + "} - A critical error happened when trying to rename the annotation " + oldName
                            + ". Please, report this error to the OpenCGA developers.");
                }
                renamedAnnotations += 1;

            }
        }
        return endQuery("Rename annotation name", startTime, Collections.singletonList(renamedAnnotations));
    }

    public QueryResult<Long> removeAnnotationField(long variableSetId, String fieldId) throws CatalogDBException {
        long startTime = startQuery();
        long removedAnnotations = 0;
        List<Document> aggregateResult = getAnnotationDocuments(variableSetId, fieldId);

        if (aggregateResult.size() > 0) {
            // Each document will be a cohort, sample or individual
            for (Document entity : aggregateResult) {
                Object entityId = entity.get(AnnotationSetParams.ID.key());
                Document annotationSet = ((Document) entity.get(AnnotationSetParams.ANNOTATION_SETS.key()));

                String annotationSetName = annotationSet.getString(AnnotationSetParams.NAME.key());

                // Build a query to look for the particular annotations
                Bson bsonQuery = Filters.and(
                        Filters.eq(PRIVATE_ID, entityId),
                        Filters.eq(AnnotationSetParams.ANNOTATION_SETS_NAME.key(), annotationSetName),
                        Filters.eq(AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_NAME.key(), fieldId)
                );

                // Extract those annotations
                Bson update = Updates.pull(AnnotationSetParams.ANNOTATION_SETS.key() + ".$." + AnnotationSetParams.ANNOTATIONS.key(),
                        Filters.eq(AnnotationSetParams.NAME.key(), fieldId));
                QueryResult<UpdateResult> queryResult = getCollection().update(bsonQuery, update, null);
                if (queryResult.first().getModifiedCount() != 1) {
                    throw new CatalogDBException("VariableSet {id: " + variableSetId + "} - AnnotationSet {name: "
                            + annotationSetName + "} - An unexpected error happened when extracting the annotation " + fieldId
                            + ". Please, report this error to the OpenCGA developers.");
                }

                removedAnnotations += 1;

            }
        }
        return endQuery("Remove annotation", startTime, Collections.singletonList(removedAnnotations));
    }

    private List<Document> getAnnotationDocuments(long variableSetId, String oldName) {
        List<Bson> aggregation = new ArrayList<>();
        aggregation.add(Aggregates.match(
                Filters.elemMatch(AnnotationSetParams.ANNOTATION_SETS.key(),
                        Filters.eq(AnnotationSetParams.VARIABLE_SET_ID.key(), variableSetId))));
        aggregation.add(Aggregates.project(Projections.include(AnnotationSetParams.ANNOTATION_SETS.key(), AnnotationSetParams.ID.key())));
        aggregation.add(Aggregates.unwind("$" + AnnotationSetParams.ANNOTATION_SETS.key()));
        aggregation.add(Aggregates.match(Filters.eq(AnnotationSetParams.ANNOTATION_SETS_VARIABLE_SET_ID.key(), variableSetId)));
        aggregation.add(Aggregates.unwind("$" + AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS.key()));
        aggregation.add(Aggregates.match(
                Filters.eq(AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_NAME.key(), oldName)));

        return getCollection().aggregate(aggregation, new QueryOptions()).getResult();
    }

    public QueryResult<VariableSummary> getAnnotationSummary(long variableSetId) throws CatalogDBException {
        long startTime = startQuery();

        List<Bson> aggregation = new ArrayList<>(6);
        aggregation.add(new Document("$project", new Document(AnnotationSetParams.ANNOTATION_SETS.key(), 1)));
        aggregation.add(new Document("$unwind", "$" + AnnotationSetParams.ANNOTATION_SETS.key()));
        aggregation.add(new Document("$unwind", "$" + AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS.key()));
        // TODO: Include annotations of type object
        // At the moment, we are excluding the annotations of type Object.
        aggregation.add(new Document("$match",
                new Document(AnnotationSetParams.ANNOTATION_SETS_VARIABLE_SET_ID.key(), variableSetId)
                    .append(AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_VALUE.key(),
                            new Document("$not", new Document("$type", "object"))
                    )
                )
        );
        aggregation.add(new Document("$group",
                new Document(
                        "_id", new Document("name", "$" + AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_NAME.key())
                            .append("value", "$" + AnnotationSetParams.ANNOTATION_SETS_ANNOTATIONS_VALUE.key()))
                        .append("count", new Document("$sum", 1))
                )
        );
        aggregation.add(new Document("$sort", new Document("_id.name", -1).append("count", -1)));

        List<Document> result = getCollection().aggregate(aggregation, new QueryOptions()).getResult();

        List<VariableSummary> variableSummaryList = new ArrayList<>();

        List<FeatureCount> featureCountList = null;
        VariableSummary v = new VariableSummary();

        for (Document document : result) {
            Document id = (Document) document.get("_id");
            String name = id.getString("name");
            Object value = id.get("value");
            int count = document.getInteger("count");

            if (!name.equals(v.getName())) {
                featureCountList = new ArrayList<>();
                v = new VariableSummary(name, featureCountList);
                variableSummaryList.add(v);
            }

            featureCountList.add(new FeatureCount(value, count));
        }

        return endQuery("Get Annotation summary", startTime, variableSummaryList);
    }

}