org.graylog2.database.PersistedServiceImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of graylog2-server Show documentation
Graylog
There is a newer version: 6.1.4
/*
 * Copyright (C) 2020 Graylog, Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the Server Side Public License, version 1,
 * as published by MongoDB, Inc.
 *
 * 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
 * Server Side Public License for more details.
 *
 * You should have received a copy of the Server Side Public License
 * along with this program. If not, see
 * .
 */
package org.graylog2.database;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.mongodb.BasicDBObject;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.client.MongoCollection;
import org.bson.Document;
import org.bson.types.ObjectId;
import org.graylog2.plugin.database.EmbeddedPersistable;
import org.graylog2.plugin.database.Persisted;
import org.graylog2.plugin.database.PersistedService;
import org.graylog2.plugin.database.ValidationException;
import org.graylog2.plugin.database.validators.ValidationResult;
import org.graylog2.plugin.database.validators.Validator;
import org.graylog2.plugin.system.NodeId;
import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class PersistedServiceImpl implements PersistedService {
    private static final Logger LOG = LoggerFactory.getLogger(PersistedServiceImpl.class);
    public final MongoConnection mongoConnection;

    protected PersistedServiceImpl(MongoConnection mongoConnection) {
        this.mongoConnection = mongoConnection;
    }

    protected DBObject get(ObjectId id, String collectionName) {
        return collection(collectionName).findOne(new BasicDBObject("_id", id));
    }

    protected  DBObject get(Class modelClass, ObjectId id) {
        return collection(modelClass).findOne(new BasicDBObject("_id", id));
    }

    protected  DBObject get(Class modelClass, String id) {
        return get(modelClass, new ObjectId(id));
    }

    protected List query(DBObject query, String collectionName) {
        return query(query, collection(collectionName));
    }

    protected List query(DBObject query, DBCollection collection) {
        return cursorToList(collection.find(query));
    }

    protected  List query(Class modelClass, DBObject query) {
        return query(query, collection(modelClass));
    }

    protected  List query(Class modelClass, DBObject query, DBObject sort) {
        return cursorToList(collection(modelClass).find(query).sort(sort));
    }

    protected  List query(Class modelClass, DBObject query, DBObject sort, int limit, int offset) {
        return cursorToList(
                collection(modelClass)
                        .find(query)
                        .sort(sort)
                        .limit(limit)
                        .skip(offset)
        );
    }

    protected long count(DBObject query, String collectionName) {
        return collection(collectionName).count(query);
    }

    protected  long count(Class modelClass, DBObject query) {
        return collection(modelClass).count(query);
    }

    private DBCollection collection(String collectionName) {
        return mongoConnection.getDatabase().getCollection(collectionName);
    }

    protected  DBCollection collection(Class modelClass) {
        return collection(collectionName(modelClass));
    }

    protected  String collectionName(final Class modelClass) {
        DbEntity dbEntityAnnotation = modelClass.getAnnotation(DbEntity.class);
        if (dbEntityAnnotation == null) {
            CollectionName collectionNameAnnotation = modelClass.getAnnotation(CollectionName.class);
            if (collectionNameAnnotation == null) {
                throw new RuntimeException("Unable to determine collection for class " + modelClass.getCanonicalName());
            } else {
                return collectionNameAnnotation.value();
            }
        } else {
            return dbEntityAnnotation.collection();
        }
    }

    protected  DBCollection collection(T model) {
        return collection(model.getClass());
    }

    protected  MongoCollection mongoCollection(final Class modelClass) {
        return mongoConnection.getMongoDatabase().getCollection(collectionName(modelClass));
    }

    protected List cursorToList(DBCursor cursor) {
        if (cursor == null) {
            return Collections.emptyList();
        }

        try {
            return Lists.newArrayList((Iterable) cursor);
        } finally {
            cursor.close();
        }
    }

    protected  DBObject findOne(Class model, DBObject query) {
        return collection(model).findOne(query);
    }

    protected  DBObject findOne(Class model, DBObject query, DBObject sort) {
        return collection(model).findOne(query, new BasicDBObject(), sort);
    }

    protected DBObject findOne(DBObject query, String collectionName) {
        return collection(collectionName).findOne(query);
    }

    protected DBObject findOne(DBObject query, DBObject sort, String collectioName) {
        return collection(collectioName).findOne(query, new BasicDBObject(), sort);
    }

    protected long totalCount(String collectionName) {
        return collection(collectionName).count();
    }

    protected  long totalCount(Class modelClass) {
        return collection(modelClass).count();
    }

    @Override
    public  int destroy(T model) {
        return collection(model).remove(new BasicDBObject("_id", new ObjectId(model.getId()))).getN();
    }

    @Override
    public  int destroyAll(Class modelClass) {
        return collection(modelClass).remove(new BasicDBObject()).getN();
    }

    protected int destroyAll(String collectionName) {
        return collection(collectionName).remove(new BasicDBObject()).getN();
    }

    protected int destroy(DBObject query, String collectionName) {
        return collection(collectionName).remove(query).getN();
    }

    protected  int destroyAll(Class modelClass, DBObject query) {
        return collection(modelClass).remove(query).getN();
    }

    @Override
    public  String save(T model) throws ValidationException {
        Map> errors = validate(model, model.getFields());
        if (!errors.isEmpty()) {
            throw new ValidationException(errors);
        }

        BasicDBObject doc = new BasicDBObject(model.getFields());
        doc.put("_id", new ObjectId(model.getId())); // ID was created in constructor or taken from original doc already.

        // Do field transformations
        fieldTransformations(doc);

        /*
         * We are running an upsert. This means that the existing
         * document will be updated if the ID already exists and
         * a new document will be created if it doesn't.
         */
        BasicDBObject q = new BasicDBObject("_id", new ObjectId(model.getId()));
        collection(model).update(q, doc, true, false);

        return model.getId();
    }

    @Override
    public  String saveWithoutValidation(T model) {
        try {
            return save(model);
        } catch (ValidationException ignored) {
            return null;
        }
    }

    @Override
    public  Map> validate(T model, Map fields) {
        return validate(model.getValidations(), fields);
    }

    @Override
    public Map> validate(Map validators, Map fields) {
        if (validators == null || validators.isEmpty()) {
            return Collections.emptyMap();
        }

        final Map> validationErrors = new HashMap<>();
        for (Map.Entry validation : validators.entrySet()) {
            Validator v = validation.getValue();
            String field = validation.getKey();

            try {
                ValidationResult validationResult = v.validate(fields.get(field));
                if (validationResult instanceof ValidationResult.ValidationFailed) {
                    LOG.debug("Validation failure: [{}] on field [{}]", v.getClass().getCanonicalName(), field);
                    validationErrors.computeIfAbsent(field, k -> new ArrayList<>());
                    validationErrors.get(field).add(validationResult);
                }
            } catch (Exception e) {
                final String error = "Error while trying to validate <" + field + ">, got exception: " + e;
                LOG.debug(error);
                validationErrors.computeIfAbsent(field, k -> new ArrayList<>());
                validationErrors.get(field).add(new ValidationResult.ValidationFailed(error));
            }
        }

        return validationErrors;
    }

    @Override
    public  Map> validate(T model) {
        return validate(model, model.getFields());
    }

    protected  void embed(T model, String key, EmbeddedPersistable o) throws ValidationException {
        Map> errors = validate(model.getEmbeddedValidations(key), o.getPersistedFields());
        if (!errors.isEmpty()) {
            throw new ValidationException(errors);
        }

        Map fields = Maps.newHashMap(o.getPersistedFields());
        fieldTransformations(fields);

        BasicDBObject dbo = new BasicDBObject(fields);
        collection(model).update(new BasicDBObject("_id", new ObjectId(model.getId())), new BasicDBObject("$push", new BasicDBObject(key, dbo)));
    }

    protected  void removeEmbedded(T model, String key, String searchId) {
        BasicDBObject aryQry = new BasicDBObject("id", searchId);
        BasicDBObject qry = new BasicDBObject("_id", new ObjectId(model.getId()));
        BasicDBObject update = new BasicDBObject("$pull", new BasicDBObject(key, aryQry));

        // http://docs.mongodb.org/manual/reference/operator/pull/

        collection(model).update(qry, update);
    }

    protected  void removeEmbedded(T model, String arrayKey, String key, String searchId) {
        BasicDBObject aryQry = new BasicDBObject(arrayKey, searchId);
        BasicDBObject qry = new BasicDBObject("_id", new ObjectId(model.getId()));
        BasicDBObject update = new BasicDBObject("$pull", new BasicDBObject(key, aryQry));

        // http://docs.mongodb.org/manual/reference/operator/pull/

        collection(model).update(qry, update);
    }

    protected void fieldTransformations(Map doc) {
        for (Map.Entry x : doc.entrySet()) {

            // Work on embedded Maps, too.
            if (x.getValue() instanceof Map) {
                x.setValue(Maps.newHashMap((Map) x.getValue()));
                fieldTransformations((Map) x.getValue());
                continue;
            }

            // JodaTime DateTime is not accepted by MongoDB. Convert to java.util.Date...
            if (x.getValue() instanceof DateTime) {
                doc.put(x.getKey(), ((DateTime) x.getValue()).toDate());
            }

            // Our own NodeID
            if (x.getValue() instanceof NodeId nodeId) {
                doc.put(x.getKey(), nodeId.getNodeId());
            }
        }
    }

}