Add proxy configuration for OkHTTPClient and NettyChannelBuilder

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

@@ -2,16 +2,16 @@
 
 import io.pinecone.configs.PineconeConfig;
 import io.pinecone.configs.PineconeConnection;
-import io.pinecone.exceptions.FailedRequestInfo;
-import io.pinecone.exceptions.HttpErrorMapper;
-import io.pinecone.exceptions.PineconeException;
-import io.pinecone.exceptions.PineconeValidationException;
+import io.pinecone.configs.ProxyConfig;
+import io.pinecone.exceptions.*;
 import okhttp3.OkHttpClient;
 import org.openapitools.client.ApiClient;
 import org.openapitools.client.ApiException;
 import org.openapitools.client.api.ManageIndexesApi;
 import org.openapitools.client.model.*;
 
+import java.net.InetSocketAddress;
+import java.net.Proxy;
 import java.util.Arrays;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -46,6 +46,10 @@ public class Pinecone {
         this.manageIndexesApi = manageIndexesApi;
     }
 
+    PineconeConfig getConfig() {
+        return config;
+    }
+
     /**
      * Creates a new serverless index with the specified parameters.
      * <p>
@@ -794,7 +798,6 @@ private void handleApiException(ApiException apiException) throws PineconeExcept
         HttpErrorMapper.mapHttpStatusError(failedRequestInfo, apiException);
     }
 
-
     /**
      * A builder class for creating a {@link Pinecone} instance. This builder allows for configuring a {@link Pinecone}
      * instance with custom parameters including an API key, a source tag, and a custom OkHttpClient.
@@ -805,7 +808,8 @@ public static class Builder {
 
         // Optional fields
         private String sourceTag;
-        private OkHttpClient okHttpClient = new OkHttpClient();
+        private ProxyConfig proxyConfig;
+        private OkHttpClient customOkHttpClient;
 
         /**
          * Constructs a new {@link Builder} with the mandatory API key.
@@ -867,7 +871,42 @@ public Builder withSourceTag(String sourceTag) {
          * @return This {@link Builder} instance for chaining method calls.
          */
         public Builder withOkHttpClient(OkHttpClient okHttpClient) {
-            this.okHttpClient = okHttpClient;
+            this.customOkHttpClient = okHttpClient;
+            return this;
+        }
+
+        /**
+         * Sets a proxy for the Pinecone client to use for control and data plane requests.
+         * <p>
+         * When a proxy is configured using this method, all control and data plane requests made by the Pinecone client
+         * will be routed through the specified proxy server.
+         * <p>
+         * It's important to note that both proxyHost and proxyPort parameters should be provided to establish
+         * the connection to the proxy server.
+         * <p>
+         * Example usage:
+         * <pre>{@code
+         *
+         * String proxyHost = System.getenv("PROXY_HOST");
+         * int proxyPort = Integer.parseInt(System.getenv("PROXY_PORT"));
+         * Pinecone pinecone = new Pinecone.Builder("PINECONE_API_KEY")
+         *     .withProxy(proxyHost, proxyPort)
+         *     .build();
+         *
+         * // Network requests for control plane operations will now be made using the specified proxy.
+         * pinecone.listIndexes();
+         *
+         * // Network requests for data plane operations will now be made using the specified proxy.
+         * Index index = pinecone.getIndexConnection("PINECONE_INDEX");
+         * index.describeIndexStats();
+         * }</pre>
+         *
+         * @param proxyHost The hostname or IP address of the proxy server. Must not be null.
+         * @param proxyPort The port number of the proxy server. Must not be null.
+         * @return This {@link Builder} instance for chaining method calls.
+         */
+        public Builder withProxy(String proxyHost, int proxyPort) {
+            this.proxyConfig = new ProxyConfig(proxyHost, proxyPort);
             return this;
         }
 
