Add user agent http header (#54)

- Use `RoundTripper` to add `User-Agent` header to HTTP request. - Fix version message when the version is dev.
tidbcloud · Dec 19, 2022 · 7114bd7 · 7114bd7
1 parent 8a73027
commit 7114bd7
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/internal/cli/root.go b/internal/cli/root.go
@@ -53,7 +53,7 @@ func Execute(ctx context.Context, ver, commit, buildDate string) {
 			if apiUrl == "" {
 				apiUrl = cloud.DefaultApiUrl
 			}
-			delegate, err := cloud.NewClientDelegate(publicKey, privateKey, apiUrl)
+			delegate, err := cloud.NewClientDelegate(publicKey, privateKey, apiUrl, ver)
 			if err != nil {
 				return nil, err
 			}

diff --git a/internal/cli/version/version.go b/internal/cli/version/version.go
@@ -39,7 +39,7 @@ func VersionCmd(h *internal.Helper, ver, commit, buildDate string) *cobra.Comman
 
 // Format formats a version string with the given information.
 func Format(ver, commit, buildDate string) string {
-	if ver == "" && buildDate == "" && commit == "" {
+	if ver == config.DevVersion && buildDate == "" && commit == "" {
 		return fmt.Sprintf("%s version (built from source)", config.CliName)
 	}
 

diff --git a/internal/service/cloud/api_client.go b/internal/service/cloud/api_client.go
@@ -15,8 +15,10 @@
 package cloud
 
 import (
+	"fmt"
 	"net/http"
 
+	"tidbcloud-cli/internal/config"
 	"tidbcloud-cli/internal/prop"
 
 	apiClient "github.com/c4pt0r/go-tidbcloud-sdk-v1/client"
@@ -29,6 +31,7 @@ import (
 
 const (
 	DefaultApiUrl = "https://api.tidbcloud.com"
+	userAgent     = "User-Agent"
 )
 
 type TiDBCloudClient interface {
@@ -49,8 +52,8 @@ type ClientDelegate struct {
 	c *apiClient.GoTidbcloud
 }
 
-func NewClientDelegate(publicKey string, privateKey string, apiUrl string) (*ClientDelegate, error) {
-	client, err := NewApiClient(publicKey, privateKey, apiUrl)
+func NewClientDelegate(publicKey string, privateKey string, apiUrl string, ver string) (*ClientDelegate, error) {
+	client, err := NewApiClient(publicKey, privateKey, apiUrl, ver)
 	if err != nil {
 		return nil, err
 	}
@@ -83,12 +86,12 @@ func (d *ClientDelegate) ListProjects(params *project.ListProjectsParams, opts .
 	return d.c.Project.ListProjects(params, opts...)
 }
 
-func NewApiClient(publicKey string, privateKey string, apiUrl string) (*apiClient.GoTidbcloud, error) {
+func NewApiClient(publicKey string, privateKey string, apiUrl string, ver string) (*apiClient.GoTidbcloud, error) {
 	httpclient := &http.Client{
-		Transport: &digest.Transport{
+		Transport: NewTransportWithAgent(&digest.Transport{
 			Username: publicKey,
 			Password: privateKey,
-		},
+		}, fmt.Sprintf("%s/%s", config.CliName, ver)),
 	}
 
 	// Parse the URL
@@ -97,5 +100,25 @@ func NewApiClient(publicKey string, privateKey string, apiUrl string) (*apiClien
 		return nil, err
 	}
 
-	return apiClient.New(httpTransport.NewWithClient(u.Host, u.Path, []string{u.Scheme}, httpclient), strfmt.Default), nil
+	transport := httpTransport.NewWithClient(u.Host, u.Path, []string{u.Scheme}, httpclient)
+	return apiClient.New(transport, strfmt.Default), nil
+}
+
+// NewTransportWithAgent returns a new http.RoundTripper that add the User-Agent header,
+// according to https://github.com/go-swagger/go-swagger/issues/1563.
+func NewTransportWithAgent(inner http.RoundTripper, userAgent string) http.RoundTripper {
+	return &UserAgentTransport{
+		inner: inner,
+		Agent: userAgent,
+	}
+}
+
+type UserAgentTransport struct {
+	inner http.RoundTripper
+	Agent string
+}
+
+func (ug *UserAgentTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	r.Header.Set(userAgent, ug.Agent)
+	return ug.inner.RoundTrip(r)
 }