eip_manage/src/main/java/com/yihu/hos/common/mongo/MongoDBPro.java

package com.yihu.hos.common.mongo;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.mongodb.MongoNamespace;
import com.mongodb.client.*;
import com.mongodb.client.model.*;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;
import org.apache.commons.lang3.StringUtils;
import org.bson.Document;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * MongoDBPro. Professional database CURD and Manager tool.
 *
 * @created Airhead 2016/2/17.
 */
public class MongoDBPro implements IMongoDBRunner, IMongoDBAdminer {

    private static final Map<String, MongoDBPro> map = new HashMap<String, MongoDBPro>();
    private final MongoDBConfig config;

    public MongoDBPro() {
        if (MongoDBKit.config == null) {
            throw new RuntimeException("The main config is null, initialize MonogDBKit first");
        }
        this.config = MongoDBKit.config;
    }

    public MongoDBPro(String configName) {
        this.config = MongoDBKit.getConfig(configName);
        if (this.config == null) {
            throw new IllegalArgumentException("Config not found by configName: " + configName);
        }
    }

    public static MongoDBPro use() {
        return use(MongoDBKit.config.name);
    }

    public static MongoDBPro use(String configName) {
        MongoDBPro result = map.get(configName);
        if (result == null) {
            result = new MongoDBPro(configName);
            map.put(configName, result);
        }

        return result;
    }

    public MongoDBPro db(String databaseName) {
        config.getDatabase(databaseName);
        return this;
    }

    @Override
    public long count(String collectionName) {
        return count(collectionName, null);
    }

    @Override
    public long count(String collectionName, String filter) {
        return count(collectionName, filter, null);
    }

    /**
     * @param collectionName
     * @param filter         the query filter
     * @param options        the options describing the count
     *                       <p>
     *                       {
     *                       limit: <integer>,
     *                       skip: <integer>,
     *                       hint: <hint>
     *                       }
     * @return
     */
    @Override

    public long count(String collectionName, String filter, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = new Document();
        if (filter != null) {
            filterDocument = Document.parse(filter);
        }

        CountOptions countOptions = new CountOptions();
        if (options != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(options, JsonNode.class);
                String hintString = rootNode.path("hint").toString();
                if (!StringUtils.isEmpty(hintString)) {
                    Document hint = Document.parse(hintString);
                    countOptions.hint(hint);
                } else {
                    countOptions.hint(new Document());
                }

                countOptions.limit(rootNode.path("limit").asInt());
                countOptions.skip(rootNode.path("skip").asInt());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        return collection.count(filterDocument, countOptions);
    }

    @Override
    public List<String> find(String collectionName) {
        return find(collectionName, null);
    }

    @Override
    public List<String> find(String collectionName, String filter) {
        return find(collectionName, filter, null);
    }

    @Override
    public List<String> find(String collectionName, String filter, String projection) {
        return find(collectionName, filter, projection, null);
    }

    @Override
    public List<String> find(String collectionName, String filter, String projection, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = new Document();
        if (filter != null) {
            filterDocument = Document.parse(filter);
        }

        Document projectionDocument = new Document();
        if (projection != null) {
            projectionDocument = Document.parse(projection);
        }

        FindIterable<Document> documents = collection.find(filterDocument).projection(projectionDocument);
        List<String> list = new ArrayList<>();
        try (MongoCursor<Document> cursor = documents.iterator()) {
            while (cursor.hasNext()) {
                Document doc = cursor.next();
                list.add(doc.toJson());
            }
        }

        return list;
    }

//    @Override
//    public List<String> aggregate(String collectionName, List<? extends String> pipeline) {
//        return null;
//    }

//    @Override
//    public List<String> mapReduce(String collectionName, String mapFunction, String reduceFunction) {
//        return null;
//    }

    @Override
    public void insertOne(String collectionName, String document) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document doc = Document.parse(document);
        collection.insertOne(doc);
    }

    @Override
    public void insertMany(String collectionName, List<String> documents) {
        insertMany(collectionName, documents, null);
    }

    /**
     * @param collectionName
     * @param documents      the documents to insert
     * @param options        the options to apply to the operation
     *                       {
     *                       orderd:<orderd>
     *                       }
     */
    @Override
    public void insertMany(String collectionName, List<String> documents, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        List<Document> list = new ArrayList<>();
        for (String document : documents) {
            Document doc = Document.parse(document);
            list.add(doc);
        }

        InsertManyOptions insertManyOptions = new InsertManyOptions();
        if (options != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(options, JsonNode.class);
                insertManyOptions.ordered(rootNode.path("ordered").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        collection.insertMany(list, insertManyOptions);
    }

    @Override
    public long deleteOne(String collectionName, String filter) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        DeleteResult deleteResult = collection.deleteOne(filterDocument);
        return deleteResult.getDeletedCount();
    }

    @Override
    public long deleteMany(String collectionName, String filter) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        DeleteResult deleteResult = collection.deleteMany(filterDocument);
        return deleteResult.getDeletedCount();
    }

