Skip to content

Commit

Permalink
Support proxy_url and ssl_ca_certs options for gRPC (#341)
Browse files Browse the repository at this point in the history
## Problem 

The PineconeGRPC client does not currently honor the proxy_url and
ssl_ca_certs arguments - they are accepted but silently ignored.

This means the gRPC client cannot be used in enviroments which require a
proxy to access the Pinecone service.

## Solution

Address by wiring up these two config options to the corresponding
options in the underlying grpc library - in addition to them still being
passed to the HTTP-based control-plane connections.

Usage is very similar to the non-gRPC client - for example:

    pc = PineconeGRPC(my_api_key,
                      proxy_url="http://localhost:8080",
                      ssl_ca_certs="/path/to/my/ca-cert.pem")

## Type of Change

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)

## Test Plan

Proxy tests expanded to also test PineconeGRPC
  • Loading branch information
daverigby authored May 13, 2024
1 parent 92b08b6 commit b174eea
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 14 deletions.
5 changes: 4 additions & 1 deletion pinecone/grpc/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,8 @@ def _gen_channel(self, options=None):
}
if self.grpc_client_config.secure:
default_options["grpc.ssl_target_name_override"] = target.split(":")[0]
if self.config.proxy_url:
default_options["grpc.http_proxy"] = self.config.proxy_url
user_provided_options = options or {}
_options = tuple((k, v) for k, v in {**default_options, **user_provided_options}.items())
_logger.debug(
Expand All @@ -109,7 +111,8 @@ def _gen_channel(self, options=None):
if not self.grpc_client_config.secure:
channel = grpc.insecure_channel(target, options=_options)
else:
root_cas = open(certifi.where(), "rb").read()
ca_certs = self.config.ssl_ca_certs if self.config.ssl_ca_certs else certifi.where()
root_cas = open(ca_certs, "rb").read()
tls = grpc.ssl_channel_credentials(root_certificates=root_cas)
channel = grpc.secure_channel(target, tls, options=_options)

Expand Down
4 changes: 3 additions & 1 deletion pinecone/grpc/pinecone.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,5 +125,7 @@ def Index(self, name: str = '', host: str = '', **kwargs):

config = ConfigBuilder.build(api_key=self.config.api_key,
host=index_host,
source_tag=self.config.source_tag)
source_tag=self.config.source_tag,
proxy_url=self.config.proxy_url,
ssl_ca_certs=self.config.ssl_ca_certs)
return GRPCIndex(index_name=name, config=config, **kwargs)
12 changes: 12 additions & 0 deletions tests/integration/proxy_config/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ def run_cmd(cmd, output):
if exit_code != 0:
raise Exception(f"Failed to run command: {cmd}")

def use_grpc():
return os.environ.get('USE_GRPC', 'false') == 'true'

@pytest.fixture(scope='session', autouse=True)
def start_docker():
with open("tests/integration/proxy_config/logs/proxyconfig-docker-start.log", "a") as output:
Expand All @@ -63,6 +66,15 @@ def proxy1():
def proxy2():
return PROXIES['proxy2']

@pytest.fixture()
def client_cls():
if use_grpc():
from pinecone.grpc import PineconeGRPC
return PineconeGRPC
else:
from pinecone import Pinecone
return Pinecone

@pytest.fixture()
def api_key():
return get_environment_var('PINECONE_API_KEY')
Expand Down
30 changes: 18 additions & 12 deletions tests/integration/proxy_config/test_proxy_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,26 +17,30 @@ def exercise_all_apis(client, index_name):
index.describe_index_stats()

class TestProxyConfig:
def test_https_proxy_with_self_signed_cert(self, api_key, index_name, proxy1):
@pytest.mark.skipif(os.getenv('USE_GRPC') != 'false',
reason="gRPC doesn't support 'https://' proxy URLs")
def test_https_proxy_with_self_signed_cert(self, client_cls, api_key, index_name, proxy1):
ssl_ca_certs = os.path.join(proxy1['ssl_ca_certs'], 'mitmproxy-ca-cert.pem')
pc = Pinecone(
pc = client_cls(
api_key=api_key,
proxy_url=PROXY1_URL_HTTPS,
ssl_ca_certs=ssl_ca_certs,
)
exercise_all_apis(pc, index_name)

def test_http_proxy_with_self_signed_cert(self, api_key, index_name, proxy1):
def test_http_proxy_with_self_signed_cert(self, client_cls, api_key, index_name, proxy1):
ssl_ca_certs = os.path.join(proxy1['ssl_ca_certs'], 'mitmproxy-ca-cert.pem')
pc = Pinecone(
pc = client_cls(
api_key=api_key,
proxy_url=PROXY1_URL_HTTP,
ssl_ca_certs=ssl_ca_certs,
)
exercise_all_apis(pc, index_name)

def test_proxy_with_ssl_verification_disabled_emits_warning(self, api_key):
pc = Pinecone(
@pytest.mark.skipif(os.getenv('USE_GRPC') != 'false',
reason="gRPC doesn't support disabling ssl_verify")
def test_proxy_with_ssl_verification_disabled_emits_warning(self, client_cls, api_key, index_name):
pc = client_cls(
api_key=api_key,
proxy_url=PROXY1_URL_HTTPS,
ssl_verify=False,
Expand All @@ -45,9 +49,9 @@ def test_proxy_with_ssl_verification_disabled_emits_warning(self, api_key):
with pytest.warns(InsecureRequestWarning):
pc.list_indexes()

def test_proxy_with_incorrect_cert_path(self, api_key):
def test_proxy_with_incorrect_cert_path(self, client_cls, api_key):
with pytest.raises(Exception) as e:
pc = Pinecone(
pc = client_cls(
api_key=api_key,
proxy_url=PROXY1_URL_HTTPS,
ssl_ca_certs='~/incorrect/path',
Expand All @@ -56,10 +60,10 @@ def test_proxy_with_incorrect_cert_path(self, api_key):

assert 'No such file or directory' in str(e.value)

def test_proxy_with_valid_path_to_incorrect_cert(self, api_key, proxy2):
def test_proxy_with_valid_path_to_incorrect_cert(self, client_cls, api_key, proxy2):
ssl_ca_certs = os.path.join(proxy2['ssl_ca_certs'], 'mitmproxy-ca-cert.pem')
with pytest.raises(Exception) as e:
pc = Pinecone(
pc = client_cls(
api_key=api_key,
proxy_url=PROXY1_URL_HTTPS,
ssl_ca_certs=ssl_ca_certs,
Expand All @@ -68,11 +72,13 @@ def test_proxy_with_valid_path_to_incorrect_cert(self, api_key, proxy2):

assert 'CERTIFICATE_VERIFY_FAILED' in str(e.value)

def test_proxy_that_requires_proxyauth(self, api_key, index_name, proxy2):
@pytest.mark.skipif(os.getenv('USE_GRPC') != 'false',
reason="gRPC doesn't support proxy auth")
def test_proxy_that_requires_proxyauth(self, client_cls, api_key, index_name, proxy2):
ssl_ca_certs = os.path.join(proxy2['ssl_ca_certs'], 'mitmproxy-ca-cert.pem')
username = proxy2['auth'][0]
password = proxy2['auth'][1]
pc = Pinecone(
pc = client_cls(
api_key=api_key,
proxy_url=PROXY2_URL,
proxy_headers=make_headers(proxy_basic_auth=f'{username}:{password}'),
Expand Down

0 comments on commit b174eea

Please sign in to comment.