mirror of https://github.com/linka-cloud/grpc.git
add cobra command utilities and log formatter
Signed-off-by: Adphi <philippe.adrien.nousse@gmail.com>
This commit is contained in:
parent
8fd170c0a8
commit
0b4d636ec2
GPG Key ID:
46BE4062DB2397FF
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
Loading…
Reference in New Issue