Merge pull request #5 from GetStream/chore-cluster-names

support for cluster names from the env

Author: darkoatanasovski

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()

pkg/server/server.go

@@ -25,19 +25,19 @@ 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(),
 	}
 
 	return &Server{
 		clusters: map[string]*ClusterConnection{
-			"default": defaultConn,
+			name: defaultConn,
 		},
-		activeCluster:  "default",
+		activeCluster:  name,
 		defaultPDAddrs: pdAddrs,
 	}
 }

pkg/types/cluster.go

@@ -0,0 +1,6 @@
+package types
+
+type Cluster struct {
+	Name    string   `json:"name"`
+	PDAddrs []string `json:"pd_addrs"`
+}

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
+}