Remove JacksonDBCollection usage from MongoInputStatusService

Graylog2 · Jan 20, 2025 · 0d7ea58 · 0d7ea58
1 parent 48b70c1
commit 0d7ea58
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 38 deletions.
diff --git a/graylog2-server/src/main/java/org/graylog2/inputs/persistence/InputStatusRecord.java b/graylog2-server/src/main/java/org/graylog2/inputs/persistence/InputStatusRecord.java
@@ -42,6 +42,8 @@ public static Builder builder() {
         return new AutoValue_InputStatusRecord.Builder();
     }
 
+    public abstract Builder toBuilder();
+
     @AutoValue.Builder
     public static abstract class Builder {
         @Id

diff --git a/graylog2-server/src/main/java/org/graylog2/inputs/persistence/MongoInputStatusService.java b/graylog2-server/src/main/java/org/graylog2/inputs/persistence/MongoInputStatusService.java
@@ -18,22 +18,18 @@
 
 import com.google.common.eventbus.EventBus;
 import com.google.common.eventbus.Subscribe;
-import com.mongodb.DB;
-import com.mongodb.DBCollection;
-import org.bson.types.ObjectId;
-import org.graylog2.bindings.providers.MongoJackObjectMapperProvider;
-import org.graylog2.database.MongoConnection;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.ReplaceOptions;
+import jakarta.inject.Inject;
+import jakarta.inject.Singleton;
+import org.graylog2.database.MongoCollections;
 import org.graylog2.database.NotFoundException;
+import org.graylog2.database.utils.MongoUtils;
 import org.graylog2.inputs.InputService;
 import org.graylog2.rest.models.system.inputs.responses.InputDeleted;
-import org.mongojack.JacksonDBCollection;
-import org.mongojack.WriteResult;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import jakarta.inject.Inject;
-import jakarta.inject.Singleton;
-
 import java.util.Optional;
 
 /**
@@ -47,42 +43,35 @@ public class MongoInputStatusService implements InputStatusService {
 
     public static final String COLLECTION_NAME = "input_status";
 
-    private final JacksonDBCollection<InputStatusRecord, ObjectId> statusCollection;
     private final InputService inputService;
+    private final MongoCollection<InputStatusRecord> collection;
 
     @Inject
-    public MongoInputStatusService(MongoConnection mongoConnection,
-                                   MongoJackObjectMapperProvider objectMapperProvider,
-                                   InputService inputService,
-                                   EventBus eventBus) {
+    public MongoInputStatusService(MongoCollections mongoCollections, InputService inputService, EventBus eventBus) {
         this.inputService = inputService;
-        DB mongoDatabase = mongoConnection.getDatabase();
-        DBCollection collection = mongoDatabase.getCollection(COLLECTION_NAME);
+        this.collection = mongoCollections.nonEntityCollection(COLLECTION_NAME, InputStatusRecord.class);
 
         eventBus.register(this);
-
-        statusCollection = JacksonDBCollection.wrap(
-                collection,
-                InputStatusRecord.class,
-                ObjectId.class,
-                objectMapperProvider.get());
     }
 
     @Override
     public Optional<InputStatusRecord> get(final String inputId) {
-        return Optional.ofNullable(statusCollection.findOneById(new ObjectId(inputId)));
+        return Optional.ofNullable(collection.find(MongoUtils.idEq(inputId)).first());
     }
 
     @Override
     public InputStatusRecord save(InputStatusRecord statusRecord) {
-        final WriteResult<InputStatusRecord, ObjectId> save = statusCollection.save(statusRecord);
-        return save.getSavedObject();
+        if (statusRecord.inputId() == null) {
+            final var insertedId = MongoUtils.insertedIdAsString(collection.insertOne(statusRecord));
+            return statusRecord.toBuilder().inputId(insertedId).build();
+        }
+        collection.replaceOne(MongoUtils.idEq(statusRecord.inputId()), statusRecord, new ReplaceOptions().upsert(true));
+        return statusRecord;
     }
 
     @Override
     public int delete(String inputId) {
-        final WriteResult<InputStatusRecord, ObjectId> delete = statusCollection.removeById(new ObjectId(inputId));
-        return delete.getN();
+        return (int) collection.deleteOne(MongoUtils.idEq(inputId)).getDeletedCount();
     }
 
     /**

diff --git a/...og2-server/src/test/java/org/graylog2/inputs/persistence/MongoInputStatusServiceTest.java b/...og2-server/src/test/java/org/graylog2/inputs/persistence/MongoInputStatusServiceTest.java
@@ -18,10 +18,10 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.eventbus.EventBus;
-import org.bson.types.ObjectId;
 import org.graylog.testing.mongodb.MongoDBFixtures;
 import org.graylog.testing.mongodb.MongoDBInstance;
 import org.graylog2.bindings.providers.MongoJackObjectMapperProvider;
+import org.graylog2.database.MongoCollections;
 import org.graylog2.database.NotFoundException;
 import org.graylog2.inputs.InputService;
 import org.graylog2.rest.models.system.inputs.responses.InputDeleted;
@@ -32,7 +32,6 @@
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
-import org.mongojack.JacksonDBCollection;
 
 import java.util.Optional;
 
@@ -58,19 +57,13 @@ public class MongoInputStatusServiceTest {
     @Mock
     private InputService inputService;
 
-    private JacksonDBCollection<InputStatusRecord, ObjectId> db;
-
     @Before
     public void setUp() {
         final ObjectMapper objectMapper = new ObjectMapperProvider().get();
         final MongoJackObjectMapperProvider mapperProvider = new MongoJackObjectMapperProvider(objectMapper);
 
-        cut = new MongoInputStatusService(mongodb.mongoConnection(), mapperProvider, inputService, mockEventBus);
-
-        db = JacksonDBCollection.wrap(mongodb.mongoConnection().getDatabase().getCollection(MongoInputStatusService.COLLECTION_NAME),
-                InputStatusRecord.class,
-                ObjectId.class,
-                mapperProvider.get());
+        cut = new MongoInputStatusService(
+                new MongoCollections(mapperProvider, mongodb.mongoConnection()), inputService, mockEventBus);
     }
 
     @Test