Skip to content

Commit

Permalink
Configure apiKey, additionalHeaders, and host through environme…
Browse files Browse the repository at this point in the history
…nt variables (#22)

## Problem
There was a request to add additional configuration validation to the Go
client as you could make requests with an empty `apiKey`. An empty check
and early error return were added in a previous PR when passing `apiKey`
directly:
https://github.com/pinecone-io/go-pinecone/pull/18/files#diff-37c3a14781b4b6e74bcc0cfc8b2462ee2ae243a720a42669256248a3cb005f01R442-R444

We're also lacking the ability to specify an api key or headers through
environment variables like the Python client, which can be a nice UX
tweak for options in how you configure the client. Additionally, there's
no way to specify a control plane host override, which is useful in
situations where we need to hit a staging URL.

## Solution
- Refactor `buildClientOptions()` into a method on the `NewClientParams`
struct rather than a function which receives the struct.
- Add ability to specify `PINECONE_API_KEY`,
`PINECONE_ADDITIONAL_HEADERS`, and `PINECONE_CONTROLLER_HOST` through
environment variables. I kept the naming aligned with Python for now.
- For these values, if the corresponding field has been passed through
`NewClientParams` directly, the passed value will override anything set
in the environment.
- Add unit tests to validate override behavior, and headers being added
as expected.

## Type of Change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [X] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan
CI and new unit tests.


---
- To see the specific tasks where the Asana app for GitHub is being
used, see below:
  - https://app.asana.com/0/0/1207208972408325
  - https://app.asana.com/0/0/1207208968992757
  • Loading branch information
austin-denoble committed May 16, 2024
1 parent cbb7f75 commit c7eda0a
Show file tree
Hide file tree
Showing 6 changed files with 322 additions and 176 deletions.
2 changes: 1 addition & 1 deletion .env.example
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
API_KEY="<Project API Key>"
PINECONE_API_KEY="<Project API Key>"
TEST_POD_INDEX_NAME="<Pod based Index name>"
TEST_SERVERLESS_INDEX_NAME="<Serverless based Index name>"
4 changes: 2 additions & 2 deletions .github/workflows/ci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ jobs:
- name: Setup Go
uses: actions/setup-go@v4
with:
go-version: '1.21.x'
go-version: "1.21.x"
- name: Install dependencies
run: |
go get ./pinecone
Expand All @@ -21,4 +21,4 @@ jobs:
env:
TEST_POD_INDEX_NAME: ${{ secrets.TEST_POD_INDEX_NAME }}
TEST_SERVERLESS_INDEX_NAME: ${{ secrets.TEST_SERVERLESS_INDEX_NAME }}
API_KEY: ${{ secrets.API_KEY }}
PINECONE_API_KEY: ${{ secrets.API_KEY }}
173 changes: 131 additions & 42 deletions pinecone/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,41 +5,78 @@ import (
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"net/url"
"os"
"strings"

"github.com/deepmap/oapi-codegen/v2/pkg/securityprovider"
"github.com/pinecone-io/go-pinecone/internal/gen/control"
"github.com/pinecone-io/go-pinecone/internal/provider"
"github.com/pinecone-io/go-pinecone/internal/useragent"
)

type Client struct {
apiKey string
headers map[string]string
restClient *control.Client
sourceTag string
headers map[string]string
}

type NewClientParams struct {
ApiKey string // required unless Authorization header provided
SourceTag string // optional
ApiKey string // required - provide through NewClientParams or environment variable PINECONE_API_KEY
Headers map[string]string // optional
Host string // optional
RestClient *http.Client // optional
SourceTag string // optional
}

type NewClientBaseParams struct {
Headers map[string]string
Host string
RestClient *http.Client
SourceTag string
}

func NewClient(in NewClientParams) (*Client, error) {
clientOptions, err := buildClientOptions(in)
if err != nil {
return nil, err
osApiKey := os.Getenv("PINECONE_API_KEY")
hasApiKey := (valueOrFallback(in.ApiKey, osApiKey) != "")

if !hasApiKey {
return nil, fmt.Errorf("no API key provided, please pass an API key for authorization through NewClientParams or set the PINECONE_API_KEY environment variable")
}

apiKeyHeader := struct{ Key, Value string }{"Api-Key", valueOrFallback(in.ApiKey, osApiKey)}

clientHeaders := in.Headers
if clientHeaders == nil {
clientHeaders = make(map[string]string)
clientHeaders[apiKeyHeader.Key] = apiKeyHeader.Value

} else {
clientHeaders[apiKeyHeader.Key] = apiKeyHeader.Value
}

client, err := control.NewClient("https://api.pinecone.io", clientOptions...)
return NewClientBase(NewClientBaseParams{Headers: clientHeaders, Host: in.Host, RestClient: in.RestClient, SourceTag: in.SourceTag})
}

func NewClientBase(in NewClientBaseParams) (*Client, error) {
clientOptions := buildClientBaseOptions(in)
var err error

controlHostOverride := valueOrFallback(in.Host, os.Getenv("PINECONE_CONTROLLER_HOST"))
if controlHostOverride != "" {
controlHostOverride, err = ensureURLScheme(controlHostOverride)
if err != nil {
return nil, err
}
}

client, err := control.NewClient(valueOrFallback(controlHostOverride, "https://api.pinecone.io"), clientOptions...)
if err != nil {
return nil, err
}

c := Client{apiKey: in.ApiKey, restClient: client, sourceTag: in.SourceTag, headers: in.Headers}
c := Client{restClient: client, sourceTag: in.SourceTag, headers: in.Headers}
return &c, nil
}

Expand All @@ -52,13 +89,42 @@ func (c *Client) IndexWithNamespace(host string, namespace string) (*IndexConnec
}

func (c *Client) IndexWithAdditionalMetadata(host string, namespace string, additionalMetadata map[string]string) (*IndexConnection, error) {
idx, err := newIndexConnection(newIndexParameters{apiKey: c.apiKey, host: host, namespace: namespace, sourceTag: c.sourceTag, additionalMetadata: additionalMetadata})
authHeader := c.extractAuthHeader()

// merge additionalMetadata with authHeader
if additionalMetadata != nil {
for _, key := range authHeader {
additionalMetadata[key] = authHeader[key]
}
} else {
additionalMetadata = authHeader
}

idx, err := newIndexConnection(newIndexParameters{host: host, namespace: namespace, sourceTag: c.sourceTag, additionalMetadata: additionalMetadata})
if err != nil {
return nil, err
}
return idx, nil
}

func (c *Client) extractAuthHeader() map[string]string {
possibleAuthKeys := []string{
"api-key",
"authorization",
"access_token",
}

for key, value := range c.headers {
for _, checkKey := range possibleAuthKeys {
if strings.ToLower(key) == checkKey {
return map[string]string{key: value}
}
}
}

return nil
}

func (c *Client) ListIndexes(ctx context.Context) ([]*Index, error) {
res, err := c.restClient.ListIndexes(ctx)
if err != nil {
Expand Down Expand Up @@ -407,53 +473,76 @@ func decodeCollection(resBody io.ReadCloser) (*Collection, error) {
return toCollection(&collectionModel), nil
}

func minOne(x int32) int32 {
if x < 1 {
return 1
}
return x
}

func derefOrDefault[T any](ptr *T, defaultValue T) T {
if ptr == nil {
return defaultValue
}
return *ptr
}

func buildClientOptions(in NewClientParams) ([]control.ClientOption, error) {
func buildClientBaseOptions(in NewClientBaseParams) []control.ClientOption {
clientOptions := []control.ClientOption{}
hasAuthorizationHeader := false
hasApiKey := in.ApiKey != ""

// build and apply user agent
userAgentProvider := provider.NewHeaderProvider("User-Agent", useragent.BuildUserAgent(in.SourceTag))
clientOptions = append(clientOptions, control.WithRequestEditorFn(userAgentProvider.Intercept))

for key, value := range in.Headers {
headerProvider := provider.NewHeaderProvider(key, value)
envAdditionalHeaders, hasEnvAdditionalHeaders := os.LookupEnv("PINECONE_ADDITIONAL_HEADERS")
additionalHeaders := make(map[string]string)

if strings.Contains(strings.ToLower(key), "authorization") {
hasAuthorizationHeader = true
// add headers from environment
if hasEnvAdditionalHeaders {
err := json.Unmarshal([]byte(envAdditionalHeaders), &additionalHeaders)
if err != nil {
log.Printf("failed to parse PINECONE_ADDITIONAL_HEADERS: %v", err)
}

clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
}

if !hasAuthorizationHeader {
apiKeyProvider, err := securityprovider.NewSecurityProviderApiKey("header", "Api-Key", in.ApiKey)
if err != nil {
return nil, err
// merge headers from parameters if passed
if in.Headers != nil {
for key, value := range in.Headers {
additionalHeaders[key] = value
}
clientOptions = append(clientOptions, control.WithRequestEditorFn(apiKeyProvider.Intercept))
}

if !hasAuthorizationHeader && !hasApiKey {
return nil, fmt.Errorf("no API key provided, please pass an API key for authorization")
// add headers to client options
for key, value := range additionalHeaders {
headerProvider := provider.NewHeaderProvider(key, value)
clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
}

// apply custom http client if provided
if in.RestClient != nil {
clientOptions = append(clientOptions, control.WithHTTPClient(in.RestClient))
}

return clientOptions, nil
return clientOptions
}

func ensureURLScheme(inputURL string) (string, error) {
parsedURL, err := url.Parse(inputURL)
if err != nil {
return "", fmt.Errorf("invalid URL: %v", err)
}

if parsedURL.Scheme == "" {
return "https://" + inputURL, nil
}
return inputURL, nil
}

func valueOrFallback[T comparable](value, fallback T) T {
var zero T
if value != zero {
return value
} else {
return fallback
}
}

func derefOrDefault[T any](ptr *T, defaultValue T) T {
if ptr == nil {
return defaultValue
}
return *ptr
}

func minOne(x int32) int32 {
if x < 1 {
return 1
}
return x
}
Loading

0 comments on commit c7eda0a

Please sign in to comment.