FEATURE: Add 'zookeeper start|stop|stat' command to manage zk ensemble

cmd/memcached/start.go

@@ -52,19 +52,12 @@ var startCmd = &cobra.Command{
 				os.Exit(1)
 			}
 
-			client, err := internal.NewSSHClient(ip)
+			session, close, err := internal.NewSSHSession(ip)
 			if err != nil {
 				fmt.Fprintln(os.Stderr, err)
 				os.Exit(1)
 			}
-			defer client.Close()
-
-			session, err := client.NewSession()
-			if err != nil {
-				fmt.Fprintln(os.Stderr, err)
-				os.Exit(1)
-			}
-			defer session.Close()
+			defer close()
 
 			memcachedPath := os.Getenv("ARCUS_PATH")
 			command := fmt.Sprintf(memcachedStartCommandTemplate,

cmd/memcached/stop.go

@@ -42,19 +42,12 @@ var stopCmd = &cobra.Command{
 				os.Exit(1)
 			}
 
-			client, err := internal.NewSSHClient(ip)
+			session, close, err := internal.NewSSHSession(ip)
 			if err != nil {
 				fmt.Fprintln(os.Stderr, err)
 				os.Exit(1)
 			}
-			defer client.Close()
-
-			session, err := client.NewSession()
-			if err != nil {
-				fmt.Fprintln(os.Stderr, err)
-				os.Exit(1)
-			}
-			defer session.Close()
+			defer close()
 
 			pidFilePath := fmt.Sprintf("%s/memcached-%s.pid", os.Getenv("ARCUS_PATH"), port)
 			command := fmt.Sprintf("kill -INT $(cat %s)", pidFilePath)

cmd/zookeeper/start.go

@@ -0,0 +1,39 @@
+package zookeeper
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/jam2in/arcus-cli/internal"
+	"github.com/spf13/cobra"
+)
+
+var startCmd = &cobra.Command{
+	Use:   "start",
+	Short: "Start the zookeeper ensemble servers",
+	Run: func(cmd *cobra.Command, args []string) {
+		zkAddr, zkPath := os.Getenv("ZK_ADDR"), os.Getenv("ZK_PATH")
+		if zkAddr == "" || zkPath == "" {
+			fmt.Fprintln(os.Stderr, "Environment variables are not provided. \nPlease set the ZK_ADDR, ZK_PATH environment variables")
+			os.Exit(1)
+		}
+		zkServers := strings.Split(zkAddr, ",")
+		for _, server := range zkServers {
+			ip := strings.Split(server, ":")[0]
+			command := fmt.Sprintf(zookeeperStartCommandTemplate, zkPath)
+			session, close, err := internal.NewSSHSession(ip)
+			if err != nil {
+				fmt.Fprintln(os.Stderr, err)
+				os.Exit(1)
+			}
+			defer close()
+
+			if err := session.Run(command); err != nil {
+				fmt.Fprintln(os.Stderr, err)
+				os.Exit(1)
+			}
+		}
+		fmt.Printf("zookeeper server started successfully!\n")
+	},
+}

cmd/zookeeper/stat.go

@@ -0,0 +1,52 @@
+package zookeeper
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+)
+
+var statCmd = &cobra.Command{
+	Use:   "stat",
+	Short: "Show the status of the Zookeeper ensemble servers",
+	Run: func(cmd *cobra.Command, args []string) {
+		zkAddr := os.Getenv("ZK_ADDR")
+		if zkAddr == "" {
+			fmt.Fprintln(os.Stderr, "Environment variable is not provided. \nPlease set the ZK_ADDR environment variable")
+			os.Exit(1)
+		}
+		serversToCheck := strings.Split(zkAddr, ",")
+		for _, server := range serversToCheck {
+			status, err := getStatusFromZK(server)
+			if err != nil {
+				fmt.Fprintln(os.Stderr, err)
+				continue
+			}
+			fmt.Println(status)
+		}
+	},
+}
+
+func getStatusFromZK(server string) (string, error) {
+	conn, err := net.Dial("tcp", server)
+	if err != nil {
+		return "", fmt.Errorf("could not connect to server %s: %v", server, err)
+	}
+	defer conn.Close()
+
+	_, err = conn.Write([]byte("stat"))
+	if err != nil {
+		return "", fmt.Errorf("could not send command to server %s: %v", server, err)
+	}
+
+	response, err := io.ReadAll(conn)
+	if err != nil {
+		return "", fmt.Errorf("could not read response from server %s: %v", server, err)
+	}
+
+	return string(response), nil
+}

cmd/zookeeper/stop.go

@@ -0,0 +1,39 @@
+package zookeeper
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/jam2in/arcus-cli/internal"
+	"github.com/spf13/cobra"
+)
+
+var stopCmd = &cobra.Command{
+	Use:   "stop",
+	Short: "Stop the zookeeper ensemble servers",
+	Run: func(cmd *cobra.Command, args []string) {
+		zkAddr, zkPath := os.Getenv("ZK_ADDR"), os.Getenv("ZK_PATH")
+		if zkAddr == "" || zkPath == "" {
+			fmt.Fprintln(os.Stderr, "Environment variables are not provided. \nPlease set the ZK_ADDR, ZK_PATH environment variables")
+			os.Exit(1)
+		}
+		zkServers := strings.Split(zkAddr, ",")
+		for _, server := range zkServers {
+			ip := strings.Split(server, ":")[0]
+			command := fmt.Sprintf(zookeeperStopCommandTemplate, zkPath)
+			session, close, err := internal.NewSSHSession(ip)
+			if err != nil {
+				fmt.Fprintln(os.Stderr, err)
+				os.Exit(1)
+			}
+			defer close()
+
+			if err := session.Run(command); err != nil {
+				fmt.Fprintln(os.Stderr, err)
+				os.Exit(1)
+			}
+		}
+		fmt.Printf("zookeeper server stopped successfully!\n")
+	},
+}

cmd/zookeeper/zookeeper.go

@@ -6,6 +6,11 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const (
+	zookeeperStartCommandTemplate = "%s/bin/zkServer.sh start"
+	zookeeperStopCommandTemplate  = "%s/bin/zkServer.sh stop"
+)
+
 var ZookeeperCmd = &cobra.Command{
 	Use:   "zookeeper",
 	Short: "A CLI tool for zookeeper commands",
@@ -29,4 +34,7 @@ var ZookeeperCmd = &cobra.Command{
 
 func init() {
 	ZookeeperCmd.AddCommand(initCmd)
+	ZookeeperCmd.AddCommand(startCmd)
+	ZookeeperCmd.AddCommand(statCmd)
+	ZookeeperCmd.AddCommand(stopCmd)
 }

internal/ssh.go

@@ -7,21 +7,21 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-func NewSSHClient(ip string) (*ssh.Client, error) {
+func NewSSHSession(ip string) (*ssh.Session, func(), error) {
 	current, err := user.Current()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	username := current.Username
 
 	keyPath := os.Getenv("HOME") + "/.ssh/id_rsa"
 	key, err := os.ReadFile(keyPath)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	signer, err := ssh.ParsePrivateKey(key)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	sshConfig := &ssh.ClientConfig{
@@ -32,5 +32,21 @@ func NewSSHClient(ip string) (*ssh.Client, error) {
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
 
-	return ssh.Dial("tcp", ip+":22", sshConfig)
+	client, err := ssh.Dial("tcp", ip+":22", sshConfig)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	session, err := client.NewSession()
+	if err != nil {
+		client.Close()
+		return nil, nil, err
+	}
+
+	close := func() {
+		client.Close()
+		session.Close()
+	}
+
+	return session, close, err
 }