From b174eea8a9a63e9c9be88bb04332eaa684830b0b Mon Sep 17 00:00:00 2001 From: Dave Rigby Date: Mon, 13 May 2024 19:55:31 +0100 Subject: [PATCH] Support proxy_url and ssl_ca_certs options for gRPC (#341) ## 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 --- pinecone/grpc/base.py | 5 +++- pinecone/grpc/pinecone.py | 4 ++- tests/integration/proxy_config/conftest.py | 12 ++++++++ .../proxy_config/test_proxy_settings.py | 30 +++++++++++-------- 4 files changed, 37 insertions(+), 14 deletions(-) diff --git a/pinecone/grpc/base.py b/pinecone/grpc/base.py index 9c0131b4..00a98e6b 100644 --- a/pinecone/grpc/base.py +++ b/pinecone/grpc/base.py @@ -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( @@ -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) diff --git a/pinecone/grpc/pinecone.py b/pinecone/grpc/pinecone.py index a7bf4817..3f2ddb46 100644 --- a/pinecone/grpc/pinecone.py +++ b/pinecone/grpc/pinecone.py @@ -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) \ No newline at end of file diff --git a/tests/integration/proxy_config/conftest.py b/tests/integration/proxy_config/conftest.py index 2c59361f..9f775efc 100644 --- a/tests/integration/proxy_config/conftest.py +++ b/tests/integration/proxy_config/conftest.py @@ -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: @@ -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') diff --git a/tests/integration/proxy_config/test_proxy_settings.py b/tests/integration/proxy_config/test_proxy_settings.py index 906c1deb..f71346ea 100644 --- a/tests/integration/proxy_config/test_proxy_settings.py +++ b/tests/integration/proxy_config/test_proxy_settings.py @@ -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, @@ -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', @@ -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, @@ -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}'),