-
Notifications
You must be signed in to change notification settings - Fork 8
/
cmd_hooks.go
224 lines (190 loc) · 7.54 KB
/
cmd_hooks.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
package cmdconfig
import (
"context"
"fmt"
"log/slog"
"os"
"runtime/debug"
"strings"
"time"
"github.com/mattn/go-isatty"
"github.com/spf13/cobra"
"github.com/spf13/viper"
"github.com/turbot/go-kit/helpers"
"github.com/turbot/pipe-fittings/app_specific"
"github.com/turbot/pipe-fittings/cloud"
"github.com/turbot/pipe-fittings/cmdconfig"
"github.com/turbot/pipe-fittings/constants"
"github.com/turbot/pipe-fittings/error_helpers"
"github.com/turbot/pipe-fittings/filepaths"
"github.com/turbot/pipe-fittings/modconfig"
"github.com/turbot/pipe-fittings/steampipeconfig"
"github.com/turbot/pipe-fittings/task"
"github.com/turbot/pipe-fittings/utils"
"github.com/turbot/powerpipe/internal/logger"
"github.com/turbot/steampipe-plugin-sdk/v5/plugin"
"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
)
var waitForTasksChannel chan struct{}
var tasksCancelFn context.CancelFunc
// postRunHook is a function that is executed after the PostRun of every command handler
func postRunHook(_ *cobra.Command, _ []string) error {
utils.LogTime("cmdhook.postRunHook start")
defer utils.LogTime("cmdhook.postRunHook end")
if waitForTasksChannel != nil {
// wait for the async tasks to finish
select {
case <-time.After(100 * time.Millisecond):
tasksCancelFn()
return nil
case <-waitForTasksChannel:
return nil
}
}
return nil
}
// postRunHook is a function that is executed before the PreRun of every command handler
func preRunHook(cmd *cobra.Command, args []string) error {
utils.LogTime("cmdhook.preRunHook start")
defer utils.LogTime("cmdhook.preRunHook end")
viper.Set(constants.ConfigKeyActiveCommand, cmd)
viper.Set(constants.ConfigKeyActiveCommandArgs, args)
viper.Set(constants.ConfigKeyIsTerminalTTY, isatty.IsTerminal(os.Stdout.Fd()))
// set up the global viper config with default values from
// config files and ENV variables
ew := initGlobalConfig()
// display any warnings
ew.ShowWarnings()
// check for error
error_helpers.FailOnError(ew.Error)
logger.Initialize()
// runScheduledTasks skips running tasks if this instance is the plugin manager
waitForTasksChannel = runScheduledTasks(cmd.Context(), cmd, args)
// set the max memory if specified
setMemoryLimit()
return nil
}
func setMemoryLimit() {
maxMemoryBytes := viper.GetInt64(constants.ArgMemoryMaxMb) * 1024 * 1024
if maxMemoryBytes > 0 {
// set the max memory
debug.SetMemoryLimit(maxMemoryBytes)
}
}
// runScheduledTasks runs the task runner and returns a channel which is closed when
// task run is complete
//
// runScheduledTasks skips running tasks if this instance is the plugin manager
func runScheduledTasks(ctx context.Context, cmd *cobra.Command, args []string) chan struct{} {
updateCheck := viper.GetBool(constants.ArgUpdateCheck)
// for now the only scheduled task we support is update check so if that is disabled, do nothing
if !updateCheck {
return nil
}
taskUpdateCtx, cancelFn := context.WithCancel(ctx)
tasksCancelFn = cancelFn
return task.RunTasks(
taskUpdateCtx,
cmd,
args,
// pass the config value in rather than runRasks querying viper directly - to avoid concurrent map access issues
// (we can use the update-check viper config here, since initGlobalConfig has already set it up
// with values from the config files and ENV settings - update-check cannot be set from the command line)
task.WithUpdateCheck(updateCheck),
)
}
// initConfig reads in config file and ENV variables if set.
func initGlobalConfig() error_helpers.ErrorAndWarnings {
utils.LogTime("cmdconfig.initGlobalConfig start")
defer utils.LogTime("cmdconfig.initGlobalConfig end")
// load workspace profile from the configured install dir
loader, err := cmdconfig.GetWorkspaceProfileLoader[*modconfig.PowerpipeWorkspaceProfile]()
if err != nil {
return error_helpers.NewErrorsAndWarning(err)
}
var cmd = viper.Get(constants.ConfigKeyActiveCommand).(*cobra.Command)
// set-up viper with defaults from the env and default workspace profile
cmdconfig.BootstrapViper(loader, cmd,
cmdconfig.WithConfigDefaults(configDefaults(cmd)),
cmdconfig.WithDirectoryEnvMappings(dirEnvMappings()))
if err != nil {
return error_helpers.NewErrorsAndWarning(err)
}
// set global containing the configured install dir (create directory if needed)
ensureInstallDirs()
// set the rest of the defaults from ENV
// ENV takes precedence over any default configuration
cmdconfig.SetDefaultsFromEnv(envMappings())
// if an explicit workspace profile was set, add to viper as highest precedence default
// NOTE: if install_dir/mod_location are set these will already have been passed to viper by BootstrapViper
// since the "ConfiguredProfile" is passed in through a cmdline flag, it will always take precedence
if loader.ConfiguredProfile != nil {
cmdconfig.SetDefaultsFromConfig(loader.ConfiguredProfile.ConfigMap(cmd))
}
// now env vars have been processed, set filepaths.PipesInstallDir
filepaths.PipesInstallDir = viper.GetString(constants.ArgPipesInstallDir)
// NOTE: we need to resolve the token separately
// - that is because we need the resolved value of ArgPipesHost in order to load any saved token
// and we cannot get this until the other config has been resolved
err = setPipesTokenDefault(loader)
if err != nil {
return error_helpers.NewErrorsAndWarning(err)
}
// now validate all config values have appropriate values
return validateConfig()
}
func setPipesTokenDefault(loader *steampipeconfig.WorkspaceProfileLoader[*modconfig.PowerpipeWorkspaceProfile]) error {
/*
saved cloud token
pipes_token in default workspace
explicit env var (PIPES_TOKEN ) wins over
pipes_token in specific workspace
*/
// set viper defaults in order of increasing precedence
// 1) saved cloud token
savedToken, err := cloud.LoadToken()
if err != nil {
return err
}
if savedToken != "" {
viper.SetDefault(constants.ArgPipesToken, savedToken)
}
// 2) default profile cloud token
if loader.DefaultProfile.PipesToken != nil {
viper.SetDefault(constants.ArgPipesToken, *loader.DefaultProfile.PipesToken)
}
// 3) env var (PIPES_TOKEN )
cmdconfig.SetDefaultFromEnv(constants.EnvPipesToken, constants.ArgPipesToken, cmdconfig.EnvVarTypeString)
// 4) explicit workspace profile
if p := loader.ConfiguredProfile; p != nil && p.PipesToken != nil {
viper.SetDefault(constants.ArgPipesToken, *p.PipesToken)
}
return nil
}
// now validate config values have appropriate values
// (currently validates telemetry)
func validateConfig() error_helpers.ErrorAndWarnings {
var res = error_helpers.ErrorAndWarnings{}
telemetry := viper.GetString(constants.ArgTelemetry)
if !helpers.StringSliceContains(constants.TelemetryLevels, telemetry) {
res.Error = sperr.New(`invalid value of 'telemetry' (%s), must be one of: %s`, telemetry, strings.Join(constants.TelemetryLevels, ", "))
return res
}
if _, legacyDiagnosticsSet := os.LookupEnv(plugin.EnvLegacyDiagnosticsLevel); legacyDiagnosticsSet {
res.AddWarning(fmt.Sprintf("Environment variable %s is deprecated - use %s", plugin.EnvLegacyDiagnosticsLevel, plugin.EnvDiagnosticsLevel))
}
res.Error = plugin.ValidateDiagnosticsEnvVar()
return res
}
// create ~/.steampipe if needed
func ensureInstallDirs() {
installDir := viper.GetString(constants.ArgInstallDir)
slog.Debug("ensureInstallDir", "installDir", installDir)
if _, err := os.Stat(installDir); os.IsNotExist(err) {
slog.Debug("creating install dir")
err = os.MkdirAll(installDir, 0755)
error_helpers.FailOnErrorWithMessage(err, fmt.Sprintf("could not create installation directory: %s", installDir))
}
// store as app_specific.InstallDir
app_specific.InstallDir = installDir
}