Add #save utility method to aid with Mongojack API migration (#20655)

* Add `#save` utility method to aid with Mongojack API migration

* Fix generics and add basic annotations to #id

* Add test

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
+ * <http://www.mongodb.com/licensing/server-side-public-license>.
+ */
+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 <T> Type of the entity that provides a #toBuilder() method
+ * @param <B> Type of the builder that allows setting the ID
+ */
+public interface BuildableMongoEntity<T, B extends BuildableMongoEntity.Builder<T, B>> extends MongoEntity {
+    B toBuilder();
+
+    interface Builder<T, B> {
+        @Id
+        @JsonProperty("id")
+        B id(String id);
+
+        T build();
+    }
+}

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.
+     * <p>
+     * This method exists to avoid the repeated implementation of this functionality during migration from the old
+     * Mongojack API.
+     * <p>
+     * <b> For new code, prefer implementing a separate "create" and "update" path instead.</b>
+     *
+     * @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<T, ?> 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

graylog2-server/src/test/java/org/graylog2/database/utils/MongoUtilsTest.java

@@ -16,13 +16,18 @@
  */
 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;
 import org.graylog.testing.mongodb.MongoDBExtension;
 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<AutoValueDTO, AutoValueDTO.Builder> {
+        @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<AutoValueDTO, 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));
+    }
 }