Adds optional source_tag to User-Agent

internal/provider/header.go

@@ -0,0 +1,20 @@
+package provider
+
+import (
+	"context"
+	"net/http"
+)
+
+type customHeader struct {
+	name  string
+	value string
+}
+
+func NewHeaderProvider(name string, value string) *customHeader {
+	return &customHeader{name: name, value: value}
+}
+
+func (h *customHeader) Intercept(ctx context.Context, req *http.Request) error {
+	req.Header.Set(h.name, h.value)
+	return nil
+}

internal/provider/header_test.go

@@ -0,0 +1,32 @@
+package provider
+
+import (
+	"context"
+	"net/http"
+	"testing"
+)
+
+func TestCustomHeaderIntercept(t *testing.T) {
+	expectedName := "X-Custom-Header"
+	expectedValue := "Custom-Value"
+	header := NewHeaderProvider(expectedName, expectedValue)
+
+	req, err := http.NewRequest("GET", "https://example.com", nil)
+	if err != nil {
+		t.Fatalf("Failed to create HTTP request: %v", err)
+	}
+
+	ctx := context.Background()
+
+	// Call the Intercept method (method being tested)
+	err = header.Intercept(ctx, req)
+	if err != nil {
+		t.Errorf("Intercept failed: %v", err)
+	}
+
+	// Verify that the custom header is set correctly
+	if req.Header.Get(expectedName) != expectedValue {
+		t.Errorf("Expected header '%s' to have value '%s', got '%s'", expectedName, expectedValue,
+			req.Header.Get(expectedName))
+	}
+}

internal/useragent/useragent.go

@@ -0,0 +1,52 @@
+package useragent
+
+import (
+	"fmt"
+	"strings"
+)
+
+func getPackageVersion() string {
+	// update at release time
+	return "v0.5.0-pre"
+}
+
+func BuildUserAgent(sourceTag string) string {
+	return buildUserAgent("go-client", sourceTag)
+}
+
+func BuildUserAgentGRPC(sourceTag string) string {
+	return buildUserAgent("go-client[grpc]", sourceTag)
+}
+
+func buildUserAgent(appName string, sourceTag string) string {
+	appVersion := getPackageVersion()
+
+	sourceTagInfo := ""
+	if sourceTag != "" {
+		sourceTagInfo = buildSourceTagField(sourceTag)
+	}
+	userAgent := fmt.Sprintf("%s/%s%s", appName, appVersion, sourceTagInfo)
+	return userAgent
+}
+
+func buildSourceTagField(userAgent string) string {
+	// Lowercase
+	userAgent = strings.ToLower(userAgent)
+
+	// Limit charset to [a-z0-9_ ]
+	var strBldr strings.Builder
+	for _, char := range userAgent {
+		if (char >= 'a' && char <= 'z') || (char >= '0' && char <= '9') || char == '_' || char == ' ' {
+			strBldr.WriteRune(char)
+		}
+	}
+	userAgent = strBldr.String()
+
+	// Trim left/right whitespace
+	userAgent = strings.TrimSpace(userAgent)
+
+	// Condense multiple spaces to one, and replace with underscore
+	userAgent = strings.Join(strings.Fields(userAgent), "_")
+
+	return fmt.Sprintf("; source_tag=%s;", userAgent)
+}

internal/useragent/useragent_test.go

