From 7ff4ed877ac3e97eb557b747a072a84e57764b91 Mon Sep 17 00:00:00 2001 From: Austin DeNoble Date: Fri, 13 Oct 2023 17:24:11 -0400 Subject: [PATCH] Support passing custom host for Index operations (#218) ## Problem We need support for supplying a custom `host` value for index operations. ## Solution `pinecone.init()` and `Config` already support the ability to supply a custom `host`. It didn't seem like it was being passed through to the `ApiClient`: - Update `_get_api_instance()` in `manage.py` to check for a `Config.CONTROLLER_HOST` value, and apply it to the `client_config` if so. This will then be used throughout the index operations. - Add a unit test file for `manage.py`, test setting `controller_host` and timeout functionality in `create_index` / `delete_index`. ## Type of Change - [X] New feature (non-breaking change which adds functionality) ## Test Plan Pull this branch down and test locally. I validated myself by passing a bogus `host` in and looking over the errors, then passing the proper `host` via `init` to validate it works properly. ```python >>> poetry shell >>> python3 >>> import pinecone >>> pinecone.init(api_key="123-456-789", environment="my_environment", host: "https:custom-host-foo.com") # Having passed a custom host that's inaccessible should fail >>> pinecone.list_indexes() # urllib3.exceptions.MaxRetryError: HTTPSConnectionPool(host='fooblahblah.io', port=443): Max retries exceeded with # url: /databases (Caused by NameResolutionError(": # Failed to resolve 'fooblahblah.io' ([Errno 8] nodename nor servname provided, or not known)")) # Re-Init - should work as expected >>> pinecone.init(api_key="123-456-789", environment="my_environment") pinecone.list_indexes() # ['austin-3', 'austin-py-test', 'resin--test-index-1', 'test-create-1'] ``` --- pinecone/manage.py | 5 ++++ tests/unit/test_grpc_index.py | 2 +- tests/unit/test_index.py | 2 ++ tests/unit/test_manage.py | 54 +++++++++++++++++++++++++++++++++++ 4 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_manage.py diff --git a/pinecone/manage.py b/pinecone/manage.py index cc470148..76541edf 100644 --- a/pinecone/manage.py +++ b/pinecone/manage.py @@ -53,6 +53,11 @@ def _get_api_instance(): client_config.api_key = client_config.api_key or {} client_config.api_key["ApiKeyAuth"] = client_config.api_key.get("ApiKeyAuth", Config.API_KEY) client_config.server_variables = {**{"environment": Config.ENVIRONMENT}, **client_config.server_variables} + + # If a custom host has been passed with initialization pass it to the client_config + if (Config.CONTROLLER_HOST): + client_config.host = Config.CONTROLLER_HOST + api_client = ApiClient(configuration=client_config) api_client.user_agent = get_user_agent() api_instance = IndexOperationsApi(api_client) diff --git a/tests/unit/test_grpc_index.py b/tests/unit/test_grpc_index.py index d4f5a7bf..2284d81f 100644 --- a/tests/unit/test_grpc_index.py +++ b/tests/unit/test_grpc_index.py @@ -39,7 +39,7 @@ def setup_method(self): sparse_values=SparseValues(indices=self.sparse_indices_2, values=self.sparse_values_2)) - # region: upsert tests + # region: upsert tests def _assert_called_once(self, vectors, async_call=False): self.index._wrap_grpc_call.assert_called_once_with( diff --git a/tests/unit/test_index.py b/tests/unit/test_index.py index fc549646..7092d71c 100644 --- a/tests/unit/test_index.py +++ b/tests/unit/test_index.py @@ -37,6 +37,8 @@ def setup_method(self): pinecone.init(api_key='example-key') self.index = pinecone.Index('example-name') + # region: upsert tests + def test_upsert_numpy_deprecation_warning(self, mocker): mocker.patch.object(self.index._vector_api, 'upsert', autospec=True) with pytest.warns(FutureWarning): diff --git a/tests/unit/test_manage.py b/tests/unit/test_manage.py new file mode 100644 index 00000000..70c7dedb --- /dev/null +++ b/tests/unit/test_manage.py @@ -0,0 +1,54 @@ +import pytest +import pinecone +import time + +class TestManage: + + def test_get_api_instance_without_host(self): + pinecone.init(api_key="123-456-789", environment="my-environment") + api_instance = pinecone.manage._get_api_instance() + assert api_instance.api_client.configuration.host == "https://controller.my-environment.pinecone.io" + + def test_get_api_instance_with_host(self): + pinecone.init(api_key="123-456-789", environment="my-environment", host="my-host") + api_instance = pinecone.manage._get_api_instance() + assert api_instance.api_client.configuration.host == "my-host" + + @pytest.mark.parametrize("timeout_value, get_status_calls, time_sleep_calls, get_status_responses", [ + # No timeout, _get_status called twice, sleep called once + (None, 2, 1, [{"ready": False}, {"ready": True}]), + # Timeout of 10 seconds, _get_status called 3 times, sleep twice + (10, 3, 2, [{"ready": False}, {"ready": False}, {"ready": True}]), + # Timeout of -1 seconds, _get_status not called, no sleep + (-1, 0, 0, [{"ready": False}]), + ]) + def test_create_index_with_timeout(self, mocker, timeout_value, get_status_calls, time_sleep_calls, get_status_responses): + mocker.patch('pinecone.manage._get_api_instance', return_value=mocker.Mock()) + mocker.patch('pinecone.manage._get_status', side_effect=get_status_responses) + mocker.patch('time.sleep') + + pinecone.manage.create_index("my-index", 10, timeout=timeout_value) + + pinecone.manage._get_api_instance.assert_called_once() + assert pinecone.manage._get_status.call_count == get_status_calls + assert time.sleep.call_count == time_sleep_calls + + @pytest.mark.parametrize("timeout_value, list_indexes_calls, time_sleep_calls, list_indexes_responses", [ + # No timeout, list_indexes called twice, sleep called once + (None, 2, 1, [["my-index", "index-1"], ["index-1"]]), + # Timeout of 10 seconds, list_indexes called 3 times, sleep twice + (10, 3, 2, [["my-index", "index-1"], ["my-index", "index-1"], ["index-1"]]), + # Timeout of -1 seconds, list_indexes not called, no sleep + (-1, 0, 0, [["my-index", "index-1"]]), + ]) + def test_delete_index_with_timeout(self, mocker, timeout_value, list_indexes_calls, time_sleep_calls, list_indexes_responses): + api_instance_mock = mocker.Mock() + api_instance_mock.list_indexes = mocker.Mock(side_effect=list_indexes_responses) + mocker.patch('pinecone.manage._get_api_instance', return_value=api_instance_mock) + mocker.patch('time.sleep') + + pinecone.manage.delete_index("my-index", timeout=timeout_value) + + pinecone.manage._get_api_instance.assert_called_once() + assert api_instance_mock.list_indexes.call_count == list_indexes_calls + assert time.sleep.call_count == time_sleep_calls \ No newline at end of file