Repository: helidon-io/helidon (java)

Issue:
4.x - Mongo DbClientServiceContext futures not completed
## Environment Details
* Helidon Version: 4.0.11
* Helidon SE
----------

## Problem Description

When using MongoDbClient, both `statementFuture` and `resultFuture` (`io.helidon.dbclient.DbClientServiceContext`) are not completed.

## Steps to reproduce

```java
@Test
void testDbClientService() {
    CompletableFuture<Long> resultFuture = new CompletableFuture<>();
    CompletableFuture<Void> statementFuture = new CompletableFuture<>();
    DbClientService service = ctx -> {
        ctx.resultFuture().whenComplete((v, ex) -> {
            if (ex != null) {
                resultFuture.completeExceptionally(ex);
            } else {
                resultFuture.complete(v);
            }
        });
        ctx.statementFuture().whenComplete((v, ex) -> {
            if (ex != null) {
                statementFuture.completeExceptionally(ex);
            } else {
                statementFuture.complete(v);
            }
        });
        return ctx;
    };
    MongoDbClient dbClient = createClient(builder -> builder.addService(service));
    DbExecute exec = dbClient.execute();
    long ignored = exec.query("{\"operation\": \"command\", \"query\": { ping: 1 }}").count();
    assertThat(resultFuture.isDone(), is(true));
    assertThat(resultFuture.isCompletedExceptionally(), is(false));
    assertThat(statementFuture.isDone(), is(true));
    assertThat(statementFuture.isCompletedExceptionally(), is(false));
}
```

[//]: # "Step by step instructions to reproduce the problem"
[//]: # "Provide sample code/application if relevant"



Pull request:
Fix DbClientService for Mongo DbClient
### Description

Complete the DbClientServiceContext futures for DbClient Mongo support.
Add tests.

Fixes #9101

### Documentation

None.

Patch:
diff --git a/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementDml.java b/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementDml.java
index 0b1f52d49f..f22d3d7adb 100644
--- a/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementDml.java
+++ b/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 class MongoDbStatementDml extends MongoDbStatement<DbStatementDml> implem
                             "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(),
diff --git a/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementQuery.java b/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementQuery.java
index fe9ea5218a..f142e39fb7 100644
--- a/dbclient/mongodb/src/main/java/io/helidon/dbclient/mongodb/MongoDbStatementQuery.java
+++ b/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 @@ package io.helidon.dbclient.mongodb;
 
 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 io.helidon.dbclient.DbStatementType;
 
 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 class MongoDbStatementQuery extends MongoDbStatement<DbStatementQuery> im
             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 @@ public class MongoDbStatementQuery extends MongoDbStatement<DbStatementQuery> im
         }
     }
 
-    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 @@ public class MongoDbStatementQuery extends MongoDbStatement<DbStatementQuery> im
             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()));
     }


Test patch:
diff --git a/dbclient/mongodb/src/test/java/io/helidon/dbclient/mongodb/MongoDbClientTest.java b/dbclient/mongodb/src/test/java/io/helidon/dbclient/mongodb/MongoDbClientTest.java
index ca9e361ea4..4bb5daa1ea 100644
--- a/dbclient/mongodb/src/test/java/io/helidon/dbclient/mongodb/MongoDbClientTest.java
+++ b/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 @@ public class MongoDbClientTest {
         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);
+                }
+            });
+        }
+    }
 }