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 all 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 }}
PINECONE_API_KEY: ${{ secrets.API_KEY }}
173 changes: 131 additions & 42 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,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) {
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)
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)}
Copy link
Contributor

Choose a reason for hiding this comment

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

Seems that valueOrFallback could be utilized for the "hasApiKey" check?


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"))
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 = 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
Loading