org.restheart.db.DocumentDAO.java Source code

Java tutorial

Introduction

Here is the source code for org.restheart.db.DocumentDAO.java

Source

/*
 * RESTHeart - the data Repository API server
 * Copyright (C) SoftInstigate Srl
 * 
 * This program 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.
 * 
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.restheart.db;

import com.mongodb.MongoClient;
import com.mongodb.bulk.BulkWriteResult;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import static com.mongodb.client.model.Filters.eq;
import com.mongodb.client.model.DeleteManyModel;
import static com.mongodb.client.model.Filters.and;
import com.mongodb.client.model.UpdateManyModel;
import com.mongodb.client.model.WriteModel;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import org.bson.BsonArray;
import org.bson.BsonDocument;
import org.bson.BsonObjectId;
import org.bson.BsonString;
import org.bson.BsonValue;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.restheart.utils.HttpStatus;
import org.bson.types.ObjectId;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * @author Andrea Di Cesare {@literal <andrea@softinstigate.com>}
 */
public class DocumentDAO implements Repository {

    private final Logger LOGGER = LoggerFactory.getLogger(DocumentDAO.class);

    private final MongoClient client;

    public DocumentDAO() {
        client = MongoDBClientSingleton.getInstance().getClient();
    }

    @Override
    public Document getDocumentEtag(final String dbName, final String collName, final Object documentId) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<Document> mcoll = mdb.getCollection(collName);

        FindIterable<Document> documents = mcoll.find(eq("_id", documentId)).projection(new Document("_etag", 1));

        return documents == null ? null : documents.iterator().tryNext();
    }

    /**
     * @param dbName
     * @param collName
     * @param documentId
     * @param shardKeys
     * @param newContent
     * @param requestEtag
     * @param patching
     * @param checkEtag
     * @return the HttpStatus code
     */
    @Override
    @SuppressWarnings("unchecked")
    public OperationResult upsertDocument(final String dbName, final String collName, final Object documentId,
            final BsonDocument shardKeys, final BsonDocument newContent, final String requestEtag,
            final boolean patching, final boolean checkEtag) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        // genereate new etag
        ObjectId newEtag = new ObjectId();

        final BsonDocument content = DAOUtils.validContent(newContent);

        content.put("_etag", new BsonObjectId(newEtag));

        OperationResult updateResult = DAOUtils.updateDocument(mcoll, documentId, shardKeys, content, !patching);

        BsonDocument oldDocument = updateResult.getOldData();

        if (patching) {
            if (oldDocument == null) {
                return new OperationResult(HttpStatus.SC_NOT_FOUND);
            } else if (checkEtag) {
                // check the old etag (in case restore the old document version)
                return optimisticCheckEtag(mcoll, shardKeys, oldDocument, newEtag, requestEtag, HttpStatus.SC_OK,
                        false);
            } else {
                BsonDocument newDocument = mcoll.find(eq("_id", documentId)).first();

                return new OperationResult(HttpStatus.SC_OK, newEtag, oldDocument, newDocument);
            }
        } else if (oldDocument != null && checkEtag) { // upsertDocument
            // check the old etag (in case restore the old document)
            return optimisticCheckEtag(mcoll, shardKeys, oldDocument, newEtag, requestEtag, HttpStatus.SC_OK,
                    false);
        } else if (oldDocument != null) { // insert
            BsonDocument newDocument = mcoll.find(eq("_id", documentId)).first();

            return new OperationResult(HttpStatus.SC_OK, newEtag, oldDocument, newDocument);
        } else {
            BsonDocument newDocument = mcoll.find(eq("_id", documentId)).first();

            return new OperationResult(HttpStatus.SC_CREATED, newEtag, null, newDocument);
        }
    }

    /**
     * @param dbName
     * @param collName
     * @param shardKeys
     * @param newContent
     * @param requestEtag
     * @param checkEtag
     * @return
     */
    @Override
    @SuppressWarnings("unchecked")
    public OperationResult upsertDocumentPost(final String dbName, final String collName, BsonDocument shardKeys,
            final BsonDocument newContent, final String requestEtag, final boolean checkEtag) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        ObjectId newEtag = new ObjectId();

        final BsonDocument content = DAOUtils.validContent(newContent);

        content.put("_etag", new BsonObjectId(newEtag));

        Object documentId;

        if (content.containsKey("_id")) {
            documentId = content.get("_id");
        } else {
            documentId = Optional.empty(); // key _id is not present
        }

        // new document since the id is missing ()
        OperationResult updateResult = DAOUtils.updateDocument(mcoll, documentId, shardKeys, content, true);

        BsonDocument oldDocument = updateResult.getOldData();
        BsonDocument newDocument = updateResult.getNewData();

        if (oldDocument == null) {
            return new OperationResult(HttpStatus.SC_CREATED, newEtag, null, newDocument);
        } else if (checkEtag) { // upsertDocument
            // check the old etag (in case restore the old document version)
            return optimisticCheckEtag(mcoll, shardKeys, oldDocument, newEtag, requestEtag, HttpStatus.SC_OK,
                    false);
        } else {
            return new OperationResult(HttpStatus.SC_OK, newEtag, oldDocument, newDocument);
        }
    }

    /**
     * @param dbName
     * @param collName
     * @param documents
     * @param shardKeys
     * @return
     */
    @Override
    @SuppressWarnings("unchecked")
    public BulkOperationResult bulkUpsertDocumentsPost(final String dbName, final String collName,
            final BsonArray documents, BsonDocument shardKeys) {
        Objects.requireNonNull(documents);

        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        BsonObjectId newEtag = new BsonObjectId(new ObjectId());

        documents.stream().filter(d -> d != null && d.isDocument()).forEachOrdered(document -> {
            document.asDocument().put("_etag", newEtag);
        });

        return DAOUtils.bulkUpsertDocuments(mcoll, documents, shardKeys);
    }

    /**
     * @param dbName
     * @param collName
     * @param documentId
     * @param shardedKeys
     * @param requestEtag
     * @param checkEtag
     * @return
     */
    @Override
    public OperationResult deleteDocument(final String dbName, final String collName, final Object documentId,
            final BsonDocument shardedKeys, final String requestEtag, final boolean checkEtag) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        BsonDocument oldDocument = mcoll.findOneAndDelete(getIdFilter(documentId, shardedKeys));

        if (oldDocument == null) {
            return new OperationResult(HttpStatus.SC_NOT_FOUND);
        } else if (checkEtag) {
            // check the old etag (in case restore the old document version)
            return optimisticCheckEtag(mcoll, null, oldDocument, null, requestEtag, HttpStatus.SC_NO_CONTENT, true);
        } else {
            return new OperationResult(HttpStatus.SC_NO_CONTENT);
        }
    }

    private Bson getIdFilter(Object documentId, BsonDocument shardedKeys) {
        if (shardedKeys != null) {
            return and(eq("_id", documentId), shardedKeys);
        } else {
            return eq("_id", documentId);
        }

    }

    @Override
    public BulkOperationResult bulkDeleteDocuments(String dbName, String collName, BsonDocument filter,
            BsonDocument shardedKeys) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        List<WriteModel<BsonDocument>> deletes = new ArrayList<>();

        Bson _filter;

        if (shardedKeys != null) {
            _filter = and(filter, shardedKeys);
        } else {
            _filter = filter;
        }

        deletes.add(new DeleteManyModel<>(_filter));

        BulkWriteResult result = mcoll.bulkWrite(deletes);

        return new BulkOperationResult(HttpStatus.SC_OK, null, result);
    }

    @Override
    public BulkOperationResult bulkPatchDocuments(String dbName, String collName, BsonDocument filter,
            BsonDocument shardedKeys, BsonDocument data) {
        MongoDatabase mdb = client.getDatabase(dbName);
        MongoCollection<BsonDocument> mcoll = mdb.getCollection(collName, BsonDocument.class);

        List<WriteModel<BsonDocument>> patches = new ArrayList<>();

        Bson _filter;

        if (shardedKeys != null) {
            _filter = and(filter, shardedKeys);
        } else {
            _filter = filter;
        }

        patches.add(new UpdateManyModel<>(_filter, DAOUtils.getUpdateDocument(data), DAOUtils.U_NOT_UPSERT_OPS));

        BulkWriteResult result = mcoll.bulkWrite(patches);

        return new BulkOperationResult(HttpStatus.SC_OK, null, result);
    }

    private OperationResult optimisticCheckEtag(final MongoCollection<BsonDocument> coll, BsonDocument shardKeys,
            final BsonDocument oldDocument, final Object newEtag, final String requestEtag,
            final int httpStatusIfOk, final boolean deleting) {

        BsonValue oldEtag = oldDocument.get("_etag");

        if (oldEtag != null && requestEtag == null) {
            // oopps, we need to restore old document
            // they call it optimistic lock strategy
            if (deleting) {
                DAOUtils.updateDocument(coll, oldDocument.get("_id"), shardKeys, oldDocument, true);
            } else {
                DAOUtils.restoreDocument(coll, oldDocument.get("_id"), shardKeys, oldDocument, newEtag);
            }

            return new OperationResult(HttpStatus.SC_CONFLICT, oldEtag, oldDocument, null);
        }

        BsonValue _requestEtag;

        if (ObjectId.isValid(requestEtag)) {
            _requestEtag = new BsonObjectId(new ObjectId(requestEtag));
        } else {
            // restheart generates ObjectId etags, but here we support
            // strings as well
            _requestEtag = new BsonString(requestEtag);
        }

        if (Objects.equals(_requestEtag, oldEtag)) {
            BsonDocument newDocument = coll.find(eq("_id", oldDocument.get("_id"))).first();

            return new OperationResult(httpStatusIfOk, newEtag, oldDocument, newDocument);
        } else {
            // oopps, we need to restore old document
            // they call it optimistic lock strategy
            if (deleting) {
                DAOUtils.updateDocument(coll, oldDocument.get("_id"), shardKeys, oldDocument, true);
            } else {
                DAOUtils.restoreDocument(coll, oldDocument.get("_id"), shardKeys, oldDocument, newEtag);
            }

            return new OperationResult(HttpStatus.SC_PRECONDITION_FAILED, oldEtag, oldDocument, null);
        }
    }
}