Add support for bookmark manager extension functions

driver/src/main/java/org/neo4j/driver/BookmarkManagerConfig.java

@@ -29,12 +29,12 @@
  */
 public final class BookmarkManagerConfig {
     private final Map<String, Set<Bookmark>> initialBookmarks;
-    private final BiConsumer<String, Set<Bookmark>> updateListener;
+    private final BiConsumer<String, Set<Bookmark>> bookmarkConsumer;
     private final BookmarkSupplier bookmarkSupplier;
 
     private BookmarkManagerConfig(BookmarkManagerConfigBuilder builder) {
         this.initialBookmarks = builder.initialBookmarks;
-        this.updateListener = builder.updateListener;
+        this.bookmarkConsumer = builder.bookmarkConsumer;
         this.bookmarkSupplier = builder.bookmarkSupplier;
     }
 
@@ -62,7 +62,7 @@ public Map<String, Set<Bookmark>> initialBookmarks() {
      * @return the update listener or {@code null}
      */
     public BiConsumer<String, Set<Bookmark>> updateListener() {
-        return updateListener;
+        return bookmarkConsumer;
     }
 
     /**
@@ -79,7 +79,7 @@ public BookmarkSupplier bookmarkSupplier() {
      */
     public static class BookmarkManagerConfigBuilder {
         private Map<String, Set<Bookmark>> initialBookmarks = Collections.emptyMap();
-        private BiConsumer<String, Set<Bookmark>> updateListener;
+        private BiConsumer<String, Set<Bookmark>> bookmarkConsumer;
         private BookmarkSupplier bookmarkSupplier;
 
         private BookmarkManagerConfigBuilder() {}
@@ -97,15 +97,15 @@ public BookmarkManagerConfigBuilder withInitialBookmarks(Map<String, Set<Bookmar
         }
 
         /**
-         * Provide a bookmarks update listener.
+         * Provide a bookmark consumer.
          * <p>
-         * The listener will be called outside bookmark manager's synchronisation lock.
+         * The consumer will be called outside bookmark manager's synchronisation lock.
          *
-         * @param updateListener update listener
+         * @param bookmarkConsumer bookmark consumer
          * @return this builder
          */
-        public BookmarkManagerConfigBuilder withUpdateListener(BiConsumer<String, Set<Bookmark>> updateListener) {
-            this.updateListener = updateListener;
+        public BookmarkManagerConfigBuilder withBookmarkConsumer(BiConsumer<String, Set<Bookmark>> bookmarkConsumer) {
+            this.bookmarkConsumer = bookmarkConsumer;
             return this;
         }
 
@@ -119,7 +119,7 @@ public BookmarkManagerConfigBuilder withUpdateListener(BiConsumer<String, Set<Bo
          * @param bookmarkSupplier the bookmarks supplier
          * @return this builder
          */
-        public BookmarkManagerConfigBuilder withBookmarksSupplier(BookmarkSupplier bookmarkSupplier) {
+        public BookmarkManagerConfigBuilder withBookmarkSupplier(BookmarkSupplier bookmarkSupplier) {
             this.bookmarkSupplier = bookmarkSupplier;
             return this;
         }

driver/src/main/java/org/neo4j/driver/BookmarkSupplier.java

@@ -29,7 +29,7 @@ public interface BookmarkSupplier {
     /**
      * Supplies a set of bookmarks for a given database.
      *
-     * @param database the database name
+     * @param database the database name, must not be {@code null}
      * @return the set of bookmarks, must not be {@code null}
      */
     Set<Bookmark> getBookmarks(String database);

driver/src/test/java/org/neo4j/driver/BookmarkManagerConfigTest.java

@@ -65,7 +65,7 @@ void shouldReturnUpdateListener() {
 
         // WHEN
         config = BookmarkManagerConfig.builder()
-                .withUpdateListener(updateListener)
+                .withBookmarkConsumer(updateListener)
                 .build();
 
         // WHEN & THEN
@@ -82,7 +82,7 @@ void shouldReturnBookmarkSupplier() {
 
         // WHEN
         config = BookmarkManagerConfig.builder()
-                .withBookmarksSupplier(bookmarkSupplier)
+                .withBookmarkSupplier(bookmarkSupplier)
                 .build();
 
         // WHEN & THEN

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/BookmarksConsumerCompleted.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package neo4j.org.testkit.backend.messages.requests;
+
+import lombok.Getter;
+import lombok.Setter;
+
+@Setter
+@Getter
+public class BookmarksConsumerCompleted implements TestkitCallbackResult {
+    private BookmarksConsumerCompletedBody data;
+
+    @Override
+    public String getCallbackId() {
+        return data.getRequestId();
+    }
+
+    @Setter
+    @Getter
+    public static class BookmarksConsumerCompletedBody {
+        private String requestId;
+    }
+}

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/BookmarksSupplierCompleted.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package neo4j.org.testkit.backend.messages.requests;
+
+import java.util.Set;
+import lombok.Getter;
+import lombok.Setter;
+
+@Setter
+@Getter
+public class BookmarksSupplierCompleted implements TestkitCallbackResult {
+    private BookmarksSupplierCompletedBody data;
+
+    @Override
+    public String getCallbackId() {
+        return data.getRequestId();
+    }
+
+    @Setter
+    @Getter
+    public static class BookmarksSupplierCompletedBody {
+        private String requestId;
+        private Set<String> bookmarks;
+    }
+}

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/NewDriver.java

@@ -125,10 +125,17 @@ public TestkitResponse process(TestkitState testkitState) {
                     .collect(Collectors.toMap(
                             Map.Entry::getKey,
                             e -> e.getValue().stream().map(Bookmark::from).collect(Collectors.toSet())));
-            var config = org.neo4j.driver.BookmarkManagerConfig.builder()
-                    .withInitialBookmarks(initialBookmarks)
-                    .build();
-            var manager = BookmarkManagers.defaultManager(config);
+            var managerConfigBuilder =
+                    org.neo4j.driver.BookmarkManagerConfig.builder().withInitialBookmarks(initialBookmarks);
+            if (managerConfig.isBookmarksSupplierRegistered()) {
+                managerConfigBuilder = managerConfigBuilder.withBookmarkSupplier(
+                        new TestkitBookmarkSupplier(testkitState, this::dispatchTestkitCallback));
+            }
+            if (managerConfig.isBookmarksConsumerRegistered()) {
+                managerConfigBuilder = managerConfigBuilder.withBookmarkConsumer(
+                        new TestkitBookmarkConsumer(testkitState, this::dispatchTestkitCallback));
+            }
+            var manager = BookmarkManagers.defaultManager(managerConfigBuilder.build());
             configBuilder.withBookmarkManager(manager);
         });
         configBuilder.withDriverMetrics();
@@ -312,8 +319,8 @@ public static class NewDriverBody {
     @Getter
     public static class BookmarkManagerConfig {
         private Map<String, Set<String>> initialBookmarks;
-        private boolean bookmarkSupplier;
-        private boolean notifyBookmarks;
+        private boolean bookmarksSupplierRegistered;
+        private boolean bookmarksConsumerRegistered;
     }
 
     @RequiredArgsConstructor

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TestkitBookmarkConsumer.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package neo4j.org.testkit.backend.messages.requests;
+
+import java.util.Set;
+import java.util.concurrent.CompletionStage;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import neo4j.org.testkit.backend.TestkitState;
+import neo4j.org.testkit.backend.messages.responses.BookmarksConsumerRequest;
+import neo4j.org.testkit.backend.messages.responses.TestkitCallback;
+import org.neo4j.driver.Bookmark;
+
+@RequiredArgsConstructor
+class TestkitBookmarkConsumer implements BiConsumer<String, Set<Bookmark>> {
+    private final TestkitState testkitState;
+    private final BiFunction<TestkitState, TestkitCallback, CompletionStage<TestkitCallbackResult>> dispatchFunction;
+
+    @Override
+    public void accept(String database, Set<Bookmark> bookmarks) {
+        var callbackId = testkitState.newId();
+        var body = BookmarksConsumerRequest.BookmarksConsumerRequestBody.builder()
+                .id(callbackId)
+                .database(database)
+                .bookmarks(bookmarks.stream().map(Bookmark::value).collect(Collectors.toUnmodifiableSet()))
+                .build();
+        var callback = BookmarksConsumerRequest.builder().data(body).build();
+
+        var callbackStage = dispatchFunction.apply(testkitState, callback);
+        try {
+            callbackStage.toCompletableFuture().get();
+        } catch (Exception e) {
+            throw new RuntimeException("Unexpected failure during Testkit callback", e);
+        }
+    }
+}

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TestkitBookmarkSupplier.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package neo4j.org.testkit.backend.messages.requests;
+
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.CompletionStage;
+import java.util.function.BiFunction;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import neo4j.org.testkit.backend.TestkitState;
+import neo4j.org.testkit.backend.messages.responses.BookmarksSupplierRequest;
+import neo4j.org.testkit.backend.messages.responses.TestkitCallback;
+import org.neo4j.driver.Bookmark;
+import org.neo4j.driver.BookmarkSupplier;
+
+@RequiredArgsConstructor
+class TestkitBookmarkSupplier implements BookmarkSupplier {
+    private final TestkitState testkitState;
+    private final BiFunction<TestkitState, TestkitCallback, CompletionStage<TestkitCallbackResult>> dispatchFunction;
+
+    @Override
+    public Set<Bookmark> getBookmarks(String database) {
+        Objects.requireNonNull(database, "database must not be null");
+        return getBookmarksFromTestkit(database);
+    }
+
+    @Override
+    public Set<Bookmark> getAllBookmarks() {
+        return getBookmarksFromTestkit(null);
+    }
+
+    private Set<Bookmark> getBookmarksFromTestkit(String database) {
+        var callbackId = testkitState.newId();
+        var bodyBuilder =
+                BookmarksSupplierRequest.BookmarksSupplierRequestBody.builder().id(callbackId);
+        if (database != null) {
+            bodyBuilder = bodyBuilder.database(database);
+        }
+        var callback =
+                BookmarksSupplierRequest.builder().data(bodyBuilder.build()).build();
+
+        var callbackStage = dispatchFunction.apply(testkitState, callback);
+        BookmarksSupplierCompleted resolutionCompleted;
+        try {
+            resolutionCompleted = (BookmarksSupplierCompleted)
+                    callbackStage.toCompletableFuture().get();
+        } catch (Exception e) {
+            throw new RuntimeException("Unexpected failure during Testkit callback", e);
+        }
+
+        return resolutionCompleted.getData().getBookmarks().stream()
+                .map(Bookmark::from)
+                .collect(Collectors.toUnmodifiableSet());
+    }
+}

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/TestkitRequest.java

@@ -27,22 +27,39 @@
 
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "name")
 @JsonSubTypes({
-    @JsonSubTypes.Type(NewDriver.class), @JsonSubTypes.Type(NewSession.class),
-    @JsonSubTypes.Type(SessionRun.class), @JsonSubTypes.Type(ResultNext.class),
-    @JsonSubTypes.Type(ResultConsume.class), @JsonSubTypes.Type(VerifyConnectivity.class),
-    @JsonSubTypes.Type(SessionClose.class), @JsonSubTypes.Type(DriverClose.class),
-    @JsonSubTypes.Type(RetryableNegative.class), @JsonSubTypes.Type(SessionReadTransaction.class),
-    @JsonSubTypes.Type(TransactionRun.class), @JsonSubTypes.Type(RetryablePositive.class),
-    @JsonSubTypes.Type(SessionBeginTransaction.class), @JsonSubTypes.Type(TransactionCommit.class),
-    @JsonSubTypes.Type(SessionLastBookmarks.class), @JsonSubTypes.Type(SessionWriteTransaction.class),
-    @JsonSubTypes.Type(ResolverResolutionCompleted.class), @JsonSubTypes.Type(CheckMultiDBSupport.class),
-    @JsonSubTypes.Type(DomainNameResolutionCompleted.class), @JsonSubTypes.Type(StartTest.class),
-    @JsonSubTypes.Type(TransactionRollback.class), @JsonSubTypes.Type(GetFeatures.class),
-    @JsonSubTypes.Type(GetRoutingTable.class), @JsonSubTypes.Type(TransactionClose.class),
-    @JsonSubTypes.Type(ResultList.class), @JsonSubTypes.Type(GetConnectionPoolMetrics.class),
-    @JsonSubTypes.Type(ResultPeek.class), @JsonSubTypes.Type(CheckDriverIsEncrypted.class),
-    @JsonSubTypes.Type(CypherTypeField.class), @JsonSubTypes.Type(ResultSingle.class),
-    @JsonSubTypes.Type(StartSubTest.class)
+    @JsonSubTypes.Type(NewDriver.class),
+    @JsonSubTypes.Type(NewSession.class),
+    @JsonSubTypes.Type(SessionRun.class),
+    @JsonSubTypes.Type(ResultNext.class),
+    @JsonSubTypes.Type(ResultConsume.class),
+    @JsonSubTypes.Type(VerifyConnectivity.class),
+    @JsonSubTypes.Type(SessionClose.class),
+    @JsonSubTypes.Type(DriverClose.class),
+    @JsonSubTypes.Type(RetryableNegative.class),
+    @JsonSubTypes.Type(SessionReadTransaction.class),
+    @JsonSubTypes.Type(TransactionRun.class),
+    @JsonSubTypes.Type(RetryablePositive.class),
+    @JsonSubTypes.Type(SessionBeginTransaction.class),
+    @JsonSubTypes.Type(TransactionCommit.class),
+    @JsonSubTypes.Type(SessionLastBookmarks.class),
+    @JsonSubTypes.Type(SessionWriteTransaction.class),
+    @JsonSubTypes.Type(ResolverResolutionCompleted.class),
+    @JsonSubTypes.Type(CheckMultiDBSupport.class),
+    @JsonSubTypes.Type(DomainNameResolutionCompleted.class),
+    @JsonSubTypes.Type(StartTest.class),
+    @JsonSubTypes.Type(TransactionRollback.class),
+    @JsonSubTypes.Type(GetFeatures.class),
+    @JsonSubTypes.Type(GetRoutingTable.class),
+    @JsonSubTypes.Type(TransactionClose.class),
+    @JsonSubTypes.Type(ResultList.class),
+    @JsonSubTypes.Type(GetConnectionPoolMetrics.class),
+    @JsonSubTypes.Type(ResultPeek.class),
+    @JsonSubTypes.Type(CheckDriverIsEncrypted.class),
+    @JsonSubTypes.Type(CypherTypeField.class),
+    @JsonSubTypes.Type(ResultSingle.class),
+    @JsonSubTypes.Type(StartSubTest.class),
+    @JsonSubTypes.Type(BookmarksSupplierCompleted.class),
+    @JsonSubTypes.Type(BookmarksConsumerCompleted.class)
 })
 public interface TestkitRequest {
     TestkitResponse process(TestkitState testkitState);