diff --git a/graylog2-server/src/main/java/org/graylog2/database/BuildableMongoEntity.java b/graylog2-server/src/main/java/org/graylog2/database/BuildableMongoEntity.java new file mode 100644 index 000000000000..fdfa646b81ac --- /dev/null +++ b/graylog2-server/src/main/java/org/graylog2/database/BuildableMongoEntity.java @@ -0,0 +1,38 @@ +/* + * 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.fasterxml.jackson.annotation.JsonProperty; +import org.mongojack.Id; + +/** + * Interface that ensures that an entity can be converted to a Builder that allows setting the ID on it. + * + * @param Type of the entity that provides a #toBuilder() method + * @param Type of the builder that allows setting the ID + */ +public interface BuildableMongoEntity> extends MongoEntity { + B toBuilder(); + + interface Builder { + @Id + @JsonProperty("id") + B id(String id); + + T build(); + } +} diff --git a/graylog2-server/src/main/java/org/graylog2/database/utils/MongoUtils.java b/graylog2-server/src/main/java/org/graylog2/database/utils/MongoUtils.java index 21bb65ab123c..931cfeba87d8 100644 --- a/graylog2-server/src/main/java/org/graylog2/database/utils/MongoUtils.java +++ b/graylog2-server/src/main/java/org/graylog2/database/utils/MongoUtils.java @@ -22,6 +22,7 @@ import com.mongodb.client.MongoIterable; import com.mongodb.client.model.Filters; import com.mongodb.client.model.FindOneAndUpdateOptions; +import com.mongodb.client.model.ReplaceOptions; import com.mongodb.client.model.ReturnDocument; import com.mongodb.client.model.Updates; import com.mongodb.client.result.InsertOneResult; @@ -32,6 +33,7 @@ import org.bson.codecs.EncoderContext; import org.bson.conversions.Bson; import org.bson.types.ObjectId; +import org.graylog2.database.BuildableMongoEntity; import org.graylog2.database.MongoEntity; import org.graylog2.database.jackson.CustomJacksonCodecRegistry; import org.mongojack.InitializationRequiredForTransformation; @@ -236,6 +238,30 @@ public T getOrCreate(T entity) { } } + /** + * Saves an entity by either inserting or replacing the document. + *

+ * This method exists to avoid the repeated implementation of this functionality during migration from the old + * Mongojack API. + *

+ * For new code, prefer implementing a separate "create" and "update" path instead. + * + * @param entity Entity to be saved, with the #id() property optionally set. + * @return Saved entity with the #id() property guaranteed to be present. + */ + public T save(BuildableMongoEntity entity) { + // going through the builder is a bit more work but avoids an unsafe cast to T + final var orig = entity.toBuilder().build(); + final var id = orig.id(); + if (id == null) { + final var insertedId = insertedIdAsString(collection.insertOne(orig)); + return entity.toBuilder().id(insertedId).build(); + } else { + collection.replaceOne(idEq(id), orig, new ReplaceOptions().upsert(true)); + return orig; + } + } + /** * Utility method to help moving away from the deprecated MongoJack Bson objects, like * {@link org.mongojack.DBQuery.Query}. These objects require initialization before they can be used as regular diff --git a/graylog2-server/src/test/java/org/graylog2/database/utils/MongoUtilsTest.java b/graylog2-server/src/test/java/org/graylog2/database/utils/MongoUtilsTest.java index 38385d204812..ac25e2ba1cb6 100644 --- a/graylog2-server/src/test/java/org/graylog2/database/utils/MongoUtilsTest.java +++ b/graylog2-server/src/test/java/org/graylog2/database/utils/MongoUtilsTest.java @@ -16,6 +16,10 @@ */ package org.graylog2.database.utils; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.google.auto.value.AutoValue; import com.mongodb.client.MongoCollection; import org.bson.RawBsonDocument; import org.bson.types.ObjectId; @@ -23,6 +27,7 @@ import org.graylog.testing.mongodb.MongoDBTestService; import org.graylog.testing.mongodb.MongoJackExtension; import org.graylog2.bindings.providers.MongoJackObjectMapperProvider; +import org.graylog2.database.BuildableMongoEntity; import org.graylog2.database.MongoCollections; import org.graylog2.database.MongoEntity; import org.junit.jupiter.api.BeforeEach; @@ -180,4 +185,77 @@ void testGetOrCreateWithNullEntityID() { .hasMessageContaining("entity ID cannot be null") .isInstanceOf(NullPointerException.class); } + + @AutoValue + @JsonDeserialize(builder = AutoValueDTO.Builder.class) + public static abstract class AutoValueDTO implements BuildableMongoEntity { + @JsonProperty("name") + public abstract String name(); + + @Override + public abstract Builder toBuilder(); + + public static AutoValueDTO.Builder builder() { + return Builder.create(); + } + + @AutoValue.Builder + public abstract static class Builder implements BuildableMongoEntity.Builder { + @JsonCreator + public static Builder create() { + return new AutoValue_MongoUtilsTest_AutoValueDTO.Builder(); + } + + @JsonProperty("name") + public abstract Builder name(String name); + + public abstract AutoValueDTO build(); + } + } + + @Test + void testSaveAutoValueDTOWithoutId() { + final var coll = mongoCollections.collection("autovalue-test", AutoValueDTO.class); + final var util = mongoCollections.utils(coll); + + final var orig = AutoValueDTO.builder().name("test").build(); + assertThat(orig.id()).isNull(); + + final var saved = util.save(orig); + final var generatedId = saved.id(); + assertThat(saved) + .isEqualTo(orig.toBuilder().id(generatedId).build()) + .isEqualTo(util.getById(generatedId).orElse(null)); + } + + @Test + void testSaveAutoValueDTOWithExistingId() { + final var coll = mongoCollections.collection("autovalue-test", AutoValueDTO.class); + final var util = mongoCollections.utils(coll); + + final var existing = util.save(AutoValueDTO.builder().name("test").build()); + final var existingId = existing.id(); + assertThat(existingId).isNotNull(); + + final var orig = existing.toBuilder().name("new name").build(); + final var saved = util.save(orig); + + assertThat(saved) + .isEqualTo(orig) + .isEqualTo(util.getById(existingId).orElse(null)); + } + + @Test + void testSaveAutoValueDTOWithNewId() { + final var coll = mongoCollections.collection("autovalue-test", AutoValueDTO.class); + final var util = mongoCollections.utils(coll); + + final var orig = AutoValueDTO.builder().id(new ObjectId().toHexString()).name("test").build(); + assertThat(util.getById(orig.id())).isEmpty(); + + final var saved = util.save(orig); + assertThat(saved) + .isEqualTo(orig) + .isEqualTo(util.getById(orig.id()).orElse(null)); + } }