-
Notifications
You must be signed in to change notification settings - Fork 88
/
root.go
159 lines (139 loc) · 5.07 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
151
152
153
154
155
156
157
158
159
package cli
import (
"context"
"fmt"
"os"
"strings"
"github.com/spf13/cobra"
"github.com/spf13/viper"
"go.opentelemetry.io/otel/trace"
"github.com/bacalhau-project/bacalhau/cmd/cli/agent"
configcli "github.com/bacalhau-project/bacalhau/cmd/cli/config"
"github.com/bacalhau-project/bacalhau/cmd/cli/deprecated"
"github.com/bacalhau-project/bacalhau/cmd/cli/devstack"
"github.com/bacalhau-project/bacalhau/cmd/cli/docker"
"github.com/bacalhau-project/bacalhau/cmd/cli/exec"
"github.com/bacalhau-project/bacalhau/cmd/cli/job"
"github.com/bacalhau-project/bacalhau/cmd/cli/node"
"github.com/bacalhau-project/bacalhau/cmd/cli/serve"
"github.com/bacalhau-project/bacalhau/cmd/cli/version"
"github.com/bacalhau-project/bacalhau/cmd/cli/wasm"
"github.com/bacalhau-project/bacalhau/cmd/util"
"github.com/bacalhau-project/bacalhau/cmd/util/flags/cliflags"
"github.com/bacalhau-project/bacalhau/cmd/util/flags/configflags"
"github.com/bacalhau-project/bacalhau/pkg/config/types"
"github.com/bacalhau-project/bacalhau/pkg/logger"
"github.com/bacalhau-project/bacalhau/pkg/system"
"github.com/bacalhau-project/bacalhau/pkg/telemetry"
)
func NewRootCmd() *cobra.Command {
// the root command: `bacalhau`
RootCmd := &cobra.Command{
Use: os.Args[0],
Short: "Compute over data",
Long: `Compute over data`,
}
RootCmd.PersistentFlags().VarP(cliflags.NewConfigFlag(), "config", "c", "config file(s) or dot separated path(s) to config values")
if err := RootCmd.RegisterFlagCompletionFunc("config", cliflags.ConfigAutoComplete); err != nil {
util.Fatal(RootCmd, err, 1)
}
// flag definitions with a corresponding field in the config file.
// when these flags are provided their value will be used instead of the value present in the config file.
// If no flg is provided, and the config file doesn't have a value defined then the default value will be used.
rootFlags := map[string][]configflags.Definition{
"api": configflags.ClientAPIFlags,
"logging": configflags.LogFlags,
"repo": configflags.DataDirFlag,
}
// register the flags on the command.
if err := configflags.RegisterFlags(RootCmd, rootFlags); err != nil {
// a failure here indicates a developer error, abort.
util.Fatal(RootCmd, err, 1)
}
// logic that must run before any child command executes
RootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
// context of the command
ctx := cmd.Context()
ctx = util.InjectCleanupManager(ctx)
ctx = injectRootSpan(cmd, ctx)
cmd.SetContext(ctx)
// Binds flags with a corresponding config file value to the root command.
// This is done in the pre-run so their values are set in the Run function.
// Cobra doesn't allow the root pre run method to run if a child command also has a prerun defined.
if err := configflags.BindFlags(viper.GetViper(), rootFlags); err != nil {
return err
}
// Configure logging
// While we allow users to configure logging via the config file, they are applied
// and will override this configuration at a later stage when the config is loaded.
// This is needed to ensure any logs before the config is loaded are captured.
logMode := viper.GetString(types.LoggingModeKey)
if logMode == "" {
logMode = string(logger.LogModeDefault)
}
logLevel := viper.GetString(types.LoggingLevelKey)
if logLevel == "" {
logLevel = "Info"
}
if err := logger.ConfigureLogging(logMode, logLevel); err != nil {
return fmt.Errorf("failed to configure logging: %w", err)
}
return nil
}
// logic that must run after any child command completes.
RootCmd.PersistentPostRunE = func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
ctx.Value(spanKey).(trace.Span).End()
ctx.Value(util.SystemManagerKey).(*system.CleanupManager).Cleanup(ctx)
return nil
}
// register child commands.
RootCmd.AddCommand(
agent.NewCmd(),
configcli.NewCmd(),
devstack.NewCmd(),
docker.NewCmd(),
exec.NewCmd(),
job.NewCmd(),
node.NewCmd(),
serve.NewCmd(),
version.NewCmd(),
wasm.NewCmd(),
// deprecated command
deprecated.NewCancelCmd(),
deprecated.NewCreateCmd(),
deprecated.NewDescribeCmd(),
deprecated.NewGetCmd(),
deprecated.NewIDCmd(),
deprecated.NewListCmd(),
deprecated.NewLogsCmd(),
deprecated.NewValidateCmd(),
)
return RootCmd
}
func Execute(ctx context.Context) {
rootCmd := NewRootCmd()
rootCmd.SetContext(ctx)
rootCmd.SilenceErrors = true
rootCmd.SilenceUsage = true
// this is needed as cobra defaults to Stderr if no output is set.
rootCmd.SetOut(os.Stdout)
rootCmd.SetErr(os.Stderr)
if err := rootCmd.Execute(); err != nil {
util.Fatal(rootCmd, err, 1)
}
}
type contextKey struct {
name string
}
var spanKey = contextKey{name: "context key for storing the root span"}
func injectRootSpan(cmd *cobra.Command, ctx context.Context) context.Context {
var names []string
root := cmd
for ; root.HasParent(); root = root.Parent() {
names = append([]string{root.Name()}, names...)
}
name := fmt.Sprintf("bacalhau.%s", strings.Join(names, "."))
ctx, span := telemetry.NewRootSpan(ctx, telemetry.GetTracer(), name)
return context.WithValue(ctx, spanKey, span)
}