Skip to content

Commit b8c12eb

Browse files
fsxfreakjhamon
andauthored
Support overriding additional_headers with PINECONE_ADDITIONAL_HEADERS environment variable (#304)
## Problem We need a way to attach special request headers in integration testing. ## Solution Add support for overriding additional_headers via the `PINECONE_ADDITIONAL_HEADERS` environment variable. ## Type of Change - [x] New feature (non-breaking change which adds functionality) ## Test Plan Unit tests ``` $ cd pinecone-python-client/ $ export PINECONE_ADDITIONAL_HEADERS='{"header": "value"}' $ python3 >>> from pinecone import Pinecone >>> client = Pinecone(api_key='key', host='host') >>> vars(client) {'config': ... additional_headers={'header': 'value'}), ...} ``` --------- Co-authored-by: Jennifer Hamon <[email protected]>
1 parent 921e0b4 commit b8c12eb

File tree

2 files changed

+19
-2
lines changed

2 files changed

+19
-2
lines changed

pinecone/config/pinecone_config.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,24 @@
11
from typing import Optional, Dict
2+
import logging
3+
import json
24
import os
35
from .config import ConfigBuilder, Config
46

7+
logger = logging.getLogger(__name__)
8+
59
DEFAULT_CONTROLLER_HOST = "https://api.pinecone.io"
610

711

812
class PineconeConfig():
913
@staticmethod
1014
def build(api_key: Optional[str] = None, host: Optional[str] = None, additional_headers: Optional[Dict[str, str]] = {}, **kwargs) -> Config:
1115
host = host or kwargs.get("host") or os.getenv("PINECONE_CONTROLLER_HOST") or DEFAULT_CONTROLLER_HOST
16+
headers_json = os.getenv("PINECONE_ADDITIONAL_HEADERS")
17+
if headers_json:
18+
try:
19+
headers = json.loads(headers_json)
20+
additional_headers = additional_headers or headers
21+
except Exception as e:
22+
logger.warn(f'Ignoring PINECONE_ADDITIONAL_HEADERS: {e}')
23+
1224
return ConfigBuilder.build(api_key=api_key, host=host, additional_headers=additional_headers, **kwargs)

tests/unit/test_config.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ def run_before_and_after_tests(tmpdir):
1414
# Defend against unexpected env vars. Since we clear these variables below
1515
# after each test execution, these should only be raised if there is
1616
# test pollution in the environment coming from some other test file/setup.
17-
known_env_vars = ["PINECONE_API_KEY", "PINECONE_ENVIRONMENT", "PINECONE_CONTROLLER_HOST"]
17+
known_env_vars = ["PINECONE_API_KEY", "PINECONE_ENVIRONMENT", "PINECONE_CONTROLLER_HOST", "PINECONE_ADDITIONAL_HEADERS"]
1818
for var in known_env_vars:
1919
if os.getenv(var):
2020
raise ValueError(f"Unexpected env var {var} found in environment. Check for test pollution.")
@@ -29,11 +29,13 @@ def run_before_and_after_tests(tmpdir):
2929
def test_init_with_environment_vars(self):
3030
os.environ["PINECONE_API_KEY"] = "test-api-key"
3131
os.environ["PINECONE_CONTROLLER_HOST"] = "https://test-controller-host"
32+
os.environ["PINECONE_ADDITIONAL_HEADERS"] = '{"header": "value"}'
3233

3334
config = PineconeConfig.build()
3435

3536
assert config.api_key == "test-api-key"
3637
assert config.host == "https://test-controller-host"
38+
assert config.additional_headers == {"header": "value"}
3739

3840
def test_init_with_positional_args(self):
3941
api_key = "my-api-key"
@@ -62,14 +64,17 @@ def test_resolution_order_kwargs_over_env_vars(self):
6264
"""
6365
os.environ["PINECONE_API_KEY"] = "env-var-api-key"
6466
os.environ["PINECONE_CONTROLLER_HOST"] = "env-var-controller-host"
67+
os.environ["PINECONE_ADDITIONAL_HEADERS"] = '{"header": "value1"}'
6568

6669
api_key = "kwargs-api-key"
6770
controller_host = "kwargs-controller-host"
71+
additional_headers = {"header": "value2"}
6872

69-
config = PineconeConfig.build(api_key=api_key, host=controller_host)
73+
config = PineconeConfig.build(api_key=api_key, host=controller_host, additional_headers=additional_headers)
7074

7175
assert config.api_key == api_key
7276
assert config.host == 'https://' + controller_host
77+
assert config.additional_headers == additional_headers
7378

7479
def test_errors_when_no_api_key_is_present(self):
7580
with pytest.raises(PineconeConfigurationError):

0 commit comments

Comments
 (0)