[cli] cobra zookeeper (#14094)

Signed-off-by: Andrew Mason <andrew@planetscale.com>

go/cmd/zk/command/add_auth.go

@@ -0,0 +1,36 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreedto in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package command
+
+import (
+	"github.com/spf13/cobra"
+)
+
+var AddAuth = &cobra.Command{
+	Use:  "addAuth <digest> <user:pass>",
+	Args: cobra.ExactArgs(2),
+	RunE: commandAddAuth,
+}
+
+func commandAddAuth(cmd *cobra.Command, args []string) error {
+	scheme, auth := cmd.Flags().Arg(0), cmd.Flags().Arg(1)
+	return fs.Conn.AddAuth(cmd.Context(), scheme, []byte(auth))
+}
+
+func init() {
+	Root.AddCommand(AddAuth)
+}

go/cmd/zk/command/cat.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreedto in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package command
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+	"github.com/z-division/go-zookeeper/zk"
+	"golang.org/x/term"
+
+	"vitess.io/vitess/go/cmd/zk/internal/zkfilepath"
+	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/zk2topo"
+)
+
+var (
+	catArgs = struct {
+		LongListing bool
+		Force       bool
+		DecodeProto bool
+	}{}
+
+	Cat = &cobra.Command{
+		Use: "cat <path1> [<path2> ...]",
+		Example: `zk cat /zk/path
+
+# List filename before file data
+zk cat -l /zk/path1 /zk/path2`,
+		Args: cobra.MinimumNArgs(1),
+		RunE: commandCat,
+	}
+)
+
+func commandCat(cmd *cobra.Command, args []string) error {
+	resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args())
+	if err != nil {
+		return fmt.Errorf("cat: invalid wildcards: %w", err)
+	}
+	if len(resolved) == 0 {
+		// the wildcards didn't result in anything, we're done
+		return nil
+	}
+
+	hasError := false
+	for _, arg := range resolved {
+		zkPath := zkfilepath.Clean(arg)
+		data, _, err := fs.Conn.Get(cmd.Context(), zkPath)
+		if err != nil {
+			hasError = true
+			if !catArgs.Force || err != zk.ErrNoNode {
+				log.Warningf("cat: cannot access %v: %v", zkPath, err)
+			}
+			continue
+		}
+
+		if catArgs.LongListing {
+			fmt.Printf("%v:\n", zkPath)
+		}
+		decoded := ""
+		if catArgs.DecodeProto {
+			decoded, err = topo.DecodeContent(zkPath, data, false)
+			if err != nil {
+				log.Warningf("cat: cannot proto decode %v: %v", zkPath, err)
+				decoded = string(data)
+			}
+		} else {
+			decoded = string(data)
+		}
+		fmt.Print(decoded)
+		if len(decoded) > 0 && decoded[len(decoded)-1] != '\n' && (term.IsTerminal(int(os.Stdout.Fd())) || catArgs.LongListing) {
+			fmt.Print("\n")
+		}
+	}
+	if hasError {
+		return fmt.Errorf("cat: some paths had errors")
+	}
+	return nil
+}
+
+func init() {
+	Cat.Flags().BoolVarP(&catArgs.LongListing, "longListing", "l", false, "long listing")
+	Cat.Flags().BoolVarP(&catArgs.Force, "force", "f", false, "no warning on nonexistent node")
+	Cat.Flags().BoolVarP(&catArgs.DecodeProto, "decodeProto", "p", false, "decode proto files and display them as text")
+
+	Root.AddCommand(Cat)
+}

go/cmd/zk/command/chmod.go

@@ -0,0 +1,91 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreedto in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package command
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/cmd/zk/internal/zkfilepath"
+	"vitess.io/vitess/go/cmd/zk/internal/zkfs"
+	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/topo/zk2topo"
+)
+
+var Chmod = &cobra.Command{
+	Use: "chmod <mode> <path>",
+	Example: `zk chmod n-mode /zk/path
+zk chmod n+mode /zk/path`,
+	Args: cobra.MinimumNArgs(2),
+	RunE: commandChmod,
+}
+
+func commandChmod(cmd *cobra.Command, args []string) error {
+	mode := cmd.Flags().Arg(0)
+	if mode[0] != 'n' {
+		return fmt.Errorf("chmod: invalid mode")
+	}
+
+	addPerms := false
+	if mode[1] == '+' {
+		addPerms = true
+	} else if mode[1] != '-' {
+		return fmt.Errorf("chmod: invalid mode")
+	}
+
+	permMask := zkfs.ParsePermMode(mode[2:])
+
+	resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()[1:])
+	if err != nil {
+		return fmt.Errorf("chmod: invalid wildcards: %w", err)
+	}
+	if len(resolved) == 0 {
+		// the wildcards didn't result in anything, we're done
+		return nil
+	}
+
+	hasError := false
+	for _, arg := range resolved {
+		zkPath := zkfilepath.Clean(arg)
+		aclv, _, err := fs.Conn.GetACL(cmd.Context(), zkPath)
+		if err != nil {
+			hasError = true
+			log.Warningf("chmod: cannot set access %v: %v", zkPath, err)
+			continue
+		}
+		if addPerms {
+			aclv[0].Perms |= permMask
+		} else {
+			aclv[0].Perms &= ^permMask
+		}
+		err = fs.Conn.SetACL(cmd.Context(), zkPath, aclv, -1)
+		if err != nil {
+			hasError = true
+			log.Warningf("chmod: cannot set access %v: %v", zkPath, err)
+			continue
+		}
+	}
+	if hasError {
+		return fmt.Errorf("chmod: some paths had errors")
+	}
+	return nil
+}
+
+func init() {
+	Root.AddCommand(Chmod)
+}

go/cmd/zk/command/cp.go

@@ -0,0 +1,43 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreedto in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package command
+
+import "github.com/spf13/cobra"
+
+var Cp = &cobra.Command{
+	Use: "cp <src> <dst>",
+	Example: `zk cp /zk/path .
+zk cp ./config /zk/path/config
+
+# Trailing slash indicates directory
+zk cp ./config /zk/path/`,
+	Args: cobra.MinimumNArgs(2),
+	RunE: commandCp,
+}
+
+func commandCp(cmd *cobra.Command, args []string) error {
+	switch cmd.Flags().NArg() {
+	case 2:
+		return fs.CopyContext(cmd.Context(), cmd.Flags().Arg(0), cmd.Flags().Arg(1))
+	default:
+		return fs.MultiCopyContext(cmd.Context(), cmd.Flags().Args())
+	}
+}
+
+func init() {
+	Root.AddCommand(Cp)
+}

go/cmd/zk/command/edit.go

@@ -0,0 +1,101 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreedto in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package command
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"os/exec"
+	"path"
+	"time"
+
+	"github.com/spf13/cobra"
+	"github.com/z-division/go-zookeeper/zk"
+
+	"vitess.io/vitess/go/cmd/zk/internal/zkfilepath"
+	"vitess.io/vitess/go/vt/log"
+)
+
+var (
+	editArgs = struct {
+		Force bool
+	}{}
+
+	Edit = &cobra.Command{
+		Use:   "edit <path>",
+		Short: "Create a local copy, edit, and write changes back to cell.",
+		Args:  cobra.ExactArgs(1),
+		RunE:  commandEdit,
+	}
+)
+
+func commandEdit(cmd *cobra.Command, args []string) error {
+	arg := cmd.Flags().Arg(0)
+	zkPath := zkfilepath.Clean(arg)
+	data, stat, err := fs.Conn.Get(cmd.Context(), zkPath)
+	if err != nil {
+		if !editArgs.Force || err != zk.ErrNoNode {
+			log.Warningf("edit: cannot access %v: %v", zkPath, err)
+		}
+		return fmt.Errorf("edit: cannot access %v: %v", zkPath, err)
+	}
+
+	name := path.Base(zkPath)
+	tmpPath := fmt.Sprintf("/tmp/zk-edit-%v-%v", name, time.Now().UnixNano())
+	f, err := os.Create(tmpPath)
+	if err == nil {
+		_, err = f.Write(data)
+		f.Close()
+	}
+	if err != nil {
+		return fmt.Errorf("edit: cannot write file %v", err)
+	}
+
+	editor := exec.Command(os.Getenv("EDITOR"), tmpPath)
+	editor.Stdin = os.Stdin
+	editor.Stdout = os.Stdout
+	editor.Stderr = os.Stderr
+	err = editor.Run()
+	if err != nil {
+		os.Remove(tmpPath)
+		return fmt.Errorf("edit: cannot start $EDITOR: %v", err)
+	}
+
+	fileData, err := os.ReadFile(tmpPath)
+	if err != nil {
+		os.Remove(tmpPath)
+		return fmt.Errorf("edit: cannot read file %v", err)
+	}
+
+	if !bytes.Equal(fileData, data) {
+		// data changed - update if we can
+		_, err = fs.Conn.Set(cmd.Context(), zkPath, fileData, stat.Version)
+		if err != nil {
+			os.Remove(tmpPath)
+			return fmt.Errorf("edit: cannot write zk file %v", err)
+		}
+	}
+	os.Remove(tmpPath)
+	return nil
+}
+
+func init() {
+	Edit.Flags().BoolVarP(&editArgs.Force, "force", "f", false, "no warning on nonexistent node")
+
+	Root.AddCommand(Edit)
+}