Skip to content

Commit c1bfc5d

Browse files
authored
Add proxy configuration for OkHTTPClient and NettyChannelBuilder (#136)
## Problem In order to configure a proxy: 1. For data plane operations, users have to first instantiate the PineconeConfig class followed by defining the custom NettyChannelBuilder and building the managedChannel to configure the proxy. And finally, after setting the customManagedChannel in PineconeConfig, PineConnection and Index/AsyncIndex classes can be instantiated. 2. For control plane operations, users have to define the OkHttpClient, configure the proxy, and set it to the Pinecone builder object. ## Solution Provide ability to configure proxies using `proxyHost` and `proxyPort` for both control plane (OkHttpClient) and data plane (gRPC calls via NettyChannelBuilder) operations without having the need to instantiate all of the classes as shown in the example below. 1. Data Plane operations via NettyChannelBuilder: Before: ```java import io.grpc.HttpConnectProxiedSocketAddress; import io.grpc.ManagedChannel; import io.grpc.ProxiedSocketAddress; import io.grpc.ProxyDetector; import io.pinecone.clients.Index; import io.pinecone.configs.PineconeConfig; import io.pinecone.configs.PineconeConnection; import io.grpc.netty.GrpcSslContexts; import io.grpc.netty.NegotiationType; import io.grpc.netty.NettyChannelBuilder; import io.pinecone.exceptions.PineconeException; import javax.net.ssl.SSLException; import java.net.InetSocketAddress; import java.net.SocketAddress; import java.util.concurrent.TimeUnit; import java.util.Arrays; ... String apiKey = System.getenv("PINECONE_API_KEY"); String proxyHost = System.getenv("PROXY_HOST"); int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT")); PineconeConfig config = new PineconeConfig(apiKey); String endpoint = System.getenv("PINECONE_HOST"); NettyChannelBuilder builder = NettyChannelBuilder.forTarget(endpoint); ProxyDetector proxyDetector = new ProxyDetector() { @OverRide public ProxiedSocketAddress proxyFor(SocketAddress targetServerAddress) { SocketAddress proxyAddress = new InetSocketAddress(proxyHost, proxyPort); return HttpConnectProxiedSocketAddress.newBuilder() .setTargetAddress((InetSocketAddress) targetServerAddress) .setProxyAddress(proxyAddress) .build(); } }; // Create custom channel try { builder = builder.overrideAuthority(endpoint) .negotiationType(NegotiationType.TLS) .keepAliveTimeout(5, TimeUnit.SECONDS) .sslContext(GrpcSslContexts.forClient().build()) .proxyDetector(proxyDetector); } catch (SSLException e) { throw new PineconeException("SSL error opening gRPC channel", e); } // Build the managed channel with the configured options ManagedChannel channel = builder.build(); config.setCustomManagedChannel(channel); PineconeConnection connection = new PineconeConnection(config); Index index = new Index(connection, "PINECONE_INDEX_NAME"); // Data plane operations // 1. Upsert data System.out.println(index.upsert("v1", Arrays.asList(1F, 2F, 3F, 4F))); // 2. Describe index stats System.out.println(index.describeIndexStats()); ``` After: ```java import io.pinecone.clients.Index; import io.pinecone.clients.Pinecone; ... String apiKey = System.getenv("PINECONE_API_KEY"); String proxyHost = System.getenv("PROXY_HOST"); int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT")); Pinecone pinecone = new Pinecone.Builder(apiKey) .withProxy(proxyHost, proxyPort) .build(); Index index = pinecone.getIndexConnection("PINECONE_INDEX_NAME"); // Data plane operation routed through the proxy server // 1. Upsert data System.out.println(index.upsert("v1", Arrays.asList(1F, 2F, 3F, 4F))); // 2. Describe index stats System.out.println(index.describeIndexStats()); ``` 2. Control Plane operations via OkHttpClient: Before: ```java import io.pinecone.clients.Pinecone; import okhttp3.OkHttpClient; import java.net.InetSocketAddress; import java.net.Proxy; ... String apiKey = System.getenv("PINECONE_API_KEY"); String proxyHost = System.getenv("PROXY_HOST"); int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT")); // Instantiate OkHttpClient instance and configure the proxy OkHttpClient client = new OkHttpClient.Builder() .proxy(new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyHost, proxyPort))) .build(); // Instantiate Pinecone class with the custom OkHttpClient object Pinecone pinecone = new Pinecone.Builder(apiKey) .withOkHttpClient(client) .build(); // Control plane operation routed through the proxy server System.out.println(pinecone.describeIndex("PINECONE_INDEX")); ``` After: ```java import io.pinecone.clients.Pinecone; ... String apiKey = System.getenv("PINECONE_API_KEY"); String proxyHost = System.getenv("PROXY_HOST"); int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT")); Pinecone pinecone = new Pinecone.Builder(apiKey) .withProxy(proxyHost, proxyPort) .build(); // Control plane operation routed through the proxy server System.out.println(pinecone.describeIndex("PINECONE_INDEX")); ``` Note: Users need to set up certificate authorities (CAs) to establish secure connections. Certificates verify server identities and encrypt data exchanged between the SDK and servers. By focusing on proxy host and port details, the SDK simplifies network setup while ensuring security. ## Type of Change - [X] New feature (non-breaking change which adds functionality) ## Test Plan Added unit tests. Given that we had issues with adding integration tests with mitm proxy in python SDK, I'm going to skip adding mitm proxy to github CI and have instead tested locally by spinning mitm proxy and successfully ran both control and data plane operations.
1 parent f7b850e commit c1bfc5d

File tree

6 files changed

+308
-34
lines changed

6 files changed

+308
-34
lines changed

src/main/java/io/pinecone/clients/Pinecone.java

+65-14
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22

33
import io.pinecone.configs.PineconeConfig;
44
import io.pinecone.configs.PineconeConnection;
5-
import io.pinecone.exceptions.FailedRequestInfo;
6-
import io.pinecone.exceptions.HttpErrorMapper;
7-
import io.pinecone.exceptions.PineconeException;
8-
import io.pinecone.exceptions.PineconeValidationException;
5+
import io.pinecone.configs.ProxyConfig;
6+
import io.pinecone.exceptions.*;
97
import okhttp3.OkHttpClient;
108
import org.openapitools.client.ApiClient;
119
import org.openapitools.client.ApiException;
1210
import org.openapitools.client.api.ManageIndexesApi;
1311
import org.openapitools.client.model.*;
1412

13+
import java.net.InetSocketAddress;
14+
import java.net.Proxy;
1515
import java.util.Arrays;
1616
import java.util.concurrent.ConcurrentHashMap;
1717

@@ -46,6 +46,10 @@ public class Pinecone {
4646
this.manageIndexesApi = manageIndexesApi;
4747
}
4848

49+
PineconeConfig getConfig() {
50+
return config;
51+
}
52+
4953
/**
5054
* Creates a new serverless index with the specified parameters.
5155
* <p>
@@ -794,7 +798,6 @@ private void handleApiException(ApiException apiException) throws PineconeExcept
794798
HttpErrorMapper.mapHttpStatusError(failedRequestInfo, apiException);
795799
}
796800

797-
798801
/**
799802
* A builder class for creating a {@link Pinecone} instance. This builder allows for configuring a {@link Pinecone}
800803
* instance with custom parameters including an API key, a source tag, and a custom OkHttpClient.
@@ -805,7 +808,8 @@ public static class Builder {
805808

806809
// Optional fields
807810
private String sourceTag;
808-
private OkHttpClient okHttpClient = new OkHttpClient();
811+
private ProxyConfig proxyConfig;
812+
private OkHttpClient customOkHttpClient;
809813

810814
/**
811815
* Constructs a new {@link Builder} with the mandatory API key.
@@ -867,7 +871,42 @@ public Builder withSourceTag(String sourceTag) {
867871
* @return This {@link Builder} instance for chaining method calls.
868872
*/
869873
public Builder withOkHttpClient(OkHttpClient okHttpClient) {
870-
this.okHttpClient = okHttpClient;
874+
this.customOkHttpClient = okHttpClient;
875+
return this;
876+
}
877+
878+
/**
879+
* Sets a proxy for the Pinecone client to use for control and data plane requests.
880+
* <p>
881+
* When a proxy is configured using this method, all control and data plane requests made by the Pinecone client
882+
* will be routed through the specified proxy server.
883+
* <p>
884+
* It's important to note that both proxyHost and proxyPort parameters should be provided to establish
885+
* the connection to the proxy server.
886+
* <p>
887+
* Example usage:
888+
* <pre>{@code
889+
*
890+
* String proxyHost = System.getenv("PROXY_HOST");
891+
* int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT"));
892+
* Pinecone pinecone = new Pinecone.Builder("PINECONE_API_KEY")
893+
* .withProxy(proxyHost, proxyPort)
894+
* .build();
895+
*
896+
* // Network requests for control plane operations will now be made using the specified proxy.
897+
* pinecone.listIndexes();
898+
*
899+
* // Network requests for data plane operations will now be made using the specified proxy.
900+
* Index index = pinecone.getIndexConnection("PINECONE_INDEX");
901+
* index.describeIndexStats();
902+
* }</pre>
903+
*
904+
* @param proxyHost The hostname or IP address of the proxy server. Must not be null.
905+
* @param proxyPort The port number of the proxy server. Must not be null.
906+
* @return This {@link Builder} instance for chaining method calls.
907+
*/
908+
public Builder withProxy(String proxyHost, int proxyPort) {
909+
this.proxyConfig = new ProxyConfig(proxyHost, proxyPort);
871910
return this;
872911
}
873912

@@ -881,13 +920,16 @@ public Builder withOkHttpClient(OkHttpClient okHttpClient) {
881920
* @return A new {@link Pinecone} instance configured based on the builder parameters.
882921
*/
883922
public Pinecone build() {
884-
PineconeConfig clientConfig = new PineconeConfig(apiKey);
885-
clientConfig.setSourceTag(sourceTag);
886-
clientConfig.validate();
923+
PineconeConfig config = new PineconeConfig(apiKey, sourceTag, proxyConfig);
924+
config.validate();
887925

888-
ApiClient apiClient = new ApiClient(okHttpClient);
889-
apiClient.setApiKey(clientConfig.getApiKey());
890-
apiClient.setUserAgent(clientConfig.getUserAgent());
926+
if (proxyConfig != null && customOkHttpClient != null) {
927+
throw new PineconeConfigurationException("Invalid configuration: Both Custom OkHttpClient and Proxy are set. Please configure only one of these options.");
928+
}
929+
930+
ApiClient apiClient = (customOkHttpClient != null) ? new ApiClient(customOkHttpClient) : new ApiClient(buildOkHttpClient());
931+
apiClient.setApiKey(config.getApiKey());
932+
apiClient.setUserAgent(config.getUserAgent());
891933

892934
if (Boolean.parseBoolean(System.getenv("PINECONE_DEBUG"))) {
893935
apiClient.setDebugging(true);
@@ -896,7 +938,16 @@ public Pinecone build() {
896938
ManageIndexesApi manageIndexesApi = new ManageIndexesApi();
897939
manageIndexesApi.setApiClient(apiClient);
898940

899-
return new Pinecone(clientConfig, manageIndexesApi);
941+
return new Pinecone(config, manageIndexesApi);
942+
}
943+
944+
private OkHttpClient buildOkHttpClient() {
945+
OkHttpClient.Builder builder = new OkHttpClient.Builder();
946+
if(proxyConfig != null) {
947+
Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyConfig.getHost(), proxyConfig.getPort()));
948+
builder.proxy(proxy);
949+
}
950+
return builder.build();
900951
}
901952
}
902953
}

