Skip to content

Commit 5eb002e

Browse files
authored
Expose missing configurations for grpc_config and pool_threads (#296)
## Problem This PR addresses a few configuration-related issues: - We moved `pool_threads` configuration up to the `Pinecone` client instance, but for backwards compatibility it would be nice to also allow it to be passed in the call to `pc.Index(host='your-host', pool_threads=5)` - We need to expose `grpc_config` as a keyword param when targeting indexes via the grpc-flavored client. ## Solution - Pass through kwargs from `.Index()` to underlying Index class implementation - Add unit tests making sure `pool_threads` and `grpc_config` handled correctly. ### Before (pool_threads) ```python from pinecone import Pinecone pc = Pinecone(api_key='key', pool_threads=5) pc.Index(host='host') pc.upsert(...) ``` ### After (pool_threads, add backwards compatibility with old way) You can use syntax as before, or alternatively can do this: ```python from pinecone import Pinecone pc = Pinecone(api_key='key') index = pc.Index(host='host', pool_threads=5) index.upsert(...) ``` ### After (grpc, allow grpc_config to be passed) ```python from pinecone.grpc import PineconeGRPC, GRPCClientConfig pc = PineconeGRPC(api_key='YOUR_API_KEY') grpc_config = GRPCClientConfig( timeout=10, secure=True, reuse_channel=True ) index = pc.Index(name='my-index', host='host', grpc_config=grpc_config) ``` ## Type of Change - [x] New feature (non-breaking change which adds functionality) - [x] None of the above: Add backwards compatibility on pool_threads config ## Test Plan New tests should be green in CI.
1 parent 2a7faa8 commit 5eb002e

File tree

6 files changed

+102
-9
lines changed

6 files changed

+102
-9
lines changed

pinecone/control/pinecone.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -94,10 +94,10 @@ def __init__(
9494
else:
9595
self.config = PineconeConfig.build(api_key=api_key, host=host, additional_headers=additional_headers, **kwargs)
9696

97+
self.pool_threads = pool_threads
9798
if index_api:
9899
self.index_api = index_api
99100
else:
100-
self.pool_threads = pool_threads
101101
api_client = ApiClient(configuration=self.config.openapi_config, pool_threads=self.pool_threads)
102102
api_client.user_agent = get_user_agent()
103103
extra_headers = self.config.additional_headers or {}
@@ -446,7 +446,7 @@ def _get_status(self, name: str):
446446
response = api_instance.describe_index(name)
447447
return response["status"]
448448

449-
def Index(self, name: str = '', host: str = ''):
449+
def Index(self, name: str = '', host: str = '', **kwargs):
450450
"""
451451
Target an index for data operations.
452452
@@ -518,12 +518,14 @@ def Index(self, name: str = '', host: str = ''):
518518
"""
519519
if name == '' and host == '':
520520
raise ValueError("Either name or host must be specified")
521+
522+
pt = kwargs.pop('pool_threads', None) or self.pool_threads
521523

522524
if host != '':
523525
# Use host url if it is provided
524-
return Index(api_key=self.config.api_key, host=normalize_host(host), pool_threads=self.pool_threads)
526+
return Index(api_key=self.config.api_key, host=normalize_host(host), pool_threads=pt, **kwargs)
525527

526528
if name != '':
527529
# Otherwise, get host url from describe_index using the index name
528530
index_host = self.index_host_store.get_host(self.index_api, self.config, name)
529-
return Index(api_key=self.config.api_key, host=index_host, pool_threads=self.pool_threads)
531+
return Index(api_key=self.config.api_key, host=index_host, pool_threads=pt, **kwargs)

pinecone/grpc/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646

4747
from .index_grpc import GRPCIndex
4848
from .pinecone import PineconeGRPC
49+
from .config import GRPCClientConfig
4950

5051
from pinecone.core.grpc.protos.vector_service_pb2 import (
5152
Vector as GRPCVector,

pinecone/grpc/config.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ class GRPCClientConfig(NamedTuple):
77
GRPC client configuration options.
88
99
:param secure: Whether to use encrypted protocol (SSL). defaults to True.
10-
:type traceroute: bool, optional
10+
:type secure: bool, optional
1111
:param timeout: defaults to 2 seconds. Fail if gateway doesn't receive response within timeout.
1212
:type timeout: int, optional
1313
:param conn_timeout: defaults to 1. Timeout to retry connection if gRPC is unavailable. 0 is no retry.

pinecone/grpc/pinecone.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ class PineconeGRPC(Pinecone):
4646
4747
"""
4848

49-
def Index(self, name: str = '', host: str = ''):
49+
def Index(self, name: str = '', host: str = '', **kwargs):
5050
"""
5151
Target an index for data operations.
5252
@@ -123,10 +123,10 @@ def Index(self, name: str = '', host: str = ''):
123123
if host != '':
124124
# Use host if it is provided
125125
config = ConfigBuilder.build(api_key=self.config.api_key, host=host)
126-
return GRPCIndex(index_name=name, config=config)
126+
return GRPCIndex(index_name=name, config=config, **kwargs)
127127

128128
if name != '':
129129
# Otherwise, get host url from describe_index using the index name
130130
index_host = self.index_host_store.get_host(self.index_api, self.config, name)
131131
config = ConfigBuilder.build(api_key=self.config.api_key, host=index_host)
132-
return GRPCIndex(index_name=name, config=config)
132+
return GRPCIndex(index_name=name, config=config, **kwargs)

tests/unit/test_control.py

+25-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import pytest
22
from pinecone import Pinecone, PodSpec, ServerlessSpec
33
from pinecone.core.client.models import IndexList, IndexModel
4+
from pinecone.core.client.api.manage_indexes_api import ManageIndexesApi
45
import time
56

67
@pytest.fixture
@@ -67,4 +68,27 @@ def test_list_indexes_returns_iterable(self, mocker, index_list_response):
6768
mocker.patch.object(p.index_api, 'list_indexes', side_effect=[index_list_response])
6869

6970
response = p.list_indexes()
70-
assert [i.name for i in response] == ["index1", "index2", "index3"]
71+
assert [i.name for i in response] == ["index1", "index2", "index3"]
72+
73+
74+
class TestIndexConfig:
75+
def test_default_pool_threads(self):
76+
pc = Pinecone(api_key="123-456-789")
77+
index = pc.Index(host='my-host.svg.pinecone.io')
78+
assert index._api_client.pool_threads == 1
79+
80+
def test_pool_threads_when_indexapi_passed(self):
81+
pc = Pinecone(api_key="123-456-789", pool_threads=2, index_api=ManageIndexesApi())
82+
index = pc.Index(host='my-host.svg.pinecone.io')
83+
assert index._api_client.pool_threads == 2
84+
85+
def test_target_index_with_pool_threads_inherited(self):
86+
pc = Pinecone(api_key="123-456-789", pool_threads=10, foo='bar')
87+
index = pc.Index(host='my-host.svg.pinecone.io')
88+
assert index._api_client.pool_threads == 10
89+
90+
def test_target_index_with_pool_threads_kwarg(self):
91+
pc = Pinecone(api_key="123-456-789", pool_threads=10)
92+
index = pc.Index(host='my-host.svg.pinecone.io', pool_threads=5)
93+
assert index._api_client.pool_threads == 5
94+
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from pinecone.grpc import PineconeGRPC, GRPCClientConfig
2+
3+
class TestGRPCIndexInitialization:
4+
def test_init_with_default_config(self):
5+
pc = PineconeGRPC(api_key='YOUR_API_KEY')
6+
index = pc.Index(name='my-index', host='host')
7+
8+
assert index.grpc_client_config.secure == True
9+
assert index.grpc_client_config.timeout == 20
10+
assert index.grpc_client_config.conn_timeout == 1
11+
assert index.grpc_client_config.reuse_channel == True
12+
assert index.grpc_client_config.retry_config == None
13+
assert index.grpc_client_config.grpc_channel_options == None
14+
15+
def test_init_with_grpc_config_from_dict(self):
16+
pc = PineconeGRPC(api_key='YOUR_API_KEY')
17+
config = GRPCClientConfig._from_dict({'timeout': 10})
18+
index = pc.Index(name='my-index', host='host', grpc_config=config)
19+
20+
assert index.grpc_client_config.timeout == 10
21+
22+
# Unset fields still get default values
23+
assert index.grpc_client_config.reuse_channel == True
24+
assert index.grpc_client_config.secure == True
25+
26+
27+
def test_init_with_grpc_config_non_dict(self):
28+
pc = PineconeGRPC(api_key='YOUR_API_KEY')
29+
config = GRPCClientConfig(timeout=10, secure=False)
30+
index = pc.Index(name='my-index', host='host', grpc_config=config)
31+
32+
assert index.grpc_client_config.timeout == 10
33+
assert index.grpc_client_config.secure == False
34+
35+
# Unset fields still get default values
36+
assert index.grpc_client_config.reuse_channel == True
37+
assert index.grpc_client_config.conn_timeout == 1
38+
39+
def test_config_passed_when_target_by_name(self):
40+
pc = PineconeGRPC(api_key='YOUR_API_KEY')
41+
42+
# Set this state in the host store to skip network call
43+
# to find host for name
44+
pc.index_host_store.set_host(pc.config, 'my-index', 'myhost')
45+
46+
config = GRPCClientConfig(timeout=10, secure=False)
47+
index = pc.Index(name='my-index', grpc_config=config)
48+
49+
assert index.grpc_client_config.timeout == 10
50+
assert index.grpc_client_config.secure == False
51+
52+
# Unset fields still get default values
53+
assert index.grpc_client_config.reuse_channel == True
54+
assert index.grpc_client_config.conn_timeout == 1
55+
56+
def test_config_passed_when_target_by_host(self):
57+
pc = PineconeGRPC(api_key='YOUR_API_KEY')
58+
config = GRPCClientConfig(timeout=5, secure=True)
59+
index = pc.Index(host='myhost', grpc_config=config)
60+
61+
assert index.grpc_client_config.timeout == 5
62+
assert index.grpc_client_config.secure == True
63+
64+
# Unset fields still get default values
65+
assert index.grpc_client_config.reuse_channel == True
66+
assert index.grpc_client_config.conn_timeout == 1

0 commit comments

Comments
 (0)