-
Notifications
You must be signed in to change notification settings - Fork 9
/
do.go
140 lines (114 loc) · 3.73 KB
/
do.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
// Copyright 2022 Namespace Labs Inc; 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.
package fncobra
import (
"context"
"fmt"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"namespacelabs.dev/foundation/internal/cli/nsboot"
"namespacelabs.dev/foundation/internal/cli/versioncheck"
"namespacelabs.dev/foundation/internal/compute"
"namespacelabs.dev/foundation/internal/console"
"namespacelabs.dev/foundation/std/tasks"
)
type CmdHandler func(context.Context, []string) error
type ArgsParser interface {
AddFlags(*cobra.Command)
Parse(ctx context.Context, args []string) error
}
func RunE(handler CmdHandler) func(*cobra.Command, []string) error {
return func(cmd *cobra.Command, args []string) error {
return handler(cmd.Context(), args)
}
}
func CheckVersion(ctx context.Context) func() {
// Navigate with care: within the detach below we don't use the graph
// context, but rather the incoming context. The reason for this is that
// we need to have the ability to cancel the version check. THIS IS OK
// BUT IT RELIES ON INTERNAL DETAILS! compute.Do invokes the callback in
// the same executor where Detach() runs. So we're guaranteed to observe
// cancelation here.
ctx, cancel := context.WithCancel(ctx)
compute.On(ctx).Detach(tasks.Action("ns.check-updated"), func(_ context.Context) error {
status, err := versioncheck.CheckRemote(ctx, nil)
if err != nil {
fmt.Fprintf(console.Debug(ctx), "failed to check remote version: %v\n", err)
return nil
}
if status == nil {
return nil
}
if newVersion, ok := nsboot.CheckInvalidate(status); ok {
compute.On(ctx).Cleanup(tasks.Action("ns.check-updated.notify").LogLevel(1), func(ctx context.Context) error {
fmt.Fprintf(console.Stdout(ctx), "\n\n A new version of ns is available (%s). It will be automatically downloaded on the next run.\n\n", newVersion)
return nil
})
}
return nil
})
return cancel
}
func RunInContext(ctx context.Context, handler func(context.Context) error, starter func(context.Context) func()) error {
ctx, cancel := WithSigIntCancel(ctx)
defer cancel()
return compute.Do(ctx, func(ctx context.Context) error {
if starter != nil {
cancel := starter(ctx)
if cancel != nil {
defer cancel()
}
}
return handler(ctx)
})
}
type CommandCtrl struct {
cmd *cobra.Command
argParsers []ArgsParser
}
func With(cmd *cobra.Command, handler func(context.Context) error) *cobra.Command {
cmd.RunE = RunE(func(ctx context.Context, _ []string) error {
return handler(ctx)
})
return cmd
}
func Cmd(cmd *cobra.Command) *CommandCtrl {
return &CommandCtrl{
cmd: cmd,
argParsers: []ArgsParser{},
}
}
func (c *CommandCtrl) With(argParser ...ArgsParser) *CommandCtrl {
c.argParsers = append(c.argParsers, argParser...)
return c
}
func (c *CommandCtrl) WithFlags(f func(flags *pflag.FlagSet)) *CommandCtrl {
return c.With(&simpleFlagParser{f})
}
func (c *CommandCtrl) Do(handler func(context.Context) error) *cobra.Command {
return c.DoWithArgs(func(ctx context.Context, args []string) error {
return handler(ctx)
})
}
func (c *CommandCtrl) DoWithArgs(handler CmdHandler) *cobra.Command {
for _, parser := range c.argParsers {
parser.AddFlags(c.cmd)
}
c.cmd.RunE = RunE(func(ctx context.Context, args []string) error {
for _, parser := range c.argParsers {
if err := parser.Parse(ctx, args); err != nil {
return err
}
}
return handler(ctx, args)
})
return c.cmd
}
type simpleFlagParser struct {
f func(flags *pflag.FlagSet)
}
func (p *simpleFlagParser) AddFlags(cmd *cobra.Command) {
p.f(cmd.Flags())
}
func (p *simpleFlagParser) Parse(ctx context.Context, args []string) error { return nil }