src/main/java/io/pinecone/configs/PineconeConfig.java

+42-3
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
/**
77
* The {@link PineconeConfig} class is responsible for managing the configuration settings
88
* required to interact with the Pinecone API. It provides methods to set and retrieve
9-
* the necessary API key, host, source tag, and custom managed channel.
9+
* the necessary API key, host, source tag, proxyConfig, and custom managed channel.
1010
* <pre>{@code
1111
*
1212
* import io.grpc.ManagedChannel;
@@ -48,6 +48,7 @@ public class PineconeConfig {
4848
// Optional fields
4949
private String host;
5050
private String sourceTag;
51+
private ProxyConfig proxyConfig;
5152
private ManagedChannel customManagedChannel;
5253

5354
/**
@@ -66,8 +67,21 @@ public PineconeConfig(String apiKey) {
6667
* @param sourceTag An optional source tag to be included in the user agent.
6768
*/
6869
public PineconeConfig(String apiKey, String sourceTag) {
70+
this(apiKey, sourceTag, null);
71+
}
72+
73+
/**
74+
* Constructs a {@link PineconeConfig} instance with the specified API key, source tag, control plane proxy
75+
* configuration, and data plane proxy configuration.
76+
*
77+
* @param apiKey The API key required to authenticate with the Pinecone API.
78+
* @param sourceTag An optional source tag to be included in the user agent.
79+
* @param proxyConfig The proxy configuration for control and data plane requests. Can be null if not set.
80+
*/
81+
public PineconeConfig(String apiKey, String sourceTag, ProxyConfig proxyConfig) {
6982
this.apiKey = apiKey;
7083
this.sourceTag = sourceTag;
84+
this.proxyConfig = proxyConfig;
7185
}
7286

