From 0ab36f465e330e2031eced8f46a56d6d2364ec54 Mon Sep 17 00:00:00 2001
From: lhy1024 <admin@liudos.us>
Date: Wed, 5 Jul 2023 11:49:05 +0800
Subject: [PATCH] use query

Signed-off-by: lhy1024 <admin@liudos.us>
---
 pkg/keyspace/tso_keyspace_group.go            |  9 ++++
 server/apiv2/handlers/tso_keyspace_group.go   | 41 +++++--------------
 .../pdctl/command/keyspace_group_command.go   |  2 +-
 3 files changed, 21 insertions(+), 31 deletions(-)

diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go
index 54887c762cb..d319798738b 100644
--- a/pkg/keyspace/tso_keyspace_group.go
+++ b/pkg/keyspace/tso_keyspace_group.go
@@ -1015,6 +1015,15 @@ func (m *GroupManager) FinishMergeKeyspaceByID(mergeTargetID uint32) error {
 
 // GetKeyspaceGroupPrimaryByID returns the primary node of the keyspace group by ID.
 func (m *GroupManager) GetKeyspaceGroupPrimaryByID(id uint32) (string, error) {
+	// check if the keyspace group exists
+	kg, err := m.GetKeyspaceGroupByID(id)
+	if err != nil {
+		return "", err
+	}
+	if kg == nil {
+		return "", ErrKeyspaceGroupNotExists(id)
+	}
+
 	// default keyspace group: "/ms/{cluster_id}/tso/00000/primary".
 	// non-default keyspace group: "/ms/{cluster_id}/tso/keyspace_groups/election/{group}/primary".
 	path := fmt.Sprintf("/ms/%d/tso/00000/primary", m.clusterID)
diff --git a/server/apiv2/handlers/tso_keyspace_group.go b/server/apiv2/handlers/tso_keyspace_group.go
index 23a5528d2ce..0e8c039d521 100644
--- a/server/apiv2/handlers/tso_keyspace_group.go
+++ b/server/apiv2/handlers/tso_keyspace_group.go
@@ -47,7 +47,6 @@ func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) {
 	router.DELETE("/:id/split", FinishSplitKeyspaceByID)
 	router.POST("/:id/merge", MergeKeyspaceGroups)
 	router.DELETE("/:id/merge", FinishMergeKeyspaceByID)
-	router.GET("/:id/primary", GetKeyspaceGroupPrimaryByID)
 }
 
 // CreateKeyspaceGroupParams defines the params for creating keyspace groups.
@@ -150,12 +149,22 @@ func GetKeyspaceGroupByID(c *gin.Context) {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, groupManagerUninitializedErr)
 		return
 	}
+
+	if c.Query("get_param") == "primary" {
+		primary, err := manager.GetKeyspaceGroupPrimaryByID(id)
+		if err != nil {
+			c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
+			return
+		}
+		c.JSON(http.StatusOK, primary)
+		return
+	}
+
 	kg, err := manager.GetKeyspaceGroupByID(id)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
 		return
 	}
-
 	c.IndentedJSON(http.StatusOK, kg)
 }
 
@@ -486,34 +495,6 @@ func SetPriorityForKeyspaceGroup(c *gin.Context) {
 	c.JSON(http.StatusOK, nil)
 }
 
-// GetKeyspaceGroupPrimaryByID gets primary of keyspace group by ID.
-func GetKeyspaceGroupPrimaryByID(c *gin.Context) {
-	id, err := validateKeyspaceGroupID(c)
-	if err != nil {
-		c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id")
-		return
-	}
-	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
-	manager := svr.GetKeyspaceGroupManager()
-	if manager == nil {
-		c.AbortWithStatusJSON(http.StatusInternalServerError, groupManagerUninitializedErr)
-		return
-	}
-	// check if keyspace group exists
-	kg, err := manager.GetKeyspaceGroupByID(id)
-	if err != nil || kg == nil {
-		c.AbortWithStatusJSON(http.StatusBadRequest, "keyspace group does not exist")
-		return
-	}
-	// get primary
-	primary, err := manager.GetKeyspaceGroupPrimaryByID(id)
-	if err != nil {
-		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
-		return
-	}
-	c.JSON(http.StatusOK, primary)
-}
-
 func validateKeyspaceGroupID(c *gin.Context) (uint32, error) {
 	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
 	if err != nil {
diff --git a/tools/pd-ctl/pdctl/command/keyspace_group_command.go b/tools/pd-ctl/pdctl/command/keyspace_group_command.go
index 665be591e9f..3b75bdc9636 100644
--- a/tools/pd-ctl/pdctl/command/keyspace_group_command.go
+++ b/tools/pd-ctl/pdctl/command/keyspace_group_command.go
@@ -357,7 +357,7 @@ func showKeyspaceGroupPrimaryCommandFunc(cmd *cobra.Command, args []string) {
 		cmd.Printf("Failed to parse the keyspace group ID: %s\n", err)
 		return
 	}
-	r, err := doRequest(cmd, fmt.Sprintf("%s/%s/primary", keyspaceGroupsPrefix, args[0]), http.MethodGet, http.Header{})
+	r, err := doRequest(cmd, fmt.Sprintf("%s/%s?get_param=primary", keyspaceGroupsPrefix, args[0]), http.MethodGet, http.Header{})
 	if err != nil {
 		cmd.Printf("Failed to get the keyspace group primary information: %s\n", err)
 		return