tools/ctl: add caller ID for pd-ctl (#8214)

ref #7300

Signed-off-by: husharp <[email protected]>

tools/pd-ctl/pdctl/command/global.go

@@ -33,7 +33,7 @@ import (
 )
 
 const (
-	pdControlCallerID = "pd-ctl"
+	PDControlCallerID = "pd-ctl"
 	clusterPrefix     = "pd/api/v1/cluster"
 )
 
@@ -107,7 +107,7 @@ func initNewPDClient(cmd *cobra.Command, opts ...pd.ClientOption) error {
 	if PDCli != nil {
 		PDCli.Close()
 	}
-	PDCli = pd.NewClient(pdControlCallerID, getEndpoints(cmd), opts...)
+	PDCli = pd.NewClient(PDControlCallerID, getEndpoints(cmd), opts...).WithCallerID(PDControlCallerID)
 	return nil
 }
 
@@ -122,7 +122,7 @@ func initNewPDClientWithTLS(cmd *cobra.Command, caPath, certPath, keyPath string
 
 // TODO: replace dialClient with the PD HTTP client completely.
 var dialClient = &http.Client{
-	Transport: apiutil.NewCallerIDRoundTripper(http.DefaultTransport, pdControlCallerID),
+	Transport: apiutil.NewCallerIDRoundTripper(http.DefaultTransport, PDControlCallerID),
 }
 
 // RequireHTTPSClient creates a HTTPS client if the related flags are set
@@ -153,7 +153,7 @@ func initHTTPSClient(caPath, certPath, keyPath string) error {
 	}
 	dialClient = &http.Client{
 		Transport: apiutil.NewCallerIDRoundTripper(
-			&http.Transport{TLSClientConfig: tlsConfig}, pdControlCallerID),
+			&http.Transport{TLSClientConfig: tlsConfig}, PDControlCallerID),
 	}
 	return nil
 }

tools/pd-ctl/tests/global_test.go

@@ -16,33 +16,44 @@ package tests
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"testing"
 
-	"github.com/pingcap/log"
+	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/assertutil"
 	"github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/server"
 	cmd "github.com/tikv/pd/tools/pd-ctl/pdctl"
-	"go.uber.org/zap"
+	"github.com/tikv/pd/tools/pd-ctl/pdctl/command"
 )
 
-const pdControlCallerID = "pd-ctl"
-
 func TestSendAndGetComponent(t *testing.T) {
 	re := require.New(t)
 	handler := func(context.Context, *server.Server) (http.Handler, apiutil.APIServiceGroup, error) {
 		mux := http.NewServeMux()
+		// check pd http sdk api
+		mux.HandleFunc("/pd/api/v1/cluster", func(w http.ResponseWriter, r *http.Request) {
+			callerID := apiutil.GetCallerIDOnHTTP(r)
+			re.Equal(command.PDControlCallerID, callerID)
+			cluster := &metapb.Cluster{Id: 1}
+			clusterBytes, err := json.Marshal(cluster)
+			re.NoError(err)
+			w.Write(clusterBytes)
+		})
+		// check http client api
+		// TODO: remove this comment after replacing dialClient with the PD HTTP client completely.
 		mux.HandleFunc("/pd/api/v1/health", func(w http.ResponseWriter, r *http.Request) {
 			callerID := apiutil.GetCallerIDOnHTTP(r)
-			for k := range r.Header {
-				log.Info("header", zap.String("key", k))
-			}
-			log.Info("caller id", zap.String("caller-id", callerID))
-			re.Equal(pdControlCallerID, callerID)
+			re.Equal(command.PDControlCallerID, callerID)
+			fmt.Fprint(w, callerID)
+		})
+		mux.HandleFunc("/pd/api/v1/stores", func(w http.ResponseWriter, r *http.Request) {
+			callerID := apiutil.GetCallerIDOnHTTP(r)
+			re.Equal(command.PDControlCallerID, callerID)
 			fmt.Fprint(w, callerID)
 		})
 		info := apiutil.APIServiceGroup{
@@ -64,8 +75,20 @@ func TestSendAndGetComponent(t *testing.T) {
 	}()
 
 	cmd := cmd.GetRootCmd()
-	args := []string{"-u", pdAddr, "health"}
+	args := []string{"-u", pdAddr, "cluster"}
 	output, err := ExecuteCommand(cmd, args...)
 	re.NoError(err)
-	re.Equal(fmt.Sprintf("%s\n", pdControlCallerID), string(output))
+	re.Equal(fmt.Sprintf("%s\n", `{
+  "id": 1
+}`), string(output))
+
+	args = []string{"-u", pdAddr, "health"}
+	output, err = ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	re.Equal(fmt.Sprintf("%s\n", command.PDControlCallerID), string(output))
+
+	args = []string{"-u", pdAddr, "store"}
+	output, err = ExecuteCommand(cmd, args...)
+	re.NoError(err)
+	re.Equal(fmt.Sprintf("%s\n", command.PDControlCallerID), string(output))
 }