7387
/**
@@ -124,6 +138,24 @@ public void setSourceTag(String sourceTag) {
124138
this.sourceTag = normalizeSourceTag(sourceTag);
125139
}
126140

141+
/**
142+
* Returns the proxy configuration for control and data plane requests.
143+
*
144+
* @return The proxy configuration for control and data plane requests, or null if not set.
145+
*/
146+
public ProxyConfig getProxyConfig() {
147+
return proxyConfig;
148+
}
149+
150+
/**
151+
* Sets the proxy configuration for control and data plane requests.
152+
*
153+
* @param proxyConfig The new proxy configuration for control and data plane requests.
154+
*/
155+
public void setProxyConfig(ProxyConfig proxyConfig) {
156+
this.proxyConfig = proxyConfig;
157+
}
158+
127159
/**
128160
* Returns the custom gRPC managed channel.
129161
*
@@ -148,13 +180,20 @@ public interface CustomChannelBuilder {
148180
}
149181

150182
/**
151-
* Validates the configuration, ensuring that the API key is not null or empty.
183+
* Validates the configuration settings of the Pinecone client.
184+
* This method ensures that the API key is not null or empty, and validates the proxy configurations if set.
185+
* Throws a PineconeConfigurationException if the API key is null or empty, or if any of the proxy configurations are invalid.
152186
*
153-
* @throws PineconeConfigurationException if the API key is null or empty.
187+
* @throws PineconeConfigurationException If the API key is null or empty, or if any of the proxy configurations are invalid.
154188
*/
155189
public void validate() {
156190
if (apiKey == null || apiKey.isEmpty())
157191
throw new PineconeConfigurationException("The API key is required and must not be empty or null");
192+
193+
// proxyConfig is set to null by default indicating the user is not interested in configuring the proxy
194+
if(proxyConfig != null) {
195+
proxyConfig.validate();
196+
}
158197
}
159198

