org.graylog2.streams.PaginatedStreamService 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.ImmutableList;
import com.mongodb.client.AggregateIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.Aggregates;
import com.mongodb.client.model.Field;
import com.mongodb.client.model.Variable;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.graylog2.bindings.providers.MongoJackObjectMapperProvider;
import org.graylog2.database.MongoConnection;
import org.graylog2.database.PaginatedDbService;
import org.graylog2.database.PaginatedList;
import org.graylog2.indexer.indexset.MongoIndexSetService;
import org.mongojack.DBQuery;

import javax.inject.Inject;
import java.util.List;
import java.util.function.Predicate;
import java.util.stream.StreamSupport;

public class PaginatedStreamService extends PaginatedDbService {
    private static final String COLLECTION_NAME = "streams";
    private static final List STRING_FIELDS = List.of("title", "description", "index_set_title");
    private final MongoCollection collection;

    @Inject
    public PaginatedStreamService(MongoConnection mongoConnection,
                                  MongoJackObjectMapperProvider mapper) {
        super(mongoConnection, mapper, StreamDTO.class, COLLECTION_NAME);
        this.collection = mongoConnection.getMongoDatabase().getCollection(COLLECTION_NAME);
    }

    public long count() {
        return db.count();
    }

    public PaginatedList findPaginated(Bson dbQuery, //query executed on DB level
                                                  Predicate predicate, //predicate executed on code level, AFTER data is fetched
                                                  int page,
                                                  int perPage,
                                                  String sortField,
                                                  String order) {


        var pipelineBuilder = ImmutableList.builder()
                .add(Aggregates.match(dbQuery));

        if (sortField.equals("index_set_title")) {
            pipelineBuilder.add(Aggregates.lookup(
                            MongoIndexSetService.COLLECTION_NAME,
                            List.of(new Variable<>("index_set_id", doc("$toObjectId", "$index_set_id"))),
                            List.of(Aggregates.match(doc("$expr", doc("$eq", List.of("$_id", "$$index_set_id"))))),
                            "index_set"
                    ))
                    .add(Aggregates.set(new Field<>("index_set_title", doc("$first", "$index_set.title"))))
                    .add(Aggregates.unset("index_set"));
        }

        if (isStringField(sortField)) {
            pipelineBuilder.add(Aggregates.set(new Field<>("lower" + sortField, doc("$toLower", "$" + sortField))))
                    .add(Aggregates.sort(getSortBuilder(order, "lower" + sortField)))
                    .add(Aggregates.unset("lower" + sortField));
        } else {
            pipelineBuilder.add(Aggregates.sort(getSortBuilder(order, sortField)));
        }

        final AggregateIterable result = collection.aggregate(pipelineBuilder.build());

        final List streamsList = StreamSupport.stream(result.spliterator(), false)
                .map(StreamDTO::fromDocument)
                .filter(predicate)
                .toList();

        final long grandTotal = db.find(DBQuery.empty()).toArray()
                .stream()
                .filter(predicate)
                .count();

        final List paginatedStreams = perPage > 0
                ? streamsList.stream()
                .skip((long) perPage * Math.max(0, page - 1))
                .limit(perPage)
                .toList()
                : streamsList;

        return new PaginatedList<>(paginatedStreams, streamsList.size(), page, perPage, grandTotal);

    }

    private boolean isStringField(String sortField) {
        return STRING_FIELDS.contains(sortField);
    }

    private Document doc(String key, Object value) {
        return new Document(key, value);
    }
}