Skip to content

Commit

Permalink
First pass at functional impl for source_tag
Browse files Browse the repository at this point in the history
  • Loading branch information
ssmith-pc committed Mar 26, 2024
1 parent bcc1b08 commit 7d6ed31
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 4 deletions.
44 changes: 44 additions & 0 deletions internal/util/user_agent.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
package util

import (
"fmt"
"regexp"
"strings"
)

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 {
// need to set to actual current version
appVersion := "0.0.1"

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_ ]
re := regexp.MustCompile(`[^a-z0-9_ ]`)
userAgent = re.ReplaceAllString(userAgent, "")

// 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)
}
28 changes: 25 additions & 3 deletions pinecone/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,17 +6,34 @@ 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/util"
"io"
"net/http"
)

type Client struct {
apiKey string
restClient *control.Client
sourceTag string
}

type CustomHeader struct {
userAgent string
}

func NewUserAgentProvider(userAgent string) *CustomHeader {
return &CustomHeader{userAgent: userAgent}
}

func (s *CustomHeader) Intercept(ctx context.Context, req *http.Request) error {
req.Header.Set("User-Agent", s.userAgent)
return nil
}

type NewClientParams struct {
ApiKey string
// optional fields
SourceTag string
}

func NewClient(in NewClientParams) (*Client, error) {
Expand All @@ -25,12 +42,17 @@ func NewClient(in NewClientParams) (*Client, error) {
return nil, err
}

client, err := control.NewClient("https://api.pinecone.io", control.WithRequestEditorFn(apiKeyProvider.Intercept))
userAgentProvider := NewUserAgentProvider(util.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
}

Expand All @@ -39,7 +61,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
}
Expand Down
4 changes: 3 additions & 1 deletion pinecone/index_connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"crypto/tls"
"fmt"
"github.com/pinecone-io/go-pinecone/internal/gen/data"
"github.com/pinecone-io/go-pinecone/internal/util"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/metadata"
Expand All @@ -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(util.BuildUserAgentGRPC(sourceTag)),
)

if err != nil {
Expand Down

0 comments on commit 7d6ed31

Please sign in to comment.