This repository has been archived by the owner on Feb 21, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 231
/
root.go
150 lines (133 loc) · 4.63 KB
/
root.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
// Copyright 2017 Pilosa Corp.
//
// 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 cmd
import (
"fmt"
"io"
"strings"
"github.com/pilosa/pilosa"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/spf13/viper"
)
// TODO maybe give this an Add method which will ensure two command
// with same name aren't added
var subcommandFns = map[string]func(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command{}
func NewRootCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command {
rc := &cobra.Command{
Use: "pilosa",
Short: "Pilosa - A Distributed In-memory Binary Bitmap Index.",
// TODO - is documentation actually there?
Long: `Pilosa is a fast index to turbocharge your database.
This binary contains Pilosa itself, as well as common
tools for administering pilosa, importing/exporting data,
backing up, and more. Complete documentation is available
at https://www.pilosa.com/docs/
Version: ` + pilosa.Version + `
Build Time: ` + pilosa.BuildTime + "\n",
PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
v := viper.New()
err := setAllConfig(v, cmd.Flags(), "PILOSA")
if err != nil {
return err
}
// return "dry run" error if "dry-run" flag is set
if ret, err := cmd.Flags().GetBool("dry-run"); ret && err == nil {
if cmd.Parent() != nil {
return fmt.Errorf("dry run")
} else if err != nil {
return fmt.Errorf("problem getting dry-run flag: %v", err)
}
}
return nil
},
}
rc.PersistentFlags().Bool("dry-run", false, "stop before executing")
_ = rc.PersistentFlags().MarkHidden("dry-run")
rc.PersistentFlags().StringP("config", "c", "", "Configuration file to read from.")
for _, subcomFn := range subcommandFns {
rc.AddCommand(subcomFn(stdin, stdout, stderr))
}
rc.SetOutput(stderr)
return rc
}
// setAllConfig takes a FlagSet to be the definition of all configuration
// options, as well as their defaults. It then reads from the command line, the
// environment, and a config file (if specified), and applies the configuration
// in that priority order. Since each flag in the set contains a pointer to
// where its value should be stored, setAllConfig can directly modify the value
// of each config variable.
//
// setAllConfig looks for environment variables which are capitalized versions
// of the flag names with dashes replaced by underscores, and prefixed with
// envPrefix plus an underscore.
func setAllConfig(v *viper.Viper, flags *pflag.FlagSet, envPrefix string) error {
// add cmd line flag def to viper
err := v.BindPFlags(flags)
if err != nil {
return err
}
// add env to viper
v.SetEnvPrefix(envPrefix)
v.SetEnvKeyReplacer(strings.NewReplacer("-", "_", ".", "_"))
v.AutomaticEnv()
c := v.GetString("config")
var flagErr error
validTags := make(map[string]bool)
flags.VisitAll(func(f *pflag.Flag) {
validTags[f.Name] = true
})
// add config file to viper
if c != "" {
v.SetConfigFile(c)
v.SetConfigType("toml")
err := v.ReadInConfig()
if err != nil {
return fmt.Errorf("error reading configuration file '%s': %v", c, err)
}
for _, key := range v.AllKeys() {
if _, ok := validTags[key]; !ok {
return fmt.Errorf("invalid option in configuration file: %v", key)
}
}
}
// set all values from viper
flags.VisitAll(func(f *pflag.Flag) {
if flagErr != nil {
return
}
var value string
if f.Value.Type() == "stringSlice" {
// special handling is needed for stringSlice as v.GetString will
// always return "" in the case that the value is an actual string
// slice from a config file rather than a comma separated string
// from a flag or env var.
vss := v.GetStringSlice(f.Name)
value = strings.Join(vss, ",")
} else {
value = v.GetString(f.Name)
}
if f.Changed {
// If f.Changed is true, that means the value has already been set
// by a flag, and we don't need to ask viper for it since the flag
// is the highest priority. This works around a problem with string
// slices where f.Value.Set(csvString) would cause the elements of
// csvString to be appended to the existing value rather than
// replacing it.
return
}
flagErr = f.Value.Set(value)
})
return flagErr
}