@@ -881,13 +920,16 @@ public Builder withOkHttpClient(OkHttpClient okHttpClient) {
          * @return A new {@link Pinecone} instance configured based on the builder parameters.
          */
         public Pinecone build() {
-            PineconeConfig clientConfig = new PineconeConfig(apiKey);
-            clientConfig.setSourceTag(sourceTag);
-            clientConfig.validate();
+            PineconeConfig config = new PineconeConfig(apiKey, sourceTag, proxyConfig);
+            config.validate();
 
-            ApiClient apiClient = new ApiClient(okHttpClient);
-            apiClient.setApiKey(clientConfig.getApiKey());
-            apiClient.setUserAgent(clientConfig.getUserAgent());
+            if (proxyConfig != null && customOkHttpClient != null) {
+                throw new PineconeConfigurationException("Invalid configuration: Both Custom OkHttpClient and Proxy are set. Please configure only one of these options.");
+            }
+
+            ApiClient apiClient = (customOkHttpClient != null) ? new ApiClient(customOkHttpClient) : new ApiClient(buildOkHttpClient());
+            apiClient.setApiKey(config.getApiKey());
+            apiClient.setUserAgent(config.getUserAgent());
 
             if (Boolean.parseBoolean(System.getenv("PINECONE_DEBUG"))) {
                 apiClient.setDebugging(true);
@@ -896,7 +938,16 @@ public Pinecone build() {
             ManageIndexesApi manageIndexesApi = new ManageIndexesApi();
             manageIndexesApi.setApiClient(apiClient);
 
-            return new Pinecone(clientConfig, manageIndexesApi);
+            return new Pinecone(config, manageIndexesApi);
+        }
+
+        private OkHttpClient buildOkHttpClient() {
+            OkHttpClient.Builder builder = new OkHttpClient.Builder();
+            if(proxyConfig != null) {
+                Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyConfig.getHost(), proxyConfig.getPort()));
+                builder.proxy(proxy);
+            }
+            return builder.build();
         }
     }
 }

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

@@ -6,7 +6,7 @@
 /**
  * The {@link PineconeConfig} class is responsible for managing the configuration settings
  * required to interact with the Pinecone API. It provides methods to set and retrieve
- * the necessary API key, host, source tag, and custom managed channel.
+ * the necessary API key, host, source tag, proxyConfig, and custom managed channel.
  * <pre>{@code
  *
  *     import io.grpc.ManagedChannel;
@@ -48,6 +48,7 @@ public class PineconeConfig {
     // Optional fields
     private String host;
     private String sourceTag;
+    private ProxyConfig proxyConfig;
     private ManagedChannel customManagedChannel;
 
     /**
@@ -66,8 +67,21 @@ public PineconeConfig(String apiKey) {
      * @param sourceTag An optional source tag to be included in the user agent.
      */
     public PineconeConfig(String apiKey, String sourceTag) {
+        this(apiKey, sourceTag, null);
+    }
+
+    /**
+     * Constructs a {@link PineconeConfig} instance with the specified API key, source tag, control plane proxy
+     * configuration, and data plane proxy configuration.
+     *
+     * @param apiKey The API key required to authenticate with the Pinecone API.
+     * @param sourceTag An optional source tag to be included in the user agent.
+     * @param proxyConfig The proxy configuration for control and data plane requests. Can be null if not set.
+     */
+    public PineconeConfig(String apiKey, String sourceTag, ProxyConfig proxyConfig) {
         this.apiKey = apiKey;
         this.sourceTag = sourceTag;
+        this.proxyConfig = proxyConfig;
     }
 
     /**
@@ -124,6 +138,24 @@ public void setSourceTag(String sourceTag) {
         this.sourceTag = normalizeSourceTag(sourceTag);
     }
 
+    /**
+     * Returns the proxy configuration for control and data plane requests.
+     *
+     * @return The proxy configuration for control and data plane requests, or null if not set.
+     */
+    public ProxyConfig getProxyConfig() {
+        return proxyConfig;
+    }
+
+    /**
+     * Sets the proxy configuration for control and data plane requests.
+     *
+     * @param proxyConfig The new proxy configuration for control and data plane requests.
+     */
+    public void setProxyConfig(ProxyConfig proxyConfig) {
+        this.proxyConfig = proxyConfig;
+    }
+
     /**
      * Returns the custom gRPC managed channel.
      *
@@ -148,13 +180,20 @@ public interface CustomChannelBuilder {
     }
 
     /**
-     * Validates the configuration, ensuring that the API key is not null or empty.
+     * Validates the configuration settings of the Pinecone client.
+     * This method ensures that the API key is not null or empty, and validates the proxy configurations if set.
+     * Throws a PineconeConfigurationException if the API key is null or empty, or if any of the proxy configurations are invalid.
      *
-     * @throws PineconeConfigurationException if the API key is null or empty.
+     * @throws PineconeConfigurationException If the API key is null or empty, or if any of the proxy configurations are invalid.
      */
     public void validate() {
         if (apiKey == null || apiKey.isEmpty())
             throw new PineconeConfigurationException("The API key is required and must not be empty or null");
+
+        // proxyConfig is set to null by default indicating the user is not interested in configuring the proxy
+        if(proxyConfig != null) {
+            proxyConfig.validate();
+        }
     }
 
     /**

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

@@ -1,7 +1,9 @@
 package io.pinecone.configs;
 
+import io.grpc.HttpConnectProxiedSocketAddress;
 import io.grpc.ManagedChannel;
 import io.grpc.Metadata;
+import io.grpc.ProxyDetector;
 import io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.NegotiationType;
 import io.grpc.netty.NettyChannelBuilder;
@@ -13,6 +15,8 @@
 import org.slf4j.LoggerFactory;
 
 import javax.net.ssl.SSLException;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -102,21 +106,6 @@ private VectorServiceGrpc.VectorServiceFutureStub generateAsyncStub(Metadata met
                 .withMaxOutboundMessageSize(DEFAULT_MAX_MESSAGE_SIZE);
     }
 
-    /**
-     * Close the connection and release all resources. A PineconeConnection's underlying gRPC components use resources
-     * like threads and TCP connections. To prevent leaking these resources the connection should be closed when it
-     * will no longer be used. If it may be used again leave it running.
-     */
-    @Override
-    public void close() {
-        try {
-            logger.debug("closing channel");
-            channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            logger.warn("Channel shutdown interrupted before termination confirmed");
-        }
-    }
-
     /**
      * Returns the gRPC channel.
      */
