Configure apiKey, additionalHeaders, and host through environment variables

.env.example

@@ -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>"

.github/workflows/ci.yaml

@@ -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
@@ -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 }}

pinecone/client.go

@@ -5,8 +5,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
-	"strings"
+	"net/url"
+	"os"
 
 	"github.com/deepmap/oapi-codegen/v2/pkg/securityprovider"
 	"github.com/pinecone-io/go-pinecone/internal/gen/control"
@@ -16,25 +18,41 @@ import (
 
 type Client struct {
 	apiKey     string
+	headers    map[string]string
 	restClient *control.Client
 	sourceTag  string
-	headers    map[string]string
+}
+
+type NewClientBaseParams struct {
+	Headers    map[string]string
+	Host       string
+	RestClient *http.Client
+	SourceTag  string
 }
 
 type NewClientParams struct {
-	ApiKey     string            // required unless Authorization header provided
-	SourceTag  string            // optional
-	Headers    map[string]string // optional
-	RestClient *http.Client      // optional
+	ApiKey     string
+	Headers    map[string]string
+	Host       string
+	RestClient *http.Client
+	SourceTag  string
 }
 
 func NewClient(in NewClientParams) (*Client, error) {
-	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"))
+	if controlHostOverride != "" {
+		controlHostOverride, err = ensureHTTPS(controlHostOverride)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	client, err := control.NewClient(valueOrFallback(controlHostOverride, "https://api.pinecone.io"), clientOptions...)
 	if err != nil {
 		return nil, err
 	}
@@ -43,6 +61,27 @@ func NewClient(in NewClientParams) (*Client, error) {
 	return &c, nil
 }
 
+func NewClientBase(in NewClientBaseParams) (*Client, error) {
+	clientOptions := in.buildClientBaseOptions()
+	var err error
+
+	controlHostOverride := valueOrFallback(in.Host, os.Getenv("PINECONE_CONTROLLER_HOST"))
+	if controlHostOverride != "" {
+		controlHostOverride, err = ensureHTTPS(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{restClient: client, sourceTag: in.SourceTag, headers: in.Headers}
+	return &c, nil
+}
+
 func (c *Client) Index(host string) (*IndexConnection, error) {
 	return c.IndexWithAdditionalMetadata(host, "", nil)
 }
@@ -421,39 +460,82 @@ 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) {
 	clientOptions := []control.ClientOption{}
-	hasAuthorizationHeader := false
-	hasApiKey := in.ApiKey != ""
+	osApiKey := os.Getenv("PINECONE_API_KEY")
+	hasApiKey := (ncp.ApiKey != "" || osApiKey != "")
 
-	userAgentProvider := provider.NewHeaderProvider("User-Agent", useragent.BuildUserAgent(in.SourceTag))
-	clientOptions = append(clientOptions, control.WithRequestEditorFn(userAgentProvider.Intercept))
+	if !hasApiKey {
+		return nil, fmt.Errorf("no API key provided, please pass an API key for authorization")
+	}
 
-	for key, value := range in.Headers {
-		headerProvider := provider.NewHeaderProvider(key, value)
+	appliedApiKey := valueOrFallback(ncp.ApiKey, osApiKey)
+	apiKeyProvider, err := securityprovider.NewSecurityProviderApiKey("header", "Api-Key", appliedApiKey)
+	if err != nil {
+		return nil, err
+	}
+	clientOptions = append(clientOptions, control.WithRequestEditorFn(apiKeyProvider.Intercept))
 
-		if strings.Contains(strings.ToLower(key), "authorization") {
-			hasAuthorizationHeader = true
-		}
+	baseClient := NewClientBaseParams{ncp.Headers, ncp.Host, ncp.RestClient, ncp.SourceTag}
+	baseClientOptions := baseClient.buildClientBaseOptions()
+	clientOptions = append(clientOptions, baseClientOptions...)
 
-		clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
-	}
+	return clientOptions, nil
+}
+
+func (ncbp *NewClientBaseParams) buildClientBaseOptions() []control.ClientOption {
+	clientOptions := []control.ClientOption{}
+
+	envAdditionalHeaders, hasEnvAdditionalHeaders := os.LookupEnv("PINECONE_ADDITIONAL_HEADERS")
+
+	// build and apply user agent
+	fmt.Printf("Source tag on its way in: %v\n", ncbp.SourceTag)
+	userAgentProvider := provider.NewHeaderProvider("User-Agent", useragent.BuildUserAgent(ncbp.SourceTag))
+	clientOptions = append(clientOptions, control.WithRequestEditorFn(userAgentProvider.Intercept))
 
-	if !hasAuthorizationHeader {
-		apiKeyProvider, err := securityprovider.NewSecurityProviderApiKey("header", "Api-Key", in.ApiKey)
+	// apply headers from parameters if passed, otherwise use environment headers
+	if ncbp.Headers != nil {
+		for key, value := range ncbp.Headers {
+			headerProvider := provider.NewHeaderProvider(key, value)
+			clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
+		}
+	} else if hasEnvAdditionalHeaders {
+		additionalHeaders := make(map[string]string)
+		err := json.Unmarshal([]byte(envAdditionalHeaders), &additionalHeaders)
 		if err != nil {
-			return nil, err
+			log.Printf("failed to parse PINECONE_ADDITIONAL_HEADERS: %v", err)
+		} else {
+			for header, value := range additionalHeaders {
+				headerProvider := provider.NewHeaderProvider(header, value)
+				clientOptions = append(clientOptions, control.WithRequestEditorFn(headerProvider.Intercept))
+			}
 		}
-		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")
+	if ncbp.RestClient != nil {
+		clientOptions = append(clientOptions, control.WithHTTPClient(ncbp.RestClient))
+	}
+
+	return clientOptions
+}
+
+func ensureHTTPS(inputURL string) (string, error) {
+	parsedURL, err := url.Parse(inputURL)
+	if err != nil {
+		return "", fmt.Errorf("invalid URL: %v", err)
 	}
 
-	if in.RestClient != nil {
-		clientOptions = append(clientOptions, control.WithHTTPClient(in.RestClient))
+	if parsedURL.Scheme == "" {
+		return "https://" + inputURL, nil
 	}
+	return inputURL, nil
+}
 
-	return clientOptions, nil
+func valueOrFallback[T comparable](value, fallback T) T {
+	var zero T
+	if value != zero {
+		return value
+	} else {
+		return fallback
+	}
 }