add cobra command utilities and log formatter

Signed-off-by: Adphi <philippe.adrien.nousse@gmail.com>

cli/clifmt/formatter.go

@@ -0,0 +1,84 @@
+// Copyright 2023 Linka Cloud  All rights reserved.
+//
+// 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 agreed to 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 clifmt
+
+import (
+	"bytes"
+	"strings"
+	"time"
+
+	"github.com/fatih/color"
+	"github.com/sirupsen/logrus"
+)
+
+type TimeFormat string
+
+const (
+	NoneTimeFormat     TimeFormat = "none"
+	FullTimeFormat     TimeFormat = "full"
+	RelativeTimeFormat TimeFormat = "relative"
+)
+
+const (
+	red    = 31
+	yellow = 33
+	blue   = 36
+	white  = 39
+	gray   = 90
+)
+
+func New(f TimeFormat) logrus.Formatter {
+	return &clifmt{start: time.Now(), format: f}
+}
+
+type clifmt struct {
+	start  time.Time
+	format TimeFormat
+}
+
+func (f *clifmt) Format(entry *logrus.Entry) ([]byte, error) {
+	var b bytes.Buffer
+	var c *color.Color
+	switch entry.Level {
+	case logrus.DebugLevel, logrus.TraceLevel:
+		c = color.New(gray)
+	case logrus.WarnLevel:
+		c = color.New(yellow)
+	case logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel:
+		c = color.New(red)
+	default:
+		c = color.New(white)
+	}
+	msg := entry.Message
+	if len(entry.Message) > 0 && entry.Level < logrus.DebugLevel {
+		msg = strings.ToTitle(string(msg[0])) + msg[1:]
+	}
+
+	var err error
+	switch f.format {
+	case FullTimeFormat:
+		_, err = c.Fprintf(&b, "[%s] %s\n", entry.Time.Format("2006-01-02 15:04:05"), entry.Message)
+	case RelativeTimeFormat:
+		_, err = c.Fprintf(&b, "[%5v] %s\n", entry.Time.Sub(f.start).Truncate(time.Second).String(), msg)
+	case NoneTimeFormat:
+		fallthrough
+	default:
+		_, err = c.Fprintln(&b, msg)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return b.Bytes(), nil
+}

cli/command.go

@@ -0,0 +1,297 @@
+// Package cli is adapted from https://github.com/rancher/wrangler-cli
+package cli
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+	"regexp"
+	"strconv"
+	"strings"
+	"unsafe"
+
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+
+	"go.linka.cloud/grpc-toolkit/signals"
+)
+
+var (
+	caseRegexp = regexp.MustCompile("([a-z])([A-Z])")
+)
+
+type PersistentPreRunnable interface {
+	PersistentPre(cmd *cobra.Command, args []string) error
+}
+
+type PreRunnable interface {
+	Pre(cmd *cobra.Command, args []string) error
+}
+
+type Runnable interface {
+	Run(cmd *cobra.Command, args []string) error
+}
+
+type customizer interface {
+	Customize(cmd *cobra.Command)
+}
+
+type fieldInfo struct {
+	FieldType  reflect.StructField
+	FieldValue reflect.Value
+}
+
+func fields(obj interface{}) []fieldInfo {
+	ptrValue := reflect.ValueOf(obj)
+	objValue := ptrValue.Elem()
+
+	var result []fieldInfo
+
+	for i := 0; i < objValue.NumField(); i++ {
+		fieldType := objValue.Type().Field(i)
+		if !fieldType.IsExported() {
+			continue
+		}
+		if fieldType.Anonymous && fieldType.Type.Kind() == reflect.Struct {
+			result = append(result, fields(objValue.Field(i).Addr().Interface())...)
+		} else if !fieldType.Anonymous {
+			result = append(result, fieldInfo{
+				FieldValue: objValue.Field(i),
+				FieldType:  objValue.Type().Field(i),
+			})
+		}
+	}
+
+	return result
+}
+
+func Name(obj interface{}) string {
+	ptrValue := reflect.ValueOf(obj)
+	objValue := ptrValue.Elem()
+	commandName := strings.Replace(objValue.Type().Name(), "Command", "", 1)
+	commandName, _ = name(commandName, "", "")
+	return commandName
+}
+
+func Main(cmd *cobra.Command) {
+	ctx := signals.SetupSignalHandler()
+	if err := cmd.ExecuteContext(ctx); err != nil {
+		logrus.Fatal(err)
+	}
+}
+
+func makeEnvVar[T comparable](to []func(), name string, vars []string, defValue T, flags *pflag.FlagSet, fn func(flag string) (T, error)) []func() {
+	for _, v := range vars {
+		to = append(to, func() {
+			v := os.Getenv(v)
+			if v == "" {
+				return
+			}
+			fv, err := fn(name)
+			if err == nil && fv == defValue {
+				flags.Set(name, v)
+			}
+		})
+	}
+	return to
+}
+
+// Command populates a obj.Command() object by extracting args from struct tags of the
+// Runnable obj passed.  Also the Run method is assigned to the RunE of the cli.
+// name = Override the struct field with
+func Command(obj Runnable, c *cobra.Command) *cobra.Command {
+	var (
+		envs     []func()
+		arrays   = map[string]reflect.Value{}
+		slices   = map[string]reflect.Value{}
+		maps     = map[string]reflect.Value{}
+		ptrValue = reflect.ValueOf(obj)
+		objValue = ptrValue.Elem()
+	)
+
+	if c.Use == "" {
+		c.Use = Name(obj)
+	}
+
+	for _, info := range fields(obj) {
+		fieldType := info.FieldType
+		v := info.FieldValue
+
+		name, alias := name(fieldType.Name, fieldType.Tag.Get("name"), fieldType.Tag.Get("short"))
+		usage := fieldType.Tag.Get("usage")
+		envVars := strings.Split(fieldType.Tag.Get("env"), ",")
+		defValue := fieldType.Tag.Get("default")
+		if len(envVars) == 1 && envVars[0] == "" {
+			envVars = nil
+		}
+		for _, v := range envVars {
+			if v == "" {
+				continue
+			}
+			usage += fmt.Sprintf("%s [$%s]", usage, v)
+		}
+		defInt, err := strconv.Atoi(defValue)
+		if err != nil {
+			defInt = 0
+		}
+		defValueLower := strings.ToLower(defValue)
+		defBool := defValueLower == "true" || defValueLower == "1" || defValueLower == "yes" || defValueLower == "y"
+
+		flags := c.PersistentFlags()
+		switch fieldType.Type.Kind() {
+		case reflect.Int:
+			flags.IntVarP((*int)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defInt, usage)
+			envs = append(envs, makeEnvVar(envs, name, envVars, defInt, flags, flags.GetInt)...)
+		case reflect.String:
+			flags.StringVarP((*string)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defValue, usage)
+			envs = append(envs, makeEnvVar(envs, name, envVars, defValue, flags, flags.GetString)...)
+		case reflect.Slice:
+			// env is not supported for slices
+			switch fieldType.Tag.Get("split") {
+			case "false":
+				arrays[name] = v
+				flags.StringArrayP(name, alias, nil, usage)
+			default:
+				slices[name] = v
+				flags.StringSliceP(name, alias, nil, usage)
+			}
+		case reflect.Map:
+			maps[name] = v
+			flags.StringSliceP(name, alias, nil, usage)
+		case reflect.Bool:
+			flags.BoolVarP((*bool)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defBool, usage)
+			envs = append(envs, makeEnvVar(envs, name, envVars, defBool, flags, flags.GetBool)...)
+		default:
+			panic("Unknown kind on field " + fieldType.Name + " on " + objValue.Type().Name())
+		}
+
+		if len(envVars) == 0 {
+			continue
+		}
+	}
+
+	if p, ok := obj.(PersistentPreRunnable); ok {
+		c.PersistentPreRunE = p.PersistentPre
+	}
+
+	if p, ok := obj.(PreRunnable); ok {
+		c.PreRunE = p.Pre
+	}
+
+	c.RunE = obj.Run
+	c.PersistentPreRunE = bind(c.PersistentPreRunE, arrays, slices, maps, envs)
+	c.PreRunE = bind(c.PreRunE, arrays, slices, maps, envs)
+	c.RunE = bind(c.RunE, arrays, slices, maps, envs)
+
+	cust, ok := obj.(customizer)
+	if ok {
+		cust.Customize(c)
+	}
+
+	return c
+}
+
+func assignMaps(app *cobra.Command, maps map[string]reflect.Value) error {
+	for k, v := range maps {
+		k = contextKey(k)
+		s, err := app.Flags().GetStringSlice(k)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			values := map[string]string{}
+			for _, part := range s {
+				parts := strings.SplitN(part, "=", 2)
+				if len(parts) == 1 {
+					values[parts[0]] = ""
+				} else {
+					values[parts[0]] = parts[1]
+				}
+			}
+			v.Set(reflect.ValueOf(values))
+		}
+	}
+	return nil
+}
+
+func assignSlices(app *cobra.Command, slices map[string]reflect.Value) error {
+	for k, v := range slices {
+		k = contextKey(k)
+		s, err := app.Flags().GetStringSlice(k)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			v.Set(reflect.ValueOf(s[:]))
+		}
+	}
+	return nil
+}
+
+func assignArrays(app *cobra.Command, arrays map[string]reflect.Value) error {
+	for k, v := range arrays {
+		k = contextKey(k)
+		s, err := app.Flags().GetStringArray(k)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			v.Set(reflect.ValueOf(s[:]))
+		}
+	}
+	return nil
+}
+
+func contextKey(name string) string {
+	parts := strings.Split(name, ",")
+	return parts[len(parts)-1]
+}
+
+func name(name, setName, short string) (string, string) {
+	if setName != "" {
+		return setName, short
+	}
+	parts := strings.Split(name, "_")
+	i := len(parts) - 1
+	name = caseRegexp.ReplaceAllString(parts[i], "$1-$2")
+	name = strings.ToLower(name)
+	result := append([]string{name}, parts[0:i]...)
+	for i := 0; i < len(result); i++ {
+		result[i] = strings.ToLower(result[i])
+	}
+	if short == "" && len(result) > 1 {
+		short = result[1]
+	}
+	return result[0], short
+}
+
+func bind(next func(*cobra.Command, []string) error,
+	arrays map[string]reflect.Value,
+	slices map[string]reflect.Value,
+	maps map[string]reflect.Value,
+	envs []func()) func(*cobra.Command, []string) error {
+	if next == nil {
+		return nil
+	}
+	return func(cmd *cobra.Command, args []string) error {
+		for _, envCallback := range envs {
+			envCallback()
+		}
+		if err := assignArrays(cmd, arrays); err != nil {
+			return err
+		}
+		if err := assignSlices(cmd, slices); err != nil {
+			return err
+		}
+		if err := assignMaps(cmd, maps); err != nil {
+			return err
+		}
+
+		if next != nil {
+			return next(cmd, args)
+		}
+
+		return nil
+	}
+}