@@ -0,0 +1,81 @@
+package useragent
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestBuildUserAgentNoSourceTag(t *testing.T) {
+	sourceTag := ""
+	expectedStartWith := fmt.Sprintf("go-client/%s", getPackageVersion())
+	result := BuildUserAgent(sourceTag)
+	if !strings.HasPrefix(result, expectedStartWith) {
+		t.Errorf("BuildUserAgent(): expected user-agent to start with %s, but got %s", expectedStartWith, result)
+	}
+	if strings.Contains(result, "source_tag") {
+		t.Errorf("BuildUserAgent(): expected user-agent to not contain 'source_tag', but got %s", result)
+	}
+}
+
+func TestBuildUserAgentWithSourceTag(t *testing.T) {
+	sourceTag := "my_source_tag"
+	expectedStartWith := fmt.Sprintf("go-client/%s", getPackageVersion())
+	result := BuildUserAgent(sourceTag)
+	if !strings.HasPrefix(result, expectedStartWith) {
+		t.Errorf("BuildUserAgent(): expected user-agent to start with %s, but got %s", expectedStartWith, result)
+	}
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(): expected user-agent to contain 'source_tag=my_source_tag', but got %s", result)
+	}
+}
+
+func TestBuildUserAgentGRPCNoSourceTag(t *testing.T) {
+	sourceTag := ""
+	expectedStartWith := fmt.Sprintf("go-client[grpc]/%s", getPackageVersion())
+	result := BuildUserAgentGRPC(sourceTag)
+	if !strings.HasPrefix(result, expectedStartWith) {
+		t.Errorf("BuildUserAgent(): expected user-agent to start with %s, but got %s", expectedStartWith, result)
+	}
+	if strings.Contains(result, "source_tag") {
+		t.Errorf("BuildUserAgent(): expected user-agent to not contain 'source_tag', but got %s", result)
+	}
+}
+
+func TestBuildUserAgentGRPCWithSourceTag(t *testing.T) {
+	sourceTag := "my_source_tag"
+	expectedStartWith := fmt.Sprintf("go-client[grpc]/%s", getPackageVersion())
+	result := BuildUserAgentGRPC(sourceTag)
+	if !strings.HasPrefix(result, expectedStartWith) {
+		t.Errorf("BuildUserAgent(): expected user-agent to start with %s, but got %s", expectedStartWith, result)
+	}
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(): expected user-agent to contain 'source_tag=my_source_tag', but got %s", result)
+	}
+}
+
+func TestBuildUserAgentSourceTagIsNormalized(t *testing.T) {
+	sourceTag := "my source tag!!!!"
+	result := BuildUserAgent(sourceTag)
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(\"%s\"): expected user-agent to contain 'source_tag=my_source_tag', but got %s", sourceTag, result)
+	}
+
+	sourceTag = "My Source Tag"
+	result = BuildUserAgent(sourceTag)
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(\"%s\"): expected user-agent to contain 'source_tag=my_source_tag', but got %s", sourceTag, result)
+	}
+
+	sourceTag = "   My Source Tag  123  "
+	result = BuildUserAgent(sourceTag)
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(\"%s\"): expected user-agent to contain 'source_tag=my_source_tag_123', but got %s", sourceTag, result)
+	}
+
+	sourceTag = "   My Source Tag  123 #### !! "
+	result = BuildUserAgent(sourceTag)
+	if !strings.Contains(result, "source_tag=my_source_tag") {
+		t.Errorf("BuildUserAgent(\"%s\"): expected user-agent to contain 'source_tag=my_source_tag_123', but got %s", sourceTag, result)
+	}
+}

pinecone/client.go

@@ -6,17 +6,22 @@ import (
 	"fmt"
 	"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"
 	"io"
 	"net/http"
 )
 
 type Client struct {
 	apiKey     string
 	restClient *control.Client
+	sourceTag  string
 }
 
 type NewClientParams struct {
 	ApiKey string
+	// optional fields
+	SourceTag string
 }
 
 func NewClient(in NewClientParams) (*Client, error) {
@@ -25,12 +30,17 @@ func NewClient(in NewClientParams) (*Client, error) {
 		return nil, err
 	}
 
-	client, err := control.NewClient("https://api.pinecone.io", control.WithRequestEditorFn(apiKeyProvider.Intercept))
+	userAgentProvider := provider.NewHeaderProvider("User-Agent", useragent.BuildUserAgent(in.SourceTag))
+
+	client, err := control.NewClient("https://api.pinecone.io",
+		control.WithRequestEditorFn(apiKeyProvider.Intercept),
+		control.WithRequestEditorFn(userAgentProvider.Intercept),
+	)
 	if err != nil {
 		return nil, err
 	}
 
-	c := Client{apiKey: in.ApiKey, restClient: client}
+	c := Client{apiKey: in.ApiKey, restClient: client, sourceTag: in.SourceTag}
 	return &c, nil
 }
 
@@ -39,7 +49,7 @@ func (c *Client) Index(host string) (*IndexConnection, error) {
 }
 
 func (c *Client) IndexWithNamespace(host string, namespace string) (*IndexConnection, error) {
-	idx, err := newIndexConnection(c.apiKey, host, namespace)
+	idx, err := newIndexConnection(c.apiKey, host, namespace, c.sourceTag)
 	if err != nil {
 		return nil, err
 	}

pinecone/index_connection.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"github.com/pinecone-io/go-pinecone/internal/gen/data"
+	"github.com/pinecone-io/go-pinecone/internal/useragent"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/metadata"
@@ -18,14 +19,15 @@ type IndexConnection struct {
 	grpcConn   *grpc.ClientConn
 }
 
-func newIndexConnection(apiKey string, host string, namespace string) (*IndexConnection, error) {
+func newIndexConnection(apiKey string, host string, namespace string, sourceTag string) (*IndexConnection, error) {
 	config := &tls.Config{}
 	target := fmt.Sprintf("%s:443", host)
 	conn, err := grpc.Dial(
 		target,
 		grpc.WithTransportCredentials(credentials.NewTLS(config)),
 		grpc.WithAuthority(target),
 		grpc.WithBlock(),
+		grpc.WithUserAgent(useragent.BuildUserAgentGRPC(sourceTag)),
 	)
 
 	if err != nil {

pinecone/index_connection_test.go

@@ -59,7 +59,7 @@ func (ts *IndexConnectionTests) SetupSuite() {
 	namespace, err := uuid.NewV7()
 	assert.NoError(ts.T(), err)
 
-	idxConn, err := newIndexConnection(ts.apiKey, ts.host, namespace.String())
+	idxConn, err := newIndexConnection(ts.apiKey, ts.host, namespace.String(), "")
 	assert.NoError(ts.T(), err)
 	ts.idxConn = idxConn