org.graylog2.streams.OutputServiceImpl 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.streams;

import com.google.common.collect.ImmutableSet;
import com.mongodb.BasicDBObject;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import org.bson.types.ObjectId;
import org.graylog2.bindings.providers.MongoJackObjectMapperProvider;
import org.graylog2.database.DbEntity;
import org.graylog2.database.MongoConnection;
import org.graylog2.database.NotFoundException;
import org.graylog2.events.ClusterEventBus;
import org.graylog2.outputs.events.OutputChangedEvent;
import org.graylog2.plugin.Tools;
import org.graylog2.plugin.database.ValidationException;
import org.graylog2.plugin.streams.Output;
import org.graylog2.rest.models.streams.outputs.requests.CreateOutputRequest;
import org.mongojack.DBQuery;
import org.mongojack.DBUpdate;
import org.mongojack.JacksonDBCollection;
import org.mongojack.WriteResult;

import javax.inject.Inject;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class OutputServiceImpl implements OutputService {
    private final JacksonDBCollection coll;
    private final DBCollection dbCollection;
    private final StreamService streamService;
    private final ClusterEventBus clusterEventBus;

    @Inject
    public OutputServiceImpl(MongoConnection mongoConnection,
                             MongoJackObjectMapperProvider mapperProvider,
                             StreamService streamService,
                             ClusterEventBus clusterEventBus) {
        this.streamService = streamService;
        final String collectionName = OutputImpl.class.getAnnotation(DbEntity.class).collection();
        this.dbCollection = mongoConnection.getDatabase().getCollection(collectionName);
        this.coll = JacksonDBCollection.wrap(dbCollection, OutputImpl.class, String.class, mapperProvider.get());
        this.clusterEventBus = clusterEventBus;
    }

    @Override
    public Output load(String streamOutputId) throws NotFoundException {
        final Output output = coll.findOneById(streamOutputId);
        if (output == null) {
            throw new NotFoundException("Couldn't find output with id " + streamOutputId);
        }

        return output;
    }

    @Override
    public Set loadAll() {
        try (org.mongojack.DBCursor outputs = coll.find()) {
            return ImmutableSet.copyOf((Iterable) outputs);
        }
    }

    @Override
    public Set loadByIds(Collection ids) {
        final DBQuery.Query query = DBQuery.in(OutputImpl.FIELD_ID, ids);
        try (org.mongojack.DBCursor dbCursor = coll.find(query)) {
            return ImmutableSet.copyOf((Iterable) dbCursor);
        }
    }

    @Override
    public Output create(Output request) throws ValidationException {
        final OutputImpl outputImpl = implOrFail(request);
        final WriteResult writeResult = coll.save(outputImpl);

        return writeResult.getSavedObject();
    }

    @Override
    public Output create(CreateOutputRequest request, String userId) throws ValidationException {
        return create(OutputImpl.create(new ObjectId().toHexString(), request.title(), request.type(), userId, request.configuration(),
                Tools.nowUTC().toDate(), request.contentPack()));
    }

    @Override
    public void destroy(Output model) throws NotFoundException {
        coll.removeById(model.getId());

        // Removing the output from all streams will emit a StreamsChangedEvent for affected streams.
        // The OutputRegistry will handle this event and stop the output.
        streamService.removeOutputFromAllStreams(model);
    }

    @Override
    public Output update(String id, Map deltas) {
        DBUpdate.Builder update = new DBUpdate.Builder();
        for (Map.Entry fields : deltas.entrySet()) {
            update = update.set(fields.getKey(), fields.getValue());
        }

        final OutputImpl updatedOutput =
                coll.findAndModify(DBQuery.is(OutputImpl.FIELD_ID, id), null, null, false, update, true, false);

        this.clusterEventBus.post(OutputChangedEvent.create(updatedOutput.getId()));

        return updatedOutput;
    }

    @Override
    public long count() {
        return coll.count();
    }

    @Override
    public Map countByType() {
        final Map outputsCountByType = new HashMap<>();
        try (DBCursor outputTypes = dbCollection.find(null, new BasicDBObject(OutputImpl.FIELD_TYPE, 1))) {

            for (DBObject outputType : outputTypes) {
                final String type = (String) outputType.get(OutputImpl.FIELD_TYPE);
                if (type != null) {
                    final Long oldValue = outputsCountByType.get(type);
                    final Long newValue = (oldValue == null) ? 1 : oldValue + 1;
                    outputsCountByType.put(type, newValue);
                }
            }
        }

        return outputsCountByType;
    }

    private OutputImpl implOrFail(Output output) {
        final OutputImpl outputImpl;
        if (output instanceof OutputImpl) {
            outputImpl = (OutputImpl) output;
            return outputImpl;
        } else {
            throw new IllegalArgumentException("Supplied output must be of implementation type OutputImpl, not " + output.getClass());
        }
    }
}