Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support overriding additional_headers with PINECONE_ADDITIONAL_HEADERS environment variable #304

Merged
merged 5 commits into from
Feb 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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