Add utility methods for MongoDB collections (#18944)

* Add support for pagination * Add utils to extract inserted id * Add stream utility method * Use stream utility method * Add "byId" convenience methods * Copy tests from PaginatedDbService * Add license * Renaming and JavaDoc * Separate helper/utils from client API * add license * Remove TODO * Support streaming of any mongo iterable * interim support for deprecated mongojack bson objects * Use shorter method names instead of bean conventions * Remove dispensable interface for utility class * Introduce MongoEntity interface * Add @nullable annotation * Add idEq helper function * Don't return nulls from #insertedId * Use idEq in utils * Add tests * Rename postProcessedPage to page * Make pagination helper immutable * Add collation support to pagination helper * Change wording Co-authored-by: Bernd Ahlers <[email protected]> * Change wording Co-authored-by: Bernd Ahlers <[email protected]> * wording Co-authored-by: Bernd Ahlers <[email protected]> * Add SortOrder enum * Remove redundant sort method from pagination helper * Use SortOrder in PageListResponse * Document immutability of pagination helper * Add locale to call of toLowerCase * Fix wrong order type used in typescript mock --------- Co-authored-by: Bernd Ahlers <[email protected]>
Graylog2 · May 8, 2024 · 1967c43 · 1967c43
1 parent 4a1cbb4
commit 1967c43
Show file tree

Hide file tree

Showing 13 changed files with 1,248 additions and 5 deletions.
diff --git a/graylog2-server/src/main/java/org/graylog2/database/MongoCollections.java b/graylog2-server/src/main/java/org/graylog2/database/MongoCollections.java
@@ -22,6 +22,9 @@
 import jakarta.inject.Singleton;
 import org.graylog2.bindings.providers.MongoJackObjectMapperProvider;
 import org.graylog2.database.jackson.CustomJacksonCodecRegistry;
+import org.graylog2.database.pagination.DefaultMongoPaginationHelper;
+import org.graylog2.database.pagination.MongoPaginationHelper;
+import org.graylog2.database.utils.MongoUtils;
 
 @Singleton
 public class MongoCollections {
@@ -37,6 +40,12 @@ public MongoCollections(MongoJackObjectMapperProvider objectMapperProvider, Mong
 
     /**
      * Get a MongoCollection configured to use Jackson for serialization/deserialization of objects.
+     * <p>
+     * <b>Prefer using {@link #collection(String, Class)} to get a more strictly typed collection!</b>
+     *
+     * @param collectionName Name of the collection
+     * @param valueType      Java type of the documents stored in the collection
+     * @return A collection using a Jackson codec for serialization and deserialization
      */
     public <T> MongoCollection<T> get(String collectionName, Class<T> valueType) {
         final MongoCollection<T> collection = mongoConnection.getMongoDatabase().getCollection(collectionName, valueType);
@@ -45,4 +54,43 @@ public <T> MongoCollection<T> get(String collectionName, Class<T> valueType) {
         jacksonCodecRegistry.addCodecForClass(valueType);
         return collection.withCodecRegistry(jacksonCodecRegistry);
     }
+
+    /**
+     * Get a MongoCollection configured to use Jackson for serialization/deserialization of objects.
+     *
+     * @param collectionName Name of the collection
+     * @param valueType      Java type of the documents stored in the collection
+     * @return A collection using a Jackson codec for serialization and deserialization
+     */
+    public <T extends MongoEntity> MongoCollection<T> collection(String collectionName, Class<T> valueType) {
+        return get(collectionName, valueType);
+    }
+
+    /**
+     * Provides a helper to perform find operations on a collection that yield pages of documents.
+     */
+    public <T extends MongoEntity> MongoPaginationHelper<T> paginationHelper(String collectionName, Class<T> valueType) {
+        return paginationHelper(collection(collectionName, valueType));
+    }
+
+    /**
+     * Provides a helper to perform find operations on a collection that yield pages of documents.
+     */
+    public <T extends MongoEntity> MongoPaginationHelper<T> paginationHelper(MongoCollection<T> collection) {
+        return new DefaultMongoPaginationHelper<>(collection);
+    }
+
+    /**
+     * Provides utility methods like getting documents by ID, etc.
+     */
+    public <T extends MongoEntity> MongoUtils<T> utils(String collectionName, Class<T> valueType) {
+        return utils(collection(collectionName, valueType));
+    }
+
+    /**
+     * Provides utility methods like getting documents by ID, etc.
+     */
+    public <T extends MongoEntity> MongoUtils<T> utils(MongoCollection<T> collection) {
+        return new MongoUtils<>(collection, objectMapper);
+    }
 }
diff --git a/graylog2-server/src/main/java/org/graylog2/database/MongoEntity.java b/graylog2-server/src/main/java/org/graylog2/database/MongoEntity.java
@@ -0,0 +1,39 @@
+/*
+ * 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
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog2.database;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.annotation.Nullable;
+import org.mongojack.Id;
+import org.mongojack.ObjectId;
+
+/**
+ * Common interface for entities stored in MongoDB.
+ */
+public interface MongoEntity {
+
+    /**
+     * ID of the entity. Will be stored as field "_id" with type ObjectId in MongoDB.
+     *
+     * @return Hex string representation of the entity's ID
+     */
+    @Nullable
+    @ObjectId
+    @Id
+    @JsonProperty("id")
+    String id();
+}
diff --git a/...2-server/src/main/java/org/graylog2/database/pagination/DefaultMongoPaginationHelper.java b/...2-server/src/main/java/org/graylog2/database/pagination/DefaultMongoPaginationHelper.java
@@ -0,0 +1,145 @@
+/*
+ * 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
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog2.database.pagination;
+
+import com.google.common.primitives.Ints;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.Collation;
+import org.bson.conversions.Bson;
+import org.graylog2.database.MongoEntity;
+import org.graylog2.database.PaginatedList;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Predicate;
+
+import static org.graylog2.database.utils.MongoUtils.stream;
+
+/**
+ * Default implementation for pagination support.
+ * <p>
+ * Objects of this class are immutable and can be re-used.
+ *
+ * @param <T> Type of documents in the underlying MongoDB collection.
+ */
+public class DefaultMongoPaginationHelper<T extends MongoEntity> implements MongoPaginationHelper<T> {
+
+    private final MongoCollection<T> collection;
+    private final Bson filter;
+    private final Bson sort;
+    private final int perPage;
+    private final boolean includeGrandTotal;
+    private final Bson grandTotalFilter;
+    private final Collation collation;
+
+    public DefaultMongoPaginationHelper(MongoCollection<T> collection) {
+        this(collection, null, null, 0, false, null, null);
+    }
+
+    private DefaultMongoPaginationHelper(MongoCollection<T> collection, Bson filter, Bson sort, int perPage,
+                                         boolean includeGrandTotal, Bson grandTotalFilter, Collation collation) {
+        this.collection = collection;
+        this.filter = filter;
+        this.sort = sort;
+        this.perPage = perPage;
+        this.includeGrandTotal = includeGrandTotal;
+        this.grandTotalFilter = grandTotalFilter;
+        this.collation = collation;
+    }
+
+    @Override
+    public MongoPaginationHelper<T> filter(Bson filter) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public MongoPaginationHelper<T> sort(Bson sort) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public MongoPaginationHelper<T> perPage(int perPage) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public MongoPaginationHelper<T> includeGrandTotal(boolean includeGrandTotal) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public MongoPaginationHelper<T> grandTotalFilter(Bson grandTotalFilter) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public MongoPaginationHelper<T> collation(Collation collation) {
+        return new DefaultMongoPaginationHelper<>(collection, filter, sort, perPage, includeGrandTotal,
+                grandTotalFilter, collation);
+    }
+
+    @Override
+    public PaginatedList<T> page(int pageNumber) {
+        final List<T> documents = collection.find()
+                .filter(filter)
+                .sort(sort)
+                .skip(perPage * Math.max(0, pageNumber - 1))
+                .limit(perPage)
+                .collation(collation)
+                .into(new ArrayList<>());
+        final int total = Ints.saturatedCast(collection.countDocuments(filter));
+
+        if (includeGrandTotal) {
+            final long grandTotal = collection.countDocuments(grandTotalFilter);
+            return new PaginatedList<>(documents, total, pageNumber, perPage, grandTotal);
+        } else {
+            return new PaginatedList<>(documents, total, pageNumber, perPage);
+        }
+    }
+
+    @Override
+    public PaginatedList<T> page(int pageNumber, Predicate<T> selector) {
+        final int total = Ints.saturatedCast(stream(collection.find()
+                .filter(filter)
+                .sort(sort)).filter(selector).count());
+
+        final List<T> documents;
+        if (perPage > 0) {
+            documents = stream(collection.find().filter(filter).sort(sort).collation(collation))
+                    .filter(selector)
+                    .skip(perPage * Math.max(0L, pageNumber - 1))
+                    .limit(perPage)
+                    .toList();
+        } else {
+            documents = stream(collection.find().filter(filter).sort(sort).collation(collation))
+                    .filter(selector).toList();
+        }
+
+        if (includeGrandTotal) {
+            final long grandTotal = collection.countDocuments(grandTotalFilter);
+            return new PaginatedList<>(documents, total, pageNumber, perPage, grandTotal);
+        } else {
+            return new PaginatedList<>(documents, total, pageNumber, perPage);
+        }
+    }
+
+}
diff --git a/graylog2-server/src/main/java/org/graylog2/database/pagination/MongoPaginationHelper.java b/graylog2-server/src/main/java/org/graylog2/database/pagination/MongoPaginationHelper.java
@@ -0,0 +1,105 @@
+/*
+ * 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
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+package org.graylog2.database.pagination;
+
+import com.mongodb.client.model.Collation;
+import org.bson.conversions.Bson;
+import org.graylog2.database.MongoEntity;
+import org.graylog2.database.PaginatedList;
+
+import java.util.function.Predicate;
+
+/**
+ * A utility class that provides paged access to a MongoDB collection.
+ * <p>
+ * Implementing classes should be immutable
+ * so that instances can be re-used.
+ *
+ * @param <T> Type of documents in the underlying MongoDB collection.
+ */
+public interface MongoPaginationHelper<T extends MongoEntity> {
+    /**
+     * Sets the query filter to apply to the query.
+     *
+     * @param filter the filter, which may be null.
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> filter(Bson filter);
+
+    /**
+     * Sets the sort criteria to apply to the query.
+     *
+     * @param sort the sort criteria, which may be null.
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> sort(Bson sort);
+
+    /**
+     * Sets the page size.
+     *
+     * @param perPage the number of documents to put on one page
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> perPage(int perPage);
+
+    /**
+     * Specifies whether to include a grand total number of all documents in the collection. No filters, except, if set,
+     * the {@link #grandTotalFilter(Bson)} will be applied to the count query.
+     *
+     * @param includeGrandTotal true if a grand total should be included. Otherwise, by default, no grand total will
+     *                          be included.
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> includeGrandTotal(boolean includeGrandTotal);
+
+    /**
+     * Sets a filter to be applied to the query to count the grand total of documents in the collection.
+     *
+     * @param grandTotalFilter the filter, which may be null
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> grandTotalFilter(Bson grandTotalFilter);
+
+    /**
+     * Sets a collation to be used in the find operation.
+     *
+     * @param collation The collation to set. If null, uses the default server collation
+     * @return A new pagination helper with the setting applied
+     */
+    MongoPaginationHelper<T> collation(Collation collation);
+
+    /**
+     * Perform the MongoDB request and return the specified page.
+     *
+     * @param pageNumber The number of the page to be returned.
+     * @return a paginated list of documents
+     */
+    PaginatedList<T> page(int pageNumber);
+
+    /**
+     * Perform the MongoDB request but apply the given predicate to each document in the list of results and only keep
+     * documents matching the predicate.
+     * <p>
+     * <b>This is a potentially expensive operation because the selector can only be applied after the documents
+     * have been fetched from MongoDB and this might result in a full collection scan. Use with care!</b>
+     *
+     * @param pageNumber The number of the page to be returned.
+     * @param selector   predicate to filter documents <b>after</b> fetching them from MongoDB.
+     * @return a paginated list of documents
+     */
+    PaginatedList<T> page(int pageNumber, Predicate<T> selector);
+}