Skip to content

Commit

Permalink
Support overriding additional_headers with `PINECONE_ADDITIONAL_HEA…
Browse files Browse the repository at this point in the history
…DERS` 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 <jhamon@pinecone.io>
  • Loading branch information
fsxfreak and jhamon committed Feb 6, 2024
1 parent 921e0b4 commit b8c12eb
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 2 deletions.
12 changes: 12 additions & 0 deletions pinecone/config/pinecone_config.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,24 @@
from typing import Optional, Dict
import logging
import json
import os
from .config import ConfigBuilder, Config

logger = logging.getLogger(__name__)

DEFAULT_CONTROLLER_HOST = "https://api.pinecone.io"


class PineconeConfig():
@staticmethod
def build(api_key: Optional[str] = None, host: Optional[str] = None, additional_headers: Optional[Dict[str, str]] = {}, **kwargs) -> Config:
host = host or kwargs.get("host") or os.getenv("PINECONE_CONTROLLER_HOST") or DEFAULT_CONTROLLER_HOST
headers_json = os.getenv("PINECONE_ADDITIONAL_HEADERS")
if headers_json:
try:
headers = json.loads(headers_json)
additional_headers = additional_headers or headers
except Exception as e:
logger.warn(f'Ignoring PINECONE_ADDITIONAL_HEADERS: {e}')

return ConfigBuilder.build(api_key=api_key, host=host, additional_headers=additional_headers, **kwargs)
9 changes: 7 additions & 2 deletions tests/unit/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ def run_before_and_after_tests(tmpdir):
# Defend against unexpected env vars. Since we clear these variables below
# after each test execution, these should only be raised if there is
# test pollution in the environment coming from some other test file/setup.
known_env_vars = ["PINECONE_API_KEY", "PINECONE_ENVIRONMENT", "PINECONE_CONTROLLER_HOST"]
known_env_vars = ["PINECONE_API_KEY", "PINECONE_ENVIRONMENT", "PINECONE_CONTROLLER_HOST", "PINECONE_ADDITIONAL_HEADERS"]
for var in known_env_vars:
if os.getenv(var):
raise ValueError(f"Unexpected env var {var} found in environment. Check for test pollution.")
Expand All @@ -29,11 +29,13 @@ def run_before_and_after_tests(tmpdir):
def test_init_with_environment_vars(self):
os.environ["PINECONE_API_KEY"] = "test-api-key"
os.environ["PINECONE_CONTROLLER_HOST"] = "https://test-controller-host"
os.environ["PINECONE_ADDITIONAL_HEADERS"] = '{"header": "value"}'

config = PineconeConfig.build()

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

def test_init_with_positional_args(self):
api_key = "my-api-key"
Expand Down Expand Up @@ -62,14 +64,17 @@ def test_resolution_order_kwargs_over_env_vars(self):
"""
os.environ["PINECONE_API_KEY"] = "env-var-api-key"
os.environ["PINECONE_CONTROLLER_HOST"] = "env-var-controller-host"
os.environ["PINECONE_ADDITIONAL_HEADERS"] = '{"header": "value1"}'

api_key = "kwargs-api-key"
controller_host = "kwargs-controller-host"
additional_headers = {"header": "value2"}

config = PineconeConfig.build(api_key=api_key, host=controller_host)
config = PineconeConfig.build(api_key=api_key, host=controller_host, additional_headers=additional_headers)

assert config.api_key == api_key
assert config.host == 'https://' + controller_host
assert config.additional_headers == additional_headers

def test_errors_when_no_api_key_is_present(self):
with pytest.raises(PineconeConfigurationError):
Expand Down

0 comments on commit b8c12eb

Please sign in to comment.