Skip to content

Commit 50f812d

Browse files
authored
Expose future stub on Pinecone connection (#32)
This commit exposes the future stub on Pinecone connection, allowing for non-blocking async calls ## Problem PineconeConnection does not currently expose a stub for making asynchronous calls to the service. Async calls are very common in high traffic services. ## Solution Exposed the generated `VectorServiceGrpc.VectorServiceFutureStub` on `PineconeConnection`. This returns `com.google.common.util.concurrent.ListenableFuture<V>` for each call, which allows chaining of asynchronous operations. The PR also removes an unused public method, `setBlockingStub`; can add this back if client API compatibility dictates. ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Run `sanityAsyncFuture` test
1 parent be78899 commit 50f812d

File tree

1 file changed

+13
-16
lines changed

1 file changed

+13
-16
lines changed

src/main/java/io/pinecone/PineconeConnection.java

+13-16
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public class PineconeConnection implements AutoCloseable {
3434
*/
3535
private VectorServiceGrpc.VectorServiceBlockingStub blockingStub;
3636

37-
private VectorServiceGrpc.VectorServiceStub asyncStub;
37+
private VectorServiceGrpc.VectorServiceFutureStub futureStub;
3838

3939
public PineconeConnection(PineconeClientConfig clientConfig, PineconeConnectionConfig connectionConfig) {
4040
this.connectionConfig = connectionConfig;
@@ -45,13 +45,17 @@ public PineconeConnection(PineconeClientConfig clientConfig, PineconeConnectionC
4545
? connectionConfig.getCustomChannelBuilder().apply(clientConfig, connectionConfig)
4646
: buildChannel(clientConfig, connectionConfig);
4747
channel.notifyWhenStateChanged(channel.getState(false), this::onConnectivityStateChanged);
48-
Metadata metadata = assembleMetadata(clientConfig, connectionConfig);
48+
Metadata metadata = assembleMetadata(clientConfig);
4949
blockingStub = VectorServiceGrpc
5050
.newBlockingStub(channel)
51-
.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(metadata));
52-
asyncStub = VectorServiceGrpc
53-
.newStub(channel)
54-
.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(metadata));
51+
.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(metadata))
52+
.withMaxInboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE)
53+
.withMaxOutboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE);
54+
futureStub = VectorServiceGrpc
55+
.newFutureStub(channel)
56+
.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(metadata))
57+
.withMaxInboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE)
58+
.withMaxOutboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE);
5559
logger.debug("created new PineconeConnection for channel: {}", channel);
5660
}
5761

@@ -79,8 +83,8 @@ public VectorServiceGrpc.VectorServiceBlockingStub getBlockingStub() {
7983
return blockingStub;
8084
}
8185

82-
public void setBlockingStub(VectorServiceGrpc.VectorServiceBlockingStub blockingStub) {
83-
this.blockingStub = blockingStub;
86+
public VectorServiceGrpc.VectorServiceFutureStub getFutureStub() {
87+
return futureStub;
8488
}
8589

8690
private void onConnectivityStateChanged() {
@@ -104,20 +108,13 @@ public static ManagedChannel buildChannel(PineconeClientConfig clientConfig,
104108
return builder.build();
105109
}
106110

107-
private static Metadata assembleMetadata(PineconeClientConfig clientConfig,
108-
PineconeConnectionConfig connectionConfig) {
111+
private static Metadata assembleMetadata(PineconeClientConfig clientConfig) {
109112
Metadata metadata = new Metadata();
110113
metadata.put(Metadata.Key.of("api-key",
111114
Metadata.ASCII_STRING_MARSHALLER), clientConfig.getApiKey());
112115
return metadata;
113116
}
114117

115-
private VectorServiceGrpc.VectorServiceBlockingStub applyDefaultBlockingStubConfig(VectorServiceGrpc.VectorServiceBlockingStub stub) {
116-
return stub
117-
.withMaxInboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE)
118-
.withMaxOutboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE);
119-
}
120-
121118
static String getEndpoint(PineconeClientConfig clientConfig, PineconeConnectionConfig connectionConfig) {
122119
String endpoint = (connectionConfig.getConnectionUrl() != null) ?
123120
connectionConfig.getConnectionUrl().replaceFirst("https?://", "") :

0 commit comments

Comments
 (0)