160199
/**

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

+37-16
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package io.pinecone.configs;
22

3+
import io.grpc.HttpConnectProxiedSocketAddress;
34
import io.grpc.ManagedChannel;
45
import io.grpc.Metadata;
6+
import io.grpc.ProxyDetector;
57
import io.grpc.netty.GrpcSslContexts;
68
import io.grpc.netty.NegotiationType;
79
import io.grpc.netty.NettyChannelBuilder;
@@ -13,6 +15,8 @@
1315
import org.slf4j.LoggerFactory;
1416

1517
import javax.net.ssl.SSLException;
18+
import java.net.InetSocketAddress;
19+
import java.net.SocketAddress;
1620
import java.util.concurrent.TimeUnit;
1721

1822
/**
@@ -102,21 +106,6 @@ private VectorServiceGrpc.VectorServiceFutureStub generateAsyncStub(Metadata met
102106
.withMaxOutboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE);
103107
}
104108

105-
/**
106-
* Close the connection and release all resources. A PineconeConnection's underlying gRPC components use resources
107-
* like threads and TCP connections. To prevent leaking these resources the connection should be closed when it
108-
* will no longer be used. If it may be used again leave it running.
109-
*/
110-
@Override
111-
public void close() {
112-
try {
113-
logger.debug("closing channel");
114-
channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
115-
} catch (InterruptedException e) {
116-
logger.warn("Channel shutdown interrupted before termination confirmed");
117-
}
118-
}
119-
120109
/**
121110
* Returns the gRPC channel.
122111
*/
@@ -145,21 +134,38 @@ private void onConnectivityStateChanged() {
145134
channel.getState(false), channel);
146135
}
147136

