Fix DbClientService for Mongo DbClient (helidon-io#9102)

Fixes helidon-io#9101

Author: romain-grecourt

dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementDml.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -63,6 +63,7 @@ public long execute() {
                             "Statement operation not yet supported: %s",
                             type.name()));
                 };
+                future.complete(result);
                 LOGGER.log(System.Logger.Level.DEBUG, () -> String.format(
                         "%s DML %s execution succeeded",
                         type.name(),

dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementQuery.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2019, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 
 import java.lang.System.Logger.Level;
 import java.util.Spliterator;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -28,6 +29,7 @@
 
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCollection;
+import com.mongodb.client.MongoCursor;
 import com.mongodb.client.MongoDatabase;
 import org.bson.Document;
 
@@ -62,8 +64,8 @@ public Stream<DbRow> execute() {
             try {
                 MongoStatement stmt = queryOrCommand(preparedStmt);
                 return switch (stmt.getOperation()) {
-                    case QUERY -> executeQuery(stmt);
-                    case COMMAND -> executeCommand(stmt);
+                    case QUERY -> executeQuery(stmt, future);
+                    case COMMAND -> executeCommand(stmt, future);
                     default -> throw new UnsupportedOperationException(String.format(
                             "Operation %s is not supported by query", stmt.getOperation().toString()));
                 };
@@ -87,14 +89,15 @@ private MongoStatement queryOrCommand(String statement) {
         }
     }
 
-    private Stream<DbRow> executeCommand(MongoStatement stmt) {
+    private Stream<DbRow> executeCommand(MongoStatement stmt, CompletableFuture<Long> future) {
         Document command = stmt.getQuery();
         LOGGER.log(Level.DEBUG, () -> String.format("Command: %s", command.toString()));
         Document doc = db().runCommand(command);
+        future.complete(1L);
         return Stream.of(new MongoDbRow(doc, context()));
     }
 
-    private Stream<DbRow> executeQuery(MongoStatement stmt) {
+    private Stream<DbRow> executeQuery(MongoStatement stmt, CompletableFuture<Long> future) {
         MongoCollection<Document> mc = db().getCollection(stmt.getCollection());
         Document query = stmt.getQuery();
         Document projection = stmt.getProjection();
@@ -109,7 +112,9 @@ private Stream<DbRow> executeQuery(MongoStatement stmt) {
             finder = finder.projection(projection);
         }
 
-        Spliterator<Document> spliterator = spliteratorUnknownSize(finder.iterator(), Spliterator.ORDERED);
+        MongoCursor<Document> it = finder.iterator();
+        future.complete(it.hasNext() ? 1L : 0L);
+        Spliterator<Document> spliterator = spliteratorUnknownSize(it, Spliterator.ORDERED);
         Stream<Document> stream = StreamSupport.stream(spliterator, false);
         return stream.map(doc -> new MongoDbRow(doc, context()));
     }

dbclient/mongodb/src/test/java/io/helidon/dbclient/mongodb/MongoDbClientTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2021, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,34 +16,39 @@
 package io.helidon.dbclient.mongodb;
 
 import java.lang.System.Logger.Level;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.function.Consumer;
 
+import io.helidon.dbclient.DbClientService;
+import io.helidon.dbclient.DbClientServiceContext;
 import io.helidon.dbclient.DbExecute;
 
 import com.mongodb.client.MongoClient;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
 
+import org.bson.Document;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.is;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 
-public class MongoDbClientTest {
+@SuppressWarnings("resource")
+class MongoDbClientTest {
 
     private static final System.Logger LOGGER = System.getLogger(MongoDbClientTest.class.getName());
     private static MongoDbClient dbClient;
 
     @BeforeAll
     static void setup() {
-        MongoClient client = Mockito.mock(MongoClient.class);
-        MongoDatabase db = Mockito.mock(MongoDatabase.class);
-        when(db.runCommand(any())).thenReturn(MongoDbStatement.EMPTY);
-        dbClient = new MongoDbClient(new MongoDbClientBuilder(), client, db);
+        dbClient = createClient(null);
     }
 
     @Test
@@ -85,4 +90,70 @@ void testUnsupportedUnwrapExecutorClass() {
         exec.query("{\"operation\": \"command\", \"query\": { ping: 1 }}");
     }
 
+    @Test
+    void testDbClientServiceQuery() {
+        TestDbClientService service = new TestDbClientService();
+        MongoDbClient dbClient = createClient(builder -> builder.addService(service));
+        DbExecute exec = dbClient.execute();
+        long ignored = exec.query("{\"operation\": \"command\", \"query\": { ping: 1 }}").count();
+        assertThat(service.resultFuture.isDone(), is(true));
+        assertThat(service.resultFuture.isCompletedExceptionally(), is(false));
+        assertThat(service.statementFuture.isDone(), is(true));
+        assertThat(service.statementFuture.isCompletedExceptionally(), is(false));
+    }
+
+    @Test
+    void testDbClientServiceDml() {
+        TestDbClientService service = new TestDbClientService();
+        MongoDbClient dbClient = createClient(builder -> builder.addService(service));
+        DbExecute exec = dbClient.execute();
+        long ignored = exec.insert("{"
+                                + "\"collection\": \"foo\","
+                                + "\"operation\": \"insert\","
+                                + "\"value\": { \"name\": \"bar\" }"
+                                + "}");
+        assertThat(service.resultFuture.isDone(), is(true));
+        assertThat(service.resultFuture.isCompletedExceptionally(), is(false));
+        assertThat(service.statementFuture.isDone(), is(true));
+        assertThat(service.statementFuture.isCompletedExceptionally(), is(false));
+    }
+
+    @SuppressWarnings("unchecked")
+    static MongoDbClient createClient(Consumer<MongoDbClientBuilder> consumer) {
+        MongoClient client = Mockito.mock(MongoClient.class);
+        MongoDatabase db = Mockito.mock(MongoDatabase.class);
+        MongoCollection<Document> collection = Mockito.mock(MongoCollection.class);
+        when(db.getCollection(any())).thenReturn(collection);
+        when(db.runCommand(any())).thenReturn(MongoDbStatement.EMPTY);
+        MongoDbClientBuilder builder = new MongoDbClientBuilder();
+        if (consumer != null) {
+            consumer.accept(builder);
+        }
+        return new MongoDbClient(builder, client, db);
+    }
+
+    record TestDbClientService(CompletableFuture<Long> resultFuture,
+                               CompletableFuture<Void> statementFuture) implements DbClientService {
+
+        TestDbClientService() {
+            this(new CompletableFuture<>(), new CompletableFuture<>());
+        }
+
+        @Override
+        public DbClientServiceContext statement(DbClientServiceContext context) {
+            setup(context.resultFuture(), resultFuture);
+            setup(context.statementFuture(), statementFuture);
+            return context;
+        }
+
+        static <T> void setup(CompletionStage<T> stage, CompletableFuture<T> future) {
+            stage.whenComplete((v, ex) -> {
+                if (ex != null) {
+                    future.completeExceptionally(ex);
+                } else {
+                    future.complete(v);
+                }
+            });
+        }
+    }
 }