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

Configure apiKey, additionalHeaders, and host through environment variables #22

Merged
merged 11 commits into from
May 16, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 }}
INTEGRATION_PINECONE_API_KEY: ${{ secrets.API_KEY }}
91 changes: 75 additions & 16 deletions pinecone/client.go
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we have two different constructors? One that requires an ApiKey and the other that requires a different auth mechanism (bearer token? oauth token?). It feels clunky that the "ApiKey is required unless..."

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That does seem reasonable. I'm less clear on what this would look like though since we don't officially offer alternative methods of authentication quite yet. The other clients require an API key for instance.

If the user is able to supply headers directly then they could technically always pass their own auth headers. I think this would also be an issue on Python and TypeScript since they don't do any checking of the headers that are provided. You could also consider this a user-error that they'd need to resolve themselves.

I'm a bit unsure of what would make the most sense here to be honest. 🤔

Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@ import (
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"net/url"
"os"
"strings"

"github.com/deepmap/oapi-codegen/v2/pkg/securityprovider"
Expand All @@ -16,25 +19,34 @@ import (

type Client struct {
apiKey string
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is apiKey still needed in the Client struct?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed from Client as well as IndexConnection.

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
Headers map[string]string // optional
Host string // optional
RestClient *http.Client // optional
SourceTag string // optional
}

func NewClient(in NewClientParams) (*Client, error) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd expect this func to just call NewClientBase and configure the ApiKey as an auth header

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It now ultimately calls NewClientBase and passes the ApiKey through as part of the Headers map.

clientOptions, err := buildClientOptions(in)
clientOptions, err := in.buildClientOptions()
if err != nil {
return nil, err
}

client, err := control.NewClient("https://api.pinecone.io", clientOptions...)
controlHostOverride := valueOrFallback(in.Host, os.Getenv("PINECONE_CONTROLLER_HOST"))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"CONTROL" vs "CONTROLLER"? (Assuming this is talking about the control plane?

Copy link
Contributor Author

@austin-denoble austin-denoble May 15, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I went with PINECONE_CONTROLLER_HOST because that's the env key that's used in the Python and TypeScript clients for the same config value:

I suppose it is a bit confusing given it's applied as an override to the control plane. I think staying consistent makes sense, although the variable name here is confusing. What do you think?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice job throwing this in for consistency with other clients.

if controlHostOverride != "" {
controlHostOverride, err = ensureHTTP(controlHostOverride)
if err != nil {
return nil, err
}
}

client, err := control.NewClient(valueOrFallback(controlHostOverride, "https://api.pinecone.io"), clientOptions...)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -421,26 +433,52 @@ func derefOrDefault[T any](ptr *T, defaultValue T) T {
return *ptr
}

func buildClientOptions(in NewClientParams) ([]control.ClientOption, error) {
func (ncp *NewClientParams) buildClientOptions() ([]control.ClientOption, error) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could buildClientOptions to be a wrapper around buildClientBaseOptions?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I slimmed things down to just buildClientBaseOptions, lemme know if that makes sense.

clientOptions := []control.ClientOption{}
hasAuthorizationHeader := false
hasApiKey := in.ApiKey != ""
osApiKey := os.Getenv("PINECONE_API_KEY")
hasApiKey := (ncp.ApiKey != "" || osApiKey != "")
envAdditionalHeaders, hasEnvAdditionalHeaders := os.LookupEnv("PINECONE_ADDITIONAL_HEADERS")

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

for key, value := range in.Headers {
headerProvider := provider.NewHeaderProvider(key, value)
// apply headers from parameters if passed, otherwise use environment headers
if ncp.Headers != nil {
for key, value := range ncp.Headers {
headerProvider := provider.NewHeaderProvider(key, value)

if strings.Contains(strings.ToLower(key), "authorization") {
hasAuthorizationHeader = true
if strings.Contains(strings.ToLower(key), "authorization") {
hasAuthorizationHeader = true
}

clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
}
} else if hasEnvAdditionalHeaders {
additionalHeaders := make(map[string]string)
err := json.Unmarshal([]byte(envAdditionalHeaders), &additionalHeaders)
if err != nil {
log.Printf("failed to parse PINECONE_ADDITIONAL_HEADERS: %v", err)
} else {
for header, value := range additionalHeaders {
headerProvider := provider.NewHeaderProvider(header, value)

if strings.Contains(strings.ToLower(header), "authorization") {
hasAuthorizationHeader = true
}

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

if !hasAuthorizationHeader {
apiKeyProvider, err := securityprovider.NewSecurityProviderApiKey("header", "Api-Key", in.ApiKey)
// if apiKey is provided and no auth header is set, add the apiKey as a header
// apiKey from parameters takes precedence over apiKey from environment
if hasApiKey && !hasAuthorizationHeader {
appliedApiKey := valueOrFallback(ncp.ApiKey, osApiKey)

apiKeyProvider, err := securityprovider.NewSecurityProviderApiKey("header", "Api-Key", appliedApiKey)
if err != nil {
return nil, err
}
Expand All @@ -451,9 +489,30 @@ func buildClientOptions(in NewClientParams) ([]control.ClientOption, error) {
return nil, fmt.Errorf("no API key provided, please pass an API key for authorization")
}

if in.RestClient != nil {
clientOptions = append(clientOptions, control.WithHTTPClient(in.RestClient))
if ncp.RestClient != nil {
clientOptions = append(clientOptions, control.WithHTTPClient(ncp.RestClient))
}

return clientOptions, nil
}

func ensureHTTP(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
}
}
151 changes: 147 additions & 4 deletions pinecone/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (

"github.com/google/uuid"
"github.com/pinecone-io/go-pinecone/internal/mocks"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
Expand All @@ -28,8 +29,8 @@ func TestClient(t *testing.T) {
}

func (ts *ClientTests) SetupSuite() {
apiKey := os.Getenv("API_KEY")
require.NotEmpty(ts.T(), apiKey, "API_KEY env variable not set")
apiKey := os.Getenv("INTEGRATION_PINECONE_API_KEY")
require.NotEmpty(ts.T(), apiKey, "INTEGRATION_PINECONE_API_KEY env variable not set")

ts.podIndex = os.Getenv("TEST_POD_INDEX_NAME")
require.NotEmpty(ts.T(), ts.podIndex, "TEST_POD_INDEX_NAME env variable not set")
Expand Down Expand Up @@ -139,9 +140,52 @@ func (ts *ClientTests) TestHeadersAppliedToRequests() {
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")

testHeaderValue := mockTransport.Req.Header.Get("test-header")
if testHeaderValue != "123456" {
ts.FailNow(fmt.Sprintf("Expected request to have header value '123456', but got '%s'", testHeaderValue))
assert.Equal(ts.T(), "123456", testHeaderValue, "Expected request to have header value '123456', but got '%s'", testHeaderValue)
}

func (ts *ClientTests) TestAdditionalHeadersAppliedToRequest() {
os.Setenv("PINECONE_ADDITIONAL_HEADERS", `{"test-header": "environment-header"}`)

apiKey := "test-api-key"

httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{ApiKey: apiKey, RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")

testHeaderValue := mockTransport.Req.Header.Get("test-header")
assert.Equal(ts.T(), "environment-header", testHeaderValue, "Expected request to have header value 'environment-header', but got '%s'", testHeaderValue)

os.Unsetenv("PINECONE_ADDITIONAL_HEADERS")
}

func (ts *ClientTests) TestHeadersOverrideAdditionalHeaders() {
os.Setenv("PINECONE_ADDITIONAL_HEADERS", `{"test-header": "environment-header"}`)

apiKey := "test-api-key"
headers := map[string]string{"test-header": "param-header"}

httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{ApiKey: apiKey, Headers: headers, RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")

testHeaderValue := mockTransport.Req.Header.Get("test-header")
assert.Equal(ts.T(), "param-header", testHeaderValue, "Expected request to have header value 'param-header', but got '%s'", testHeaderValue)

os.Unsetenv("PINECONE_ADDITIONAL_HEADERS")
}

func (ts *ClientTests) TestAuthorizationHeaderOverridesApiKey() {
Expand Down Expand Up @@ -169,6 +213,105 @@ func (ts *ClientTests) TestAuthorizationHeaderOverridesApiKey() {
}
}

func (ts *ClientTests) TestClientReadsApiKeyFromEnv() {
os.Setenv("PINECONE_API_KEY", "test-env-api-key")

httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")

testHeaderValue := mockTransport.Req.Header.Get("Api-Key")
assert.Equal(ts.T(), "test-env-api-key", testHeaderValue, "Expected request to have header value 'test-env-api-key', but got '%s'", testHeaderValue)

os.Unsetenv("PINECONE_API_KEY")
}

func (ts *ClientTests) TestControllerHostOverride() {
apiKey := "test-api-key"
httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{ApiKey: apiKey, Host: "https://test-controller-host.io", RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")
assert.Equal(ts.T(), "test-controller-host.io", mockTransport.Req.Host, "Expected request to be made to 'test-controller-host.io', but got '%s'", mockTransport.Req.URL.Host)
}

func (ts *ClientTests) TestControllerHostOverrideFromEnv() {
os.Setenv("PINECONE_CONTROLLER_HOST", "https://env-controller-host.io")

apiKey := "test-api-key"
httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{ApiKey: apiKey, RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")
assert.Equal(ts.T(), "env-controller-host.io", mockTransport.Req.Host, "Expected request to be made to 'env-controller-host.io', but got '%s'", mockTransport.Req.URL.Host)

os.Unsetenv("PINECONE_CONTROLLER_HOST")
}

func (ts *ClientTests) TestControllerHostNormalization() {
tests := []struct {
name string
host string
wantHost string
wantScheme string
}{
{
name: "Test with https prefix",
host: "https://pinecone-api.io",
wantHost: "pinecone-api.io",
wantScheme: "https",
}, {
name: "Test with http prefix",
host: "http://pinecone-api.io",
wantHost: "pinecone-api.io",
wantScheme: "http",
}, {
name: "Test without prefix",
host: "pinecone-api.io",
wantHost: "pinecone-api.io",
wantScheme: "https",
},
}

for _, tt := range tests {
ts.Run(tt.name, func() {
apiKey := "test-api-key"
httpClient := mocks.CreateMockClient(`{"indexes": []}`)
client, err := NewClient(NewClientParams{ApiKey: apiKey, Host: tt.host, RestClient: httpClient})
if err != nil {
ts.FailNow(err.Error())
}
mockTransport := httpClient.Transport.(*mocks.MockTransport)

_, err = client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
require.NotNil(ts.T(), mockTransport.Req, "Expected request to be made")

assert.Equal(ts.T(), tt.wantHost, mockTransport.Req.URL.Host, "Expected request to be made to host '%s', but got '%s'", tt.wantHost, mockTransport.Req.URL.Host)
assert.Equal(ts.T(), tt.wantScheme, mockTransport.Req.URL.Scheme, "Expected request to be made to host '%s, but got '%s'", tt.wantScheme, mockTransport.Req.URL.Host)
})
}
}

func (ts *ClientTests) TestListIndexes() {
indexes, err := ts.client.ListIndexes(context.Background())
require.NoError(ts.T(), err)
Expand Down
1 change: 0 additions & 1 deletion pinecone/index_connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,6 @@ func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*Fe
for id, vector := range res.Vectors {
vectors[id] = toVector(vector)
}
fmt.Printf("VECTORS: %+v\n", vectors)

return &FetchVectorsResponse{
Vectors: vectors,
Expand Down
2 changes: 1 addition & 1 deletion pinecone/index_connection_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ type IndexConnectionTests struct {

// Runs the test suite with `go test`
func TestIndexConnection(t *testing.T) {
apiKey := os.Getenv("API_KEY")
apiKey := os.Getenv("INTEGRATION_PINECONE_API_KEY")
assert.NotEmptyf(t, apiKey, "API_KEY env variable not set")

client, err := NewClient(NewClientParams{ApiKey: apiKey})
Expand Down
Loading