@@ -145,21 +134,38 @@ private void onConnectivityStateChanged() {
                 channel.getState(false), channel);
     }
 
-    public static ManagedChannel buildChannel(String host) {
+    private ManagedChannel buildChannel(String host) {
         String endpoint = formatEndpoint(host);
         NettyChannelBuilder builder = NettyChannelBuilder.forTarget(endpoint);
 
         try {
             builder = builder.overrideAuthority(endpoint)
                     .negotiationType(NegotiationType.TLS)
                     .sslContext(GrpcSslContexts.forClient().build());
+
+            if(config.getProxyConfig() != null) {
+                ProxyDetector proxyDetector = getProxyDetector();
+                builder.proxyDetector(proxyDetector);
+            }
         } catch (SSLException e) {
             throw new PineconeException("SSL error opening gRPC channel", e);
         }
 
         return builder.build();
     }
 
+    private ProxyDetector getProxyDetector() {
+        ProxyConfig proxyConfig = config.getProxyConfig();
+        return (targetServerAddress) -> {
+            SocketAddress proxyAddress = new InetSocketAddress(proxyConfig.getHost(), proxyConfig.getPort());
+
+            return HttpConnectProxiedSocketAddress.newBuilder()
+                    .setTargetAddress((InetSocketAddress) targetServerAddress)
+                    .setProxyAddress(proxyAddress)
+                    .build();
+        };
+    }
+
     private static Metadata assembleMetadata(PineconeConfig config) {
         Metadata metadata = new Metadata();
         metadata.put(Metadata.Key.of("api-key", Metadata.ASCII_STRING_MARSHALLER), config.getApiKey());
@@ -174,4 +180,19 @@ public static String formatEndpoint(String host) {
             throw new PineconeValidationException("Index host cannot be null or empty");
         }
     }
+
+    /**
+     * Close the connection and release all resources. A PineconeConnection's underlying gRPC components use resources
+     * like threads and TCP connections. To prevent leaking these resources the connection should be closed when it
+     * will no longer be used. If it may be used again leave it running.
+     */
+    @Override
+    public void close() {
+        try {
+            logger.debug("closing channel");
+            channel.shutdownNow().awaitTermination(5, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            logger.warn("Channel shutdown interrupted before termination confirmed");
+        }
+    }
 }