diff --git a/cmd/main.go b/cmd/main.go index 8ecf397..1d9c743 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -2,12 +2,10 @@ package main import ( "context" - "fmt" "log" "net/http" "os" "os/signal" - "strings" "syscall" "time" @@ -19,29 +17,30 @@ import ( ) func main() { - // Read PD addresses from env: TIKV_PD_ADDRS="127.0.0.1:2379,127.0.0.1:2381" + // Read PD addresses from env: TIKV_PD_ADDRS="127.0.0.1:2379|My Cluster 1,127.0.0.1:2381;server-2.com:2379|Cluster 2" pdAddrsEnv := os.Getenv("TIKV_PD_ADDRS") if pdAddrsEnv == "" { log.Fatal("TIKV_PD_ADDRS env var is required (comma-separated PD addresses)") } - clusterStrs := strings.Split(strings.Trim(pdAddrsEnv, ";"), ";") - pdAddrs := utils.SplitAndTrim(clusterStrs[0], ",") - + clusters := utils.GetClusters(pdAddrsEnv) + if len(clusters) == 0 { + log.Fatal("no clusters found") + } ctx := context.Background() // Create TiKV RawKV client for default cluster - cli, err := rawkv.NewClient(ctx, pdAddrs, config.DefaultConfig().Security) + cli, err := rawkv.NewClient(ctx, clusters[0].PDAddrs, config.DefaultConfig().Security) if err != nil { log.Fatalf("failed to create TiKV RawKV client: %v", err) } log.Printf("Connected to default TiKV cluster ID: %d", cli.ClusterID()) - srv := server.New(cli, pdAddrs) + srv := server.New(cli, clusters[0].PDAddrs, clusters[0].Name) defer srv.Close() - for i, cluster := range clusterStrs[1:] { - srv.AddCluster(ctx, fmt.Sprintf("cluster-%d", time.Now().Unix()+int64(i)), utils.SplitAndTrim(cluster, ",")) + for _, cluster := range clusters[1:] { + srv.AddCluster(ctx, cluster.Name, cluster.PDAddrs) } mux := http.NewServeMux() diff --git a/pkg/server/server.go b/pkg/server/server.go index e5ae077..724d4cc 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -25,9 +25,9 @@ type Server struct { } // New creates a new Server instance with an initial connection -func New(client *rawkv.Client, pdAddrs []string) *Server { +func New(client *rawkv.Client, pdAddrs []string, name string) *Server { defaultConn := &ClusterConnection{ - Name: "default", + Name: name, PDAddrs: pdAddrs, Client: client, ClusterID: client.ClusterID(), @@ -35,9 +35,9 @@ func New(client *rawkv.Client, pdAddrs []string) *Server { return &Server{ clusters: map[string]*ClusterConnection{ - "default": defaultConn, + name: defaultConn, }, - activeCluster: "default", + activeCluster: name, defaultPDAddrs: pdAddrs, } } diff --git a/pkg/types/cluster.go b/pkg/types/cluster.go new file mode 100644 index 0000000..3ff793c --- /dev/null +++ b/pkg/types/cluster.go @@ -0,0 +1,6 @@ +package types + +type Cluster struct { + Name string `json:"name"` + PDAddrs []string `json:"pd_addrs"` +} diff --git a/pkg/utils/string.go b/pkg/utils/string.go index f5b6078..35f81a6 100644 --- a/pkg/utils/string.go +++ b/pkg/utils/string.go @@ -1,6 +1,13 @@ package utils -import "strings" +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "strings" + + "github.com/GetStream/tikv-ui/pkg/types" +) // SplitAndTrim splits a string by separator and trims whitespace from each part func SplitAndTrim(s, sep string) []string { @@ -14,3 +21,35 @@ func SplitAndTrim(s, sep string) []string { } return out } + +func GetClusters(s string) []types.Cluster { + parts := SplitAndTrim(s, ";") + clusters := make([]types.Cluster, 0, len(parts)) + for _, part := range parts { + clusters = append(clusters, GetCluster(part)) + } + return clusters +} + +func GetCluster(s string) types.Cluster { + parts := SplitAndTrim(s, "|") + if len(parts) < 2 { + hex, _ := RandHex(5) + return types.Cluster{ + Name: fmt.Sprintf("cluster-%s", hex), + PDAddrs: SplitAndTrim(parts[0], ","), + } + } + return types.Cluster{ + Name: parts[1], + PDAddrs: SplitAndTrim(parts[0], ","), + } +} + +func RandHex(n int) (string, error) { + b := make([]byte, n) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +}