Merge pull request #1 from ken8203/refactor/cmd

cmd/delete.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"io"
 
 	"github.com/spf13/cobra"
 )
@@ -11,17 +12,22 @@ var deleteCmd = &cobra.Command{
 	Use:     "delete",
 	Aliases: []string{"del"},
 	Short:   "Delete a key",
-	RunE:    deleteRunE,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return delete(cmd.Context(), cmd.OutOrStdout(), args)
+	},
 }
 
-func deleteRunE(cmd *cobra.Command, args []string) error {
-	return delete(cmd.Context(), args)
-}
-
-func delete(ctx context.Context, args []string) error {
+func delete(ctx context.Context, w io.Writer, args []string) error {
 	if len(args) < 1 {
 		return fmt.Errorf("%w 'DELETE'", ErrInvalidArgs)
 	}
 
-	return c.Delete(ctx, []byte(args[0]))
+	if err := c.Delete(ctx, []byte(args[0])); err != nil {
+		return err
+	}
+
+	if _, err := fmt.Fprintln(w, "OK"); err != nil {
+		return err
+	}
+	return nil
 }

cmd/get.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"os"
+	"io"
 
 	"github.com/spf13/cobra"
 	tikverror "github.com/tikv/client-go/v2/error"
@@ -13,33 +13,26 @@ import (
 var getCmd = &cobra.Command{
 	Use:   "get",
 	Short: "Get a key",
-	RunE:  getRunE,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return get(cmd.Context(), cmd.OutOrStdout(), args)
+	},
 }
 
-func getRunE(cmd *cobra.Command, args []string) error {
-	value, err := get(cmd.Context(), args)
-	if err != nil {
-		return err
-	}
-
-	if _, err := fmt.Fprintln(os.Stdout, string(value)); err != nil {
-		return err
-	}
-	return nil
-}
-
-func get(ctx context.Context, args []string) ([]byte, error) {
+func get(ctx context.Context, w io.Writer, args []string) error {
 	if len(args) < 1 {
-		return nil, fmt.Errorf("%w 'GET'", ErrInvalidArgs)
+		return fmt.Errorf("%w 'GET'", ErrInvalidArgs)
 	}
 
 	value, err := c.Get(ctx, []byte(args[0]))
 	if err != nil {
 		if errors.Is(err, tikverror.ErrNotExist) {
-			return nil, ErrNotExist
+			return ErrNotExist
 		}
-		return nil, err
+		return err
 	}
 
-	return value, nil
+	if _, err := fmt.Fprintln(w, string(value)); err != nil {
+		return err
+	}
+	return nil
 }

cmd/put.go

@@ -3,27 +3,30 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"io"
 
 	"github.com/spf13/cobra"
 )
 
 var putCmd = &cobra.Command{
 	Use:   "put",
 	Short: "Put a key",
-	RunE:  putRunE,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return put(cmd.Context(), cmd.OutOrStdout(), args)
+	},
 }
 
-func putRunE(cmd *cobra.Command, args []string) error {
-	return put(cmd.Context(), args)
-}
-
-func put(ctx context.Context, args []string) error {
+func put(ctx context.Context, w io.Writer, args []string) error {
 	if len(args) < 2 {
 		return fmt.Errorf("%w 'PUT'", ErrInvalidArgs)
 	}
 
 	if err := c.Put(ctx, []byte(args[0]), []byte(args[1])); err != nil {
 		return err
 	}
+
+	if _, err := fmt.Fprintln(w, "OK"); err != nil {
+		return err
+	}
 	return nil
 }

cmd/shell.go

@@ -17,41 +17,43 @@ func shellRunE(cmd *cobra.Command, _ []string) error {
 	for command := range ch {
 		switch strings.ToLower(command.Cmd) {
 		case putCmd.Name():
-			if err := put(ctx, command.Args); err != nil {
+			if err := put(ctx, term, command.Args); err != nil {
 				fmt.Fprintln(term, err.Error())
 				break
 			}
 
-			fmt.Fprintln(term, "OK")
 			break
 
 		case getCmd.Name():
-			value, err := get(ctx, command.Args)
-			if err != nil {
+			if err := get(ctx, term, command.Args); err != nil {
 				fmt.Fprintln(term, err.Error())
 				break
 			}
 
-			fmt.Fprintln(term, string(value))
 			break
 
 		case deleteCmd.Name():
-			if err := delete(ctx, command.Args); err != nil {
+			if err := delete(ctx, term, command.Args); err != nil {
 				fmt.Fprintln(term, err.Error())
 				break
 			}
 
-			fmt.Fprintln(term, "OK")
 			break
 
 		case ttlCmd.Name():
-			ttl, err := ttl(ctx, command.Args)
-			if err != nil {
+			if err := ttl(ctx, term, command.Args); err != nil {
+				fmt.Fprintln(term, err.Error())
+				break
+			}
+
+			break
+
+		case versionCmd.Name():
+			if err := version(term); err != nil {
 				fmt.Fprintln(term, err.Error())
 				break
 			}
 
-			fmt.Fprintln(term, ttl)
 			break
 
 		default:

cmd/ttl.go

@@ -3,47 +3,35 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"os"
+	"io"
 
+	"github.com/ken8203/tikv-cli/internal/client"
 	"github.com/spf13/cobra"
 )
 
 var ttlCmd = &cobra.Command{
 	Use:   "ttl",
 	Short: "Get the TTL of a key",
-	RunE:  ttlRunE,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return ttl(cmd.Context(), cmd.OutOrStdout(), args)
+	},
 }
 
-func ttlRunE(cmd *cobra.Command, args []string) error {
-	client, err := newClient()
-	if err != nil {
-		return fmt.Errorf("new client: %v", err)
+func ttl(ctx context.Context, w io.Writer, args []string) error {
+	if Mode == string(client.ModeTxn) {
+		return fmt.Errorf("%w 'TTL'", ErrCommandNotSupported)
+	}
+	if len(args) < 1 {
+		return fmt.Errorf("%w 'TTL'", ErrInvalidArgs)
 	}
-	defer client.Close(cmd.Context())
 
-	value, err := ttl(cmd.Context(), args)
+	ttl, err := c.TTL(ctx, []byte(args[0]))
 	if err != nil {
 		return err
 	}
 
-	if _, err := fmt.Fprintln(os.Stdout, value); err != nil {
+	if _, err := fmt.Fprintln(w, ttl); err != nil {
 		return err
 	}
 	return nil
 }
-
-func ttl(ctx context.Context, args []string) (uint64, error) {
-	if Mode == "txn" {
-		return 0, fmt.Errorf("%w 'TTL'", ErrCommandNotSupported)
-	}
-	if len(args) < 1 {
-		return 0, fmt.Errorf("%w 'TTL'", ErrInvalidArgs)
-	}
-
-	ttl, err := c.TTL(ctx, []byte(args[0]))
-	if err != nil {
-		return 0, err
-	}
-
-	return ttl, nil
-}

cmd/version.go

@@ -1,7 +1,8 @@
 package cmd
 
 import (
-	"os"
+	"fmt"
+	"io"
 
 	"github.com/spf13/cobra"
 )
@@ -12,8 +13,12 @@ var Version string
 var versionCmd = &cobra.Command{
 	Use:   "version",
 	Short: "Print the version number of tikv-cli",
-	RunE: func(_ *cobra.Command, _ []string) error {
-		_, err := os.Stdout.Write([]byte("tikv-cli " + Version))
-		return err
+	RunE: func(cmd *cobra.Command, _ []string) error {
+		return version(cmd.OutOrStdout())
 	},
 }
+
+func version(w io.Writer) error {
+	_, err := fmt.Fprintln(w, "tikv-cli", Version)
+	return err
+}