cli/command_test.go

@@ -0,0 +1,49 @@
+package cli
+
+import (
+	"os"
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type testCmd struct {
+	String      string   `name:"string" short:"s" usage:"string flag" env:"STRING" default:"string"`
+	Int         int      `name:"int" short:"i" usage:"int flag" env:"INT" default:"1"`
+	Bool        bool     `name:"bool" short:"b" usage:"bool flag" env:"BOOL" default:"true"`
+	StringSlice []string `name:"string-slice" short:"S" usage:"string slice flag"`
+}
+
+func (c *testCmd) Run(cmd *cobra.Command, args []string) error {
+	return nil
+}
+
+func TestCommand(t *testing.T) {
+	var c testCmd
+	cmd := Command(&c, &cobra.Command{
+		Short: "test",
+	})
+	require.NoError(t, cmd.Execute())
+	assert.Equal(t, "string", c.String)
+	assert.Equal(t, 1, c.Int)
+	assert.Equal(t, true, c.Bool)
+	assert.Equal(t, []string{}, c.StringSlice)
+}
+
+func TestCommandEnv(t *testing.T) {
+	require.NoError(t, os.Setenv("STRING", "env-string"))
+	require.NoError(t, os.Setenv("INT", "2"))
+	require.NoError(t, os.Setenv("BOOL", "false"))
+	require.NoError(t, os.Setenv("STRING_SLICE", "env-string1,env-string2"))
+	var c testCmd
+	cmd := Command(&c, &cobra.Command{
+		Short: "test",
+	})
+	require.NoError(t, cmd.Execute())
+	assert.Equal(t, "env-string", c.String)
+	assert.Equal(t, 2, c.Int)
+	assert.Equal(t, false, c.Bool)
+	assert.Equal(t, []string{}, c.StringSlice)
+}