148-
public static ManagedChannel buildChannel(String host) {
137+
private ManagedChannel buildChannel(String host) {
149138
String endpoint = formatEndpoint(host);
150139
NettyChannelBuilder builder = NettyChannelBuilder.forTarget(endpoint);
151140

152141
try {
153142
builder = builder.overrideAuthority(endpoint)
154143
.negotiationType(NegotiationType.TLS)
155144
.sslContext(GrpcSslContexts.forClient().build());
145+
146+
if(config.getProxyConfig() != null) {
147+
ProxyDetector proxyDetector = getProxyDetector();
148+
builder.proxyDetector(proxyDetector);
149+
}
156150
} catch (SSLException e) {
157151
throw new PineconeException("SSL error opening gRPC channel", e);
158152
}
159153

160154
return builder.build();
161155
}
162156

157+
private ProxyDetector getProxyDetector() {
158+
ProxyConfig proxyConfig = config.getProxyConfig();
159+
return (targetServerAddress) -> {
160+
SocketAddress proxyAddress = new InetSocketAddress(proxyConfig.getHost(), proxyConfig.getPort());
161+
162+
return HttpConnectProxiedSocketAddress.newBuilder()
163+
.setTargetAddress((InetSocketAddress) targetServerAddress)
164+
.setProxyAddress(proxyAddress)
165+
.build();
166+
};
167+
}
168+
163169
private static Metadata assembleMetadata(PineconeConfig config) {
164170
Metadata metadata = new Metadata();
165171
metadata.put(Metadata.Key.of("api-key", Metadata.ASCII_STRING_MARSHALLER), config.getApiKey());
@@ -174,4 +180,19 @@ public static String formatEndpoint(String host) {
174180
throw new PineconeValidationException("Index host cannot be null or empty");
175181
}
176182
}
183+
184+
/**
185+
* Close the connection and release all resources. A PineconeConnection's underlying gRPC components use resources
186+
* like threads and TCP connections. To prevent leaking these resources the connection should be closed when it
187+
* will no longer be used. If it may be used again leave it running.
188+
*/
189+
@Override
190+
public void close() {
191+
try {
192+
logger.debug("closing channel");
193+
channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
194+
} catch (InterruptedException e) {
195+
logger.warn("Channel shutdown interrupted before termination confirmed");
196+
}
197+
}
177198
}

0 commit comments

Comments
 (0)