diff --git a/cmd/backup_yaml.go b/cmd/backup_yaml.go new file mode 100644 index 00000000..20e44927 --- /dev/null +++ b/cmd/backup_yaml.go @@ -0,0 +1,102 @@ +package cmd + +import ( + "fmt" + "strings" + "github.com/spf13/cobra" + "github.com/zilliztech/milvus-backup/core/paramtable" + "gopkg.in/yaml.v3" +) + +var configCmd = &cobra.Command{ + Use: "backup_yaml", + Short: "backup_yaml is a subcommand to check. It prints the current backup config file in yaml format to stdio.", + + Run: func(cmd *cobra.Command, args []string) { + var params paramtable.BackupParams + params.GlobalInitWithYaml(config) + params.Init() + + printParams(¶ms) + }, +} + +type YAMLConFig struct { + Log struct { + Level string `yaml:"level"` + Console bool `yaml:"console"` + File struct { + RootPath string `yaml:"rootPath"` + } + Http struct { + SimpleResponse bool `yaml:"simpleResponse"` + } `yaml:"http"` + } `yaml:"log"` + Milvus struct { + Address string `yaml:"address"` + Port int `yaml:"port"` + AuthorizationEnabled bool `yaml:"authorizationEnabled"` + TlsMode int `yaml:"tlsMode"` + User string `yaml:"user"` + Password string `yaml:"password"` + } `yaml:"milvus"` + Minio struct { + Address string `yaml:"address"` + Port int `yaml:"port"` + AccessKeyID string `yaml:"accessKeyID"` + secretAccessKey string `yaml:"secretAccessKey"` + UseSSL bool `yaml:"useSSL"` + UseIAM bool `yaml:"useIAM"` + CloudProvider string `yaml:"cloudProvider"` + IamEndpoint string `yaml:"iamEndpoint"` + BucketName string `yaml:"bucketName"` + RootPath string `yaml:"rootPath"` + BackupBucketName string `yaml:"backupBucketName"` + BackupRootPath string `yaml:"backupRootPath"` + } `yaml:"minio"` + Backup struct { + MaxSegmentGroupSize string `yaml:"maxSegmentGroupSize"` + } `yaml:"backup"` +} + +func init() { + checkCmd.AddCommand(configCmd) +} + +func printParams(base *paramtable.BackupParams) { + + yml := YAMLConFig{} + + yml.Log.Level = base.BaseTable.LoadWithDefault("log.level", "debug") + yml.Log.Console = base.ParseBool("log.console", false) + yml.Log.File.RootPath = base.LoadWithDefault("log.file.rootPath", "backup.log") + + yml.Milvus.Address = base.LoadWithDefault("milvus.address", "localhost") + yml.Milvus.Port = base.ParseIntWithDefault("milvus.port", 19530) + yml.Milvus.AuthorizationEnabled = base.ParseBool("milvus.authorizationEnabled", false) + yml.Milvus.TlsMode = base.ParseIntWithDefault("milvus.tlsMode", 0) + yml.Milvus.User = base.BaseTable.LoadWithDefault("milvus.user", "") + yml.Milvus.Password = base.BaseTable.LoadWithDefault("milvus.password", "") + + yml.Minio.Address = base.LoadWithDefault("minio.address", "localhost") + yml.Minio.Port = base.ParseIntWithDefault("minio.port", 9000) + yml.Minio.AccessKeyID = base.BaseTable.LoadWithDefault("minio.accessKeyID", "") + yml.Minio.secretAccessKey = base.BaseTable.LoadWithDefault("minio.secretAccessKey", "") + yml.Minio.UseSSL = base.ParseBool("minio.useSSL", false) + yml.Minio.UseIAM = base.ParseBool("minio.useIAM", false) + yml.Minio.CloudProvider = base.BaseTable.LoadWithDefault("minio.cloudProvider", "aws") + yml.Minio.IamEndpoint = base.BaseTable.LoadWithDefault("minio.iamEndpoint", "") + yml.Minio.BucketName = base.BaseTable.LoadWithDefault("minio.bucketName", "") + yml.Minio.RootPath = base.LoadWithDefault("minio.rootPath", "") + yml.Minio.BackupBucketName = base.LoadWithDefault("minio.backupBucketName", "") + yml.Minio.BackupRootPath = base.LoadWithDefault("minio.backupRootPath", "") + + yml.Backup.MaxSegmentGroupSize = base.LoadWithDefault("backup.maxSegmentGroupSize", "5G") + + bytes, err := yaml.Marshal(yml) + if err != nil { + panic(err) + } + + fmt.Printf("%s\n%s", strings.Repeat("-", 80), string(bytes)) +} diff --git a/cmd/root.go b/cmd/root.go index ca7f5ff8..93e4a985 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -3,6 +3,8 @@ package cmd import ( "errors" "fmt" + "os" + "strings" "github.com/spf13/cobra" @@ -10,7 +12,8 @@ import ( ) var ( - config string + config string + yamlOverrides []string ) var rootCmd = &cobra.Command{ @@ -20,10 +23,14 @@ var rootCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { Error(cmd, args, errors.New("unrecognized command")) }, + PersistentPreRun: func(cmd *cobra.Command, args []string){ + setEnvs(yamlOverrides) + }, } func Execute() { rootCmd.PersistentFlags().StringVarP(&config, "config", "", "backup.yaml", "config YAML file of milvus") + rootCmd.PersistentFlags().StringSliceVar(&yamlOverrides, "set", []string{}, "Override yaml values using a capitalized snake case format (--set MILVUS_USER=Marco)") rootCmd.CompletionOptions.DisableDefaultCmd = true rootCmd.Execute() } @@ -33,3 +40,12 @@ func SetVersionInfo(version, commit, date string) { println(rootCmd.Version) log.Info(fmt.Sprintf("Milvus backup version: %s", rootCmd.Version)) } + +// Set environment variables from yamlOverrides +func setEnvs(envs []string) { + for _, e := range envs { + env := strings.Split(e, "=") + os.Setenv(env[0], env[1]) + } + +} diff --git a/core/paramtable/base_table_test.go b/core/paramtable/base_table_test.go index 62042192..aac7cc44 100644 --- a/core/paramtable/base_table_test.go +++ b/core/paramtable/base_table_test.go @@ -1,9 +1,11 @@ package paramtable import ( + "os" + "testing" + "github.com/stretchr/testify/assert" memkv "github.com/zilliztech/milvus-backup/internal/kv/mem" - "testing" ) func TestParseDataSizeWithDefault(t *testing.T) { @@ -37,3 +39,13 @@ func TestParseDataSizeWithDefault(t *testing.T) { assert.NoError(t, err) assert.Equal(t, size1024000, int64(1024000)) } + +func TestTryLoadFromEnv(t *testing.T) { + base := &BaseTable{ + params: memkv.NewMemoryKV(), + } + os.Setenv("MILVUS_USER", "Marco") + base.tryLoadFromEnv() + + assert.Equal(t, "Marco", base.params.LoadWithDefault("milvus.user", "")) +} diff --git a/go.mod b/go.mod index 85e42a73..c9c8dd52 100644 --- a/go.mod +++ b/go.mod @@ -37,6 +37,7 @@ require ( golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 google.golang.org/grpc v1.48.0 gopkg.in/natefinch/lumberjack.v2 v2.0.0 + gopkg.in/yaml.v3 v3.0.1 ) require ( @@ -113,7 +114,6 @@ require ( google.golang.org/protobuf v1.30.0 // indirect gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect ) replace github.com/milvus-io/milvus-sdk-go/v2 => github.com/wayblink/milvus-sdk-go/v2 v2.3.6-beta3