    @Override
    public long replaceOne(String collectionName, String filter, String replacement) {
        return replaceOne(collectionName, filter, replacement, null);
    }

    /**
     * @param collectionName
     * @param filter         the query filter to apply the the replace operation
     * @param replacement    the replacement document
     * @param updateOptions  the options to apply to the replace operation
     *                       {
     *                       upsert:<upsert>
     *                       }
     * @return
     */
    @Override
    public long replaceOne(String collectionName, String filter, String replacement, String updateOptions) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        Document document = Document.parse(replacement);
        UpdateOptions options = new UpdateOptions();
        if (updateOptions != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(updateOptions, JsonNode.class);
                options.upsert(rootNode.path("upsert").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        UpdateResult updateResult = collection.replaceOne(filterDocument, document, options);
        return updateResult.getModifiedCount();
    }

    @Override
    public long updateOne(String collectionName, String filter, String update) {
        return updateOne(collectionName, filter, update, null);
    }

    /**
     * @param collectionName
     * @param filter         a document describing the query filter, which may not be null.
     * @param update         a document describing the update, which may not be null. The update to apply must include only update operators.
     * @param updateOptions  the options to apply to the update operation
     *                       {
     *                       upsert:<upsert>
     *                       }
     * @return
     */
    @Override
    public long updateOne(String collectionName, String filter, String update, String updateOptions) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        Document document = Document.parse(update);
        UpdateOptions options = new UpdateOptions();
        if (updateOptions != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(updateOptions, JsonNode.class);
                options.upsert(rootNode.path("upsert").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        UpdateResult updateResult = collection.updateOne(filterDocument, document, options);
        return updateResult.getModifiedCount();
    }

    @Override
    public long updateMany(String collectionName, String filter, String update) {
        return updateMany(collectionName, filter, update, null);
    }

    /**
     * @param collectionName
     * @param filter         a document describing the query filter, which may not be null.
     * @param update         a document describing the update, which may not be null. The update to apply must include only update operators.
     * @param updateOptions  the options to apply to the update operation
     *                       {
     *                       upsert:<upsert>
     *                       }
     * @return
     */
    @Override
    public long updateMany(String collectionName, String filter, String update, String updateOptions) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        Document document = Document.parse(update);
        UpdateOptions options = new UpdateOptions();
        if (updateOptions != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(updateOptions, JsonNode.class);
                options.upsert(rootNode.path("upsert").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        UpdateResult updateResult = collection.updateMany(filterDocument, document, options);
        return updateResult.getModifiedCount();
    }

    @Override
    public String findOneAndDelete(String collectionName, String filter) {
        return findOneAndDelete(collectionName, filter, null);
    }

    /**
     * @param collectionName
     * @param filter         the query filter to find the document with
     * @param options        the options to apply to the operation
     *                       {
     *                       projection:<document>,
     *                       sort:<document>
     *                       }
     * @return
     */
    @Override
    public String findOneAndDelete(String collectionName, String filter, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        FindOneAndDeleteOptions findOneAndDeleteOptions = new FindOneAndDeleteOptions();
        if (options != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(options, JsonNode.class);
                String projection = rootNode.path("projection").toString();
                Document projectionDoc = new Document();
                if (!StringUtils.isEmpty(projection)) {
                    Document.parse(projection);
                }

                String sort = rootNode.path("sort").toString();
                Document sortDoc = new Document();
                if (!StringUtils.isEmpty(sort)) {
                    Document.parse(sort);
                }

                findOneAndDeleteOptions.projection(projectionDoc);
                findOneAndDeleteOptions.sort(sortDoc);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        Document document = collection.findOneAndDelete(filterDocument, findOneAndDeleteOptions);
        return document == null ? "{}" : document.toJson();
    }

    @Override
    public String findOneAndReplace(String collectionName, String filter, String replacement) {
        return findOneAndReplace(collectionName, filter, replacement, null);
    }

    /**
     * @param collectionName
     * @param filter         the query filter to apply the the replace operation
     * @param replacement    the replacement document
     * @param options        the options to apply to the operation
     *                       {
     *                       projection:<document>,
     *                       sort:<document>,
     *                       upsert:<upsert>
     *                       }
     * @return
     */
    @Override
    public String findOneAndReplace(String collectionName, String filter, String replacement, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        Document replacementDocument = Document.parse(replacement);
        FindOneAndReplaceOptions findOneAndReplaceOptions = new FindOneAndReplaceOptions();
        if (options != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(options, JsonNode.class);
                String projection = rootNode.path("projection").toString();
                Document projectionDoc = new Document();
                if (!StringUtils.isEmpty(projection)) {
                    Document.parse(projection);
                }

                String sort = rootNode.path("sort").toString();
                Document sortDoc = new Document();
                if (!StringUtils.isEmpty(sort)) {
                    Document.parse(sort);
                }

                findOneAndReplaceOptions.projection(projectionDoc);
                findOneAndReplaceOptions.sort(sortDoc);
                findOneAndReplaceOptions.upsert(rootNode.path("upsert").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        Document document = collection.findOneAndReplace(filterDocument, replacementDocument, findOneAndReplaceOptions);
        return document == null ? "{}" : document.toJson();
    }

    @Override
    public String findOneAndUpdate(String collectionName, String filter, String update) {
        return findOneAndUpdate(collectionName, filter, update, null);
    }

    /**
     * @param collectionName
     * @param filter         a document describing the query filter, which may not be null.
     * @param update         a document describing the update, which may not be null. The update to apply must include only update operators.
     * @param options        the options to apply to the operation
     *                       {
     *                       projection:<document>,
     *                       sort:<document>,
     *                       upsert:<upsert>
     *                       }
     * @return
     */
    @Override
    public String findOneAndUpdate(String collectionName, String filter, String update, String options) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document filterDocument = Document.parse(filter);
        Document updateDocument = Document.parse(update);
        FindOneAndUpdateOptions findOneAndUpdateOptions = new FindOneAndUpdateOptions();
        if (options != null) {
            ObjectMapper mapper = new ObjectMapper();
            try {
                JsonNode rootNode = mapper.readValue(options, JsonNode.class);
                String projection = rootNode.path("projection").asText();
                Document projectionDoc = Document.parse(projection);
                String sort = rootNode.path("sort").asText();
                Document sortDoc = Document.parse(sort);

                findOneAndUpdateOptions.projection(projectionDoc);
                findOneAndUpdateOptions.sort(sortDoc);
                findOneAndUpdateOptions.upsert(rootNode.path("upsert").asBoolean());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        Document document = collection.findOneAndUpdate(filterDocument, updateDocument, findOneAndUpdateOptions);
        return document.toJson();
    }

    @Override
    public void drop(String collectionName) {
        getCollection(collectionName).drop();
    }

    @Override
    public String createIndex(String collectionName, String keys) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document keysDocument = Document.parse(keys);
        return collection.createIndex(keysDocument);
    }

    @Override
    public String createIndex(String collectionName, String keys, String indexOptions) {
        MongoCollection<Document> collection = getCollection(collectionName);
        Document keysDocument = Document.parse(keys);

        IndexOptions options = new IndexOptions();
//TODO:解析indexOptions
//        try {
//            ObjectMapper mapper = new ObjectMapper();
//            JsonNode rootNode = mapper.readValue(indexOptions, JsonNode.class);
//
//
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
        return collection.createIndex(keysDocument, options);
    }

    @Override
    public List<String> listIndexes(String collectionName) {
        MongoCollection<Document> collection = getCollection(collectionName);
        ListIndexesIterable<Document> indexes = collection.listIndexes();

        List<String> list = new ArrayList<>();
        try (MongoCursor<Document> cursor = indexes.iterator()) {
            while (cursor.hasNext()) {
                Document doc = cursor.next();
                list.add(doc.toJson());
            }
        }

        return list;
    }

    @Override
    public void dropIndex(String collectionName, String indexName) {
        getCollection(collectionName).dropIndex(indexName);
    }

    @Override
    public void dropIndexes(String collectionName) {
        getCollection(collectionName).dropIndexes();
    }

    @Override
    public void renameCollection(String collectionName, String newCollectionName) {
        MongoCollection<Document> collection = getCollection(collectionName);
        MongoNamespace namespace = collection.getNamespace();
        collection.renameCollection(new MongoNamespace(namespace.getDatabaseName(), newCollectionName));
    }

    @Override
    public void renameCollection(String collectionName, String newCollectionName, String renameCollectionOptions) {
        MongoCollection<Document> collection = getCollection(collectionName);
        MongoNamespace namespace = collection.getNamespace();
        RenameCollectionOptions options = new RenameCollectionOptions();
        try {
            ObjectMapper mapper = new ObjectMapper();
            JsonNode rootNode = mapper.readValue(renameCollectionOptions, JsonNode.class);

            options.dropTarget(rootNode.path("dropTarget").asBoolean());
        } catch (IOException e) {
            e.printStackTrace();
        }

        collection.renameCollection(new MongoNamespace(namespace.getDatabaseName(), newCollectionName), options);
    }

    public MongoCollection<Document> getCollection(String collectionName) {
        MongoDatabase database = config.getDatabase();
        return database.getCollection(collectionName);
    }

    public List<String> listCollectionNames() {

        MongoDatabase database = config.getDatabase();
        MongoIterable<String> listCollectionNames = database.listCollectionNames();

        List<String> list = new ArrayList<>();
        for (String collectionName : listCollectionNames) {
            list.add(collectionName);
        }

        return list;

    }
}