Skip to content

Commit fbf1c89

Browse files
injectivesinjectives
authored
Add support for bookmark manager extension functions (#1290)
* Add support for bookmark manager extension functions The extension option `updateListener` has been updated to `bookmarkConsumer`. * Ensure database is not null in explicit request
1 parent d63e21f commit fbf1c89

File tree

11 files changed

+357
-35
lines changed

11 files changed

+357
-35
lines changed

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,12 @@
2929
*/
3030
public final class BookmarkManagerConfig {
3131
private final Map<String, Set<Bookmark>> initialBookmarks;
32-
private final BiConsumer<String, Set<Bookmark>> updateListener;
32+
private final BiConsumer<String, Set<Bookmark>> bookmarkConsumer;
3333
private final BookmarkSupplier bookmarkSupplier;
3434

3535
private BookmarkManagerConfig(BookmarkManagerConfigBuilder builder) {
3636
this.initialBookmarks = builder.initialBookmarks;
37-
this.updateListener = builder.updateListener;
37+
this.bookmarkConsumer = builder.bookmarkConsumer;
3838
this.bookmarkSupplier = builder.bookmarkSupplier;
3939
}
4040

@@ -62,7 +62,7 @@ public Map<String, Set<Bookmark>> initialBookmarks() {
6262
* @return the update listener or {@code null}
6363
*/
6464
public BiConsumer<String, Set<Bookmark>> updateListener() {
65-
return updateListener;
65+
return bookmarkConsumer;
6666
}
6767

6868
/**
@@ -79,7 +79,7 @@ public BookmarkSupplier bookmarkSupplier() {
7979
*/
8080
public static class BookmarkManagerConfigBuilder {
8181
private Map<String, Set<Bookmark>> initialBookmarks = Collections.emptyMap();
82-
private BiConsumer<String, Set<Bookmark>> updateListener;
82+
private BiConsumer<String, Set<Bookmark>> bookmarkConsumer;
8383
private BookmarkSupplier bookmarkSupplier;
8484

8585
private BookmarkManagerConfigBuilder() {}
@@ -97,15 +97,15 @@ public BookmarkManagerConfigBuilder withInitialBookmarks(Map<String, Set<Bookmar
9797
}
9898

9999
/**
100-
* Provide a bookmarks update listener.
100+
* Provide a bookmark consumer.
101101
* <p>
102-
* The listener will be called outside bookmark manager's synchronisation lock.
102+
* The consumer will be called outside bookmark manager's synchronisation lock.
103103
*
104-
* @param updateListener update listener
104+
* @param bookmarkConsumer bookmark consumer
105105
* @return this builder
106106
*/
107-
public BookmarkManagerConfigBuilder withUpdateListener(BiConsumer<String, Set<Bookmark>> updateListener) {
108-
this.updateListener = updateListener;
107+
public BookmarkManagerConfigBuilder withBookmarkConsumer(BiConsumer<String, Set<Bookmark>> bookmarkConsumer) {
108+
this.bookmarkConsumer = bookmarkConsumer;
109109
return this;
110110
}
111111

@@ -119,7 +119,7 @@ public BookmarkManagerConfigBuilder withUpdateListener(BiConsumer<String, Set<Bo
119119
* @param bookmarkSupplier the bookmarks supplier
120120
* @return this builder
121121
*/
122-
public BookmarkManagerConfigBuilder withBookmarksSupplier(BookmarkSupplier bookmarkSupplier) {
122+
public BookmarkManagerConfigBuilder withBookmarkSupplier(BookmarkSupplier bookmarkSupplier) {
123123
this.bookmarkSupplier = bookmarkSupplier;
124124
return this;
125125
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ public interface BookmarkSupplier {
2929
/**
3030
* Supplies a set of bookmarks for a given database.
3131
*
32-
* @param database the database name
32+
* @param database the database name, must not be {@code null}
3333
* @return the set of bookmarks, must not be {@code null}
3434
*/
3535
Set<Bookmark> getBookmarks(String database);

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ void shouldReturnUpdateListener() {
6565

6666
// WHEN
6767
config = BookmarkManagerConfig.builder()
68-
.withUpdateListener(updateListener)
68+
.withBookmarkConsumer(updateListener)
6969
.build();
7070

7171
// WHEN & THEN
@@ -82,7 +82,7 @@ void shouldReturnBookmarkSupplier() {
8282

8383
// WHEN
8484
config = BookmarkManagerConfig.builder()
85-
.withBookmarksSupplier(bookmarkSupplier)
85+
.withBookmarkSupplier(bookmarkSupplier)
8686
.build();
8787

8888
// WHEN & THEN

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package neo4j.org.testkit.backend.messages.requests;
20+
21+
import lombok.Getter;
22+
import lombok.Setter;
23+
24+
@Setter
25+
@Getter
26+
public class BookmarksConsumerCompleted implements TestkitCallbackResult {
27+
private BookmarksConsumerCompletedBody data;
28+
29+
@Override
30+
public String getCallbackId() {
31+
return data.getRequestId();
32+
}
33+
34+
@Setter
35+
@Getter
36+
public static class BookmarksConsumerCompletedBody {
37+
private String requestId;
38+
}
39+
}

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package neo4j.org.testkit.backend.messages.requests;
20+
21+
import java.util.Set;
22+
import lombok.Getter;
23+
import lombok.Setter;
24+
25+
@Setter
26+
@Getter
27+
public class BookmarksSupplierCompleted implements TestkitCallbackResult {
28+
private BookmarksSupplierCompletedBody data;
29+
30+
@Override
31+
public String getCallbackId() {
32+
return data.getRequestId();
33+
}
34+
35+
@Setter
36+
@Getter
37+
public static class BookmarksSupplierCompletedBody {
38+
private String requestId;
39+
private Set<String> bookmarks;
40+
}
41+
}

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

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,10 +125,17 @@ public TestkitResponse process(TestkitState testkitState) {
125125
.collect(Collectors.toMap(
126126
Map.Entry::getKey,
127127
e -> e.getValue().stream().map(Bookmark::from).collect(Collectors.toSet())));
128-
var config = org.neo4j.driver.BookmarkManagerConfig.builder()
129-
.withInitialBookmarks(initialBookmarks)
130-
.build();
131-
var manager = BookmarkManagers.defaultManager(config);
128+
var managerConfigBuilder =
129+
org.neo4j.driver.BookmarkManagerConfig.builder().withInitialBookmarks(initialBookmarks);
130+
if (managerConfig.isBookmarksSupplierRegistered()) {
131+
managerConfigBuilder = managerConfigBuilder.withBookmarkSupplier(
132+
new TestkitBookmarkSupplier(testkitState, this::dispatchTestkitCallback));
133+
}
134+
if (managerConfig.isBookmarksConsumerRegistered()) {
135+
managerConfigBuilder = managerConfigBuilder.withBookmarkConsumer(
136+
new TestkitBookmarkConsumer(testkitState, this::dispatchTestkitCallback));
137+
}
138+
var manager = BookmarkManagers.defaultManager(managerConfigBuilder.build());
132139
configBuilder.withBookmarkManager(manager);
133140
});
134141
configBuilder.withDriverMetrics();
@@ -312,8 +319,8 @@ public static class NewDriverBody {
312319
@Getter
313320
public static class BookmarkManagerConfig {
314321
private Map<String, Set<String>> initialBookmarks;
315-
private boolean bookmarkSupplier;
316-
private boolean notifyBookmarks;
322+
private boolean bookmarksSupplierRegistered;
323+
private boolean bookmarksConsumerRegistered;
317324
}
318325

319326
@RequiredArgsConstructor

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package neo4j.org.testkit.backend.messages.requests;
20+
21+
import java.util.Set;
22+
import java.util.concurrent.CompletionStage;
23+
import java.util.function.BiConsumer;
24+
import java.util.function.BiFunction;
25+
import java.util.stream.Collectors;
26+
import lombok.RequiredArgsConstructor;
27+
import neo4j.org.testkit.backend.TestkitState;
28+
import neo4j.org.testkit.backend.messages.responses.BookmarksConsumerRequest;
29+
import neo4j.org.testkit.backend.messages.responses.TestkitCallback;
30+
import org.neo4j.driver.Bookmark;
31+
32+
@RequiredArgsConstructor
33+
class TestkitBookmarkConsumer implements BiConsumer<String, Set<Bookmark>> {
34+
private final TestkitState testkitState;
35+
private final BiFunction<TestkitState, TestkitCallback, CompletionStage<TestkitCallbackResult>> dispatchFunction;
36+
37+
@Override
38+
public void accept(String database, Set<Bookmark> bookmarks) {
39+
var callbackId = testkitState.newId();
40+
var body = BookmarksConsumerRequest.BookmarksConsumerRequestBody.builder()
41+
.id(callbackId)
42+
.database(database)
43+
.bookmarks(bookmarks.stream().map(Bookmark::value).collect(Collectors.toUnmodifiableSet()))
44+
.build();
45+
var callback = BookmarksConsumerRequest.builder().data(body).build();
46+
47+
var callbackStage = dispatchFunction.apply(testkitState, callback);
48+
try {
49+
callbackStage.toCompletableFuture().get();
50+
} catch (Exception e) {
51+
throw new RuntimeException("Unexpected failure during Testkit callback", e);
52+
}
53+
}
54+
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package neo4j.org.testkit.backend.messages.requests;
20+
21+
import java.util.Objects;
22+
import java.util.Set;
23+
import java.util.concurrent.CompletionStage;
24+
import java.util.function.BiFunction;
25+
import java.util.stream.Collectors;
26+
import lombok.RequiredArgsConstructor;
27+
import neo4j.org.testkit.backend.TestkitState;
28+
import neo4j.org.testkit.backend.messages.responses.BookmarksSupplierRequest;
29+
import neo4j.org.testkit.backend.messages.responses.TestkitCallback;
30+
import org.neo4j.driver.Bookmark;
31+
import org.neo4j.driver.BookmarkSupplier;
32+
33+
@RequiredArgsConstructor
34+
class TestkitBookmarkSupplier implements BookmarkSupplier {
35+
private final TestkitState testkitState;
36+
private final BiFunction<TestkitState, TestkitCallback, CompletionStage<TestkitCallbackResult>> dispatchFunction;
37+
38+
@Override
39+
public Set<Bookmark> getBookmarks(String database) {
40+
Objects.requireNonNull(database, "database must not be null");
41+
return getBookmarksFromTestkit(database);
42+
}
43+
44+
@Override
45+
public Set<Bookmark> getAllBookmarks() {
46+
return getBookmarksFromTestkit(null);
47+
}
48+
49+
private Set<Bookmark> getBookmarksFromTestkit(String database) {
50+
var callbackId = testkitState.newId();
51+
var bodyBuilder =
52+
BookmarksSupplierRequest.BookmarksSupplierRequestBody.builder().id(callbackId);
53+
if (database != null) {
54+
bodyBuilder = bodyBuilder.database(database);
55+
}
56+
var callback =
57+
BookmarksSupplierRequest.builder().data(bodyBuilder.build()).build();
58+
59+
var callbackStage = dispatchFunction.apply(testkitState, callback);
60+
BookmarksSupplierCompleted resolutionCompleted;
61+
try {
62+
resolutionCompleted = (BookmarksSupplierCompleted)
63+
callbackStage.toCompletableFuture().get();
64+
} catch (Exception e) {
65+
throw new RuntimeException("Unexpected failure during Testkit callback", e);
66+
}
67+
68+
return resolutionCompleted.getData().getBookmarks().stream()
69+
.map(Bookmark::from)
70+
.collect(Collectors.toUnmodifiableSet());
71+
}
72+
}

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

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,22 +27,39 @@
2727

2828
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "name")
2929
@JsonSubTypes({
30-
@JsonSubTypes.Type(NewDriver.class), @JsonSubTypes.Type(NewSession.class),
31-
@JsonSubTypes.Type(SessionRun.class), @JsonSubTypes.Type(ResultNext.class),
32-
@JsonSubTypes.Type(ResultConsume.class), @JsonSubTypes.Type(VerifyConnectivity.class),
33-
@JsonSubTypes.Type(SessionClose.class), @JsonSubTypes.Type(DriverClose.class),
34-
@JsonSubTypes.Type(RetryableNegative.class), @JsonSubTypes.Type(SessionReadTransaction.class),
35-
@JsonSubTypes.Type(TransactionRun.class), @JsonSubTypes.Type(RetryablePositive.class),
36-
@JsonSubTypes.Type(SessionBeginTransaction.class), @JsonSubTypes.Type(TransactionCommit.class),
37-
@JsonSubTypes.Type(SessionLastBookmarks.class), @JsonSubTypes.Type(SessionWriteTransaction.class),
38-
@JsonSubTypes.Type(ResolverResolutionCompleted.class), @JsonSubTypes.Type(CheckMultiDBSupport.class),
39-
@JsonSubTypes.Type(DomainNameResolutionCompleted.class), @JsonSubTypes.Type(StartTest.class),
40-
@JsonSubTypes.Type(TransactionRollback.class), @JsonSubTypes.Type(GetFeatures.class),
41-
@JsonSubTypes.Type(GetRoutingTable.class), @JsonSubTypes.Type(TransactionClose.class),
42-
@JsonSubTypes.Type(ResultList.class), @JsonSubTypes.Type(GetConnectionPoolMetrics.class),
43-
@JsonSubTypes.Type(ResultPeek.class), @JsonSubTypes.Type(CheckDriverIsEncrypted.class),
44-
@JsonSubTypes.Type(CypherTypeField.class), @JsonSubTypes.Type(ResultSingle.class),
45-
@JsonSubTypes.Type(StartSubTest.class)
30+
@JsonSubTypes.Type(NewDriver.class),
31+
@JsonSubTypes.Type(NewSession.class),
32+
@JsonSubTypes.Type(SessionRun.class),
33+
@JsonSubTypes.Type(ResultNext.class),
34+
@JsonSubTypes.Type(ResultConsume.class),
35+
@JsonSubTypes.Type(VerifyConnectivity.class),
36+
@JsonSubTypes.Type(SessionClose.class),
37+
@JsonSubTypes.Type(DriverClose.class),
38+
@JsonSubTypes.Type(RetryableNegative.class),
39+
@JsonSubTypes.Type(SessionReadTransaction.class),
40+
@JsonSubTypes.Type(TransactionRun.class),
41+
@JsonSubTypes.Type(RetryablePositive.class),
42+
@JsonSubTypes.Type(SessionBeginTransaction.class),
43+
@JsonSubTypes.Type(TransactionCommit.class),
44+
@JsonSubTypes.Type(SessionLastBookmarks.class),
45+
@JsonSubTypes.Type(SessionWriteTransaction.class),
46+
@JsonSubTypes.Type(ResolverResolutionCompleted.class),
47+
@JsonSubTypes.Type(CheckMultiDBSupport.class),
48+
@JsonSubTypes.Type(DomainNameResolutionCompleted.class),
49+
@JsonSubTypes.Type(StartTest.class),
50+
@JsonSubTypes.Type(TransactionRollback.class),
51+
@JsonSubTypes.Type(GetFeatures.class),
52+
@JsonSubTypes.Type(GetRoutingTable.class),
53+
@JsonSubTypes.Type(TransactionClose.class),
54+
@JsonSubTypes.Type(ResultList.class),
55+
@JsonSubTypes.Type(GetConnectionPoolMetrics.class),
56+
@JsonSubTypes.Type(ResultPeek.class),
57+
@JsonSubTypes.Type(CheckDriverIsEncrypted.class),
58+
@JsonSubTypes.Type(CypherTypeField.class),
59+
@JsonSubTypes.Type(ResultSingle.class),
60+
@JsonSubTypes.Type(StartSubTest.class),
61+
@JsonSubTypes.Type(BookmarksSupplierCompleted.class),
62+
@JsonSubTypes.Type(BookmarksConsumerCompleted.class)
4663
})
4764
public interface TestkitRequest {
4865
TestkitResponse process(TestkitState testkitState);

0 commit comments

Comments
 (0)