Skip to content

Commit 3a5e13f

Browse files
breezewishbreezewish
authored
playground-ng: fix arg0 and polish (#2673)
* playground-ng: render help/examples with correct invocation * Simplify Signed-off-by: Wish <[email protected]> * tui: preserve argv0 path in cobra usage * playground-ng: colorize help via coloredcobra * Small polish Signed-off-by: Wish <[email protected]> * playground-ng: Fix output after snapshot Signed-off-by: Wish <[email protected]> * Fix tidy Signed-off-by: Wish <[email protected]> --------- Signed-off-by: Wish <[email protected]>
1 parent 740a506 commit 3a5e13f

File tree

15 files changed

+321
-42
lines changed

15 files changed

+321
-42
lines changed

‎components/playground-ng/artifacts.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ func (p *Playground) printClusterInfoCallout(tidbSucc, tiproxySucc []string) boo
9494
if p.destroyDataAfterExit {
9595
blocks = append(blocks, block{
9696
colorstr.Sprintf("[dim]Cluster data will be destroyed after exit.[reset]"),
97-
colorstr.Sprintf("[dim]Use [reset][bold]--tag <name>[reset][dim] to persist data after exit.[reset]"),
97+
colorstr.Sprintf("[dim]Use [bold]--tag <name>[reset][dim] to persist data after exit.[reset]"),
9898
})
9999
}
100100

‎components/playground-ng/boot.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ func ValidateBootOptionsPure(options *BootOptions) error {
243243
case proc.ModeCSE, proc.ModeDisAgg, proc.ModeNextGen:
244244
if utils.Version(options.Version).IsValid() && !tidbver.TiFlashPlaygroundNewStartMode(options.Version) {
245245
// For simplicity, currently we only implemented disagg mode when TiFlash can run without config.
246-
return fmt.Errorf("tiup playground-ng only supports CSE/Disagg mode for TiDB cluster >= v7.1.0 (or nightly)")
246+
return fmt.Errorf("%s only supports CSE/Disagg mode for TiDB cluster >= v7.1.0 (or nightly)", playgroundCLIArg0())
247247
}
248248
}
249249

@@ -260,7 +260,7 @@ func ValidateBootOptionsPure(options *BootOptions) error {
260260

261261
// Currently we always assign region=local. Other regions are not supported.
262262
if strings.Contains(hostname, "amazonaws.com") {
263-
return fmt.Errorf("tiup playground-ng only supports local S3 (like minio); S3 on AWS regions is not supported")
263+
return fmt.Errorf("%s only supports local S3 (like minio); S3 on AWS regions is not supported", playgroundCLIArg0())
264264
}
265265
}
266266

‎components/playground-ng/cli_invocation.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"os"
6+
"strings"
7+
8+
"github.com/pingcap/tiup/pkg/localdata"
9+
)
10+
11+
const playgroundComponentName = "playground-ng"
12+
13+
func resolvePlaygroundCLIArg0(tiupVersion, userInputComponentVersion, argv0 string) string {
14+
// When running as a TiUP component, reconstruct the user-facing command line
15+
// prefix so help/examples are copy-pasteable:
16+
// tiup playground-ng[:<component-version>]
17+
//
18+
// NOTE: TIUP_USER_INPUT_VERSION is only set when users specify the component
19+
// version explicitly (e.g. `tiup playground-ng:nightly ...`). It is empty for
20+
// `tiup playground-ng ...`.
21+
if tiupVersion != "" {
22+
if userInputComponentVersion != "" {
23+
return fmt.Sprintf("tiup %s:%s", playgroundComponentName, userInputComponentVersion)
24+
}
25+
return fmt.Sprintf("tiup %s", playgroundComponentName)
26+
}
27+
28+
// Standalone mode: keep argv[0] as-is so it matches how users invoked the
29+
// binary (e.g. `bin/tiup-playground-ng`).
30+
if argv0 != "" {
31+
return argv0
32+
}
33+
return playgroundComponentName
34+
}
35+
36+
func playgroundCLIArg0() string {
37+
argv0 := ""
38+
if len(os.Args) > 0 {
39+
argv0 = os.Args[0]
40+
}
41+
return resolvePlaygroundCLIArg0(
42+
os.Getenv(localdata.EnvNameTiUPVersion),
43+
os.Getenv(localdata.EnvNameUserInputVersion),
44+
argv0,
45+
)
46+
}
47+
48+
func playgroundCLICommand(subcommand string) string {
49+
arg0 := playgroundCLIArg0()
50+
if subcommand == "" {
51+
return arg0
52+
}
53+
return arg0 + " " + subcommand
54+
}
55+
56+
func rewriteCobraUseLine(arg0, useLine string) string {
57+
if arg0 == "" {
58+
return useLine
59+
}
60+
i := strings.Index(useLine, " ")
61+
if i <= 0 {
62+
return arg0
63+
}
64+
return arg0 + useLine[i:]
65+
}
66+
67+
func rewriteCobraCommandPath(arg0, commandPath string) string {
68+
if arg0 == "" {
69+
return commandPath
70+
}
71+
i := strings.Index(commandPath, " ")
72+
if i <= 0 {
73+
return arg0
74+
}
75+
return arg0 + commandPath[i:]
76+
}

‎components/playground-ng/cli_invocation_test.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package main
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
)
8+
9+
func TestResolvePlaygroundCLIArg0(t *testing.T) {
10+
t.Run("StandaloneUsesArgv0", func(t *testing.T) {
11+
got := resolvePlaygroundCLIArg0("", "", "bin/playground-ng")
12+
require.Equal(t, "bin/playground-ng", got)
13+
})
14+
15+
t.Run("StandaloneFallsBackToComponentName", func(t *testing.T) {
16+
got := resolvePlaygroundCLIArg0("", "", "")
17+
require.Equal(t, playgroundComponentName, got)
18+
})
19+
20+
t.Run("TiUPImplicitComponentVersionOmitsColon", func(t *testing.T) {
21+
got := resolvePlaygroundCLIArg0("v1.0.0", "", "/tmp/ignored")
22+
require.Equal(t, "tiup playground-ng", got)
23+
})
24+
25+
t.Run("TiUPExplicitComponentVersionIncludesColon", func(t *testing.T) {
26+
got := resolvePlaygroundCLIArg0("v1.0.0", "nightly", "/tmp/ignored")
27+
require.Equal(t, "tiup playground-ng:nightly", got)
28+
})
29+
}
30+
31+
func TestRewriteCobraUseLine(t *testing.T) {
32+
arg0 := "tiup playground-ng:v1.2.3"
33+
34+
t.Run("RewritesRootUseLine", func(t *testing.T) {
35+
got := rewriteCobraUseLine(arg0, "tiup-playground-ng [version]")
36+
require.Equal(t, "tiup playground-ng:v1.2.3 [version]", got)
37+
})
38+
39+
t.Run("RewritesSubcommandUseLine", func(t *testing.T) {
40+
got := rewriteCobraUseLine(arg0, "tiup-playground-ng stop [flags]")
41+
require.Equal(t, "tiup playground-ng:v1.2.3 stop [flags]", got)
42+
})
43+
44+
t.Run("NoSuffix", func(t *testing.T) {
45+
got := rewriteCobraUseLine(arg0, "tiup-playground-ng")
46+
require.Equal(t, arg0, got)
47+
})
48+
}
49+
50+
func TestRewriteCobraCommandPath(t *testing.T) {
51+
arg0 := "tiup playground-ng:v1.2.3"
52+
53+
t.Run("RewritesRootCommandPath", func(t *testing.T) {
54+
got := rewriteCobraCommandPath(arg0, "tiup-playground-ng")
55+
require.Equal(t, arg0, got)
56+
})
57+
58+
t.Run("RewritesSubcommandCommandPath", func(t *testing.T) {
59+
got := rewriteCobraCommandPath(arg0, "tiup-playground-ng help")
60+
require.Equal(t, "tiup playground-ng:v1.2.3 help", got)
61+
})
62+
}

‎components/playground-ng/command.go‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,8 @@ func scaleOutServiceIDs() []proc.ServiceID {
276276
}
277277

278278
func newScaleOut(state *cliState) *cobra.Command {
279+
arg0 := playgroundCLIArg0()
280+
279281
var services []string
280282
var counts []int
281283
var cfg proc.Config
@@ -294,7 +296,7 @@ func newScaleOut(state *cliState) *cobra.Command {
294296
cmd := &cobra.Command{
295297
Use: "scale-out",
296298
Short: "Scale out instances in a running playground",
297-
Example: "tiup playground-ng scale-out --service tidb --count 1",
299+
Example: fmt.Sprintf("%s scale-out --service tidb --count 1", arg0),
298300
RunE: func(cmd *cobra.Command, args []string) error {
299301
var reqs []ScaleOutRequest
300302
switch {
@@ -367,13 +369,15 @@ func newScaleOut(state *cliState) *cobra.Command {
367369
}
368370

369371
func newScaleIn(state *cliState) *cobra.Command {
372+
arg0 := playgroundCLIArg0()
373+
370374
var names []string
371375
var pids []int
372376

373377
cmd := &cobra.Command{
374378
Use: "scale-in",
375379
Short: "Scale in one or more instances by name or pid",
376-
Example: " tiup playground-ng scale-in --name tidb-0\n tiup playground-ng scale-in --pid 12345",
380+
Example: fmt.Sprintf(" %s scale-in --name tidb-0\n %s scale-in --pid 12345", arg0, arg0),
377381
RunE: func(cmd *cobra.Command, args []string) error {
378382
if len(names) == 0 && len(pids) == 0 {
379383
return cmd.Help()
@@ -412,8 +416,8 @@ func newScaleIn(state *cliState) *cobra.Command {
412416
Hidden: false,
413417
}
414418

415-
cmd.Flags().StringSliceVar(&names, "name", nil, "Instance name(s) to scale in (get from tiup playground-ng display)")
416-
cmd.Flags().IntSliceVar(&pids, "pid", nil, "Instance PID(s) to scale in (get from tiup playground-ng display --verbose)")
419+
cmd.Flags().StringSliceVar(&names, "name", nil, fmt.Sprintf("Instance name(s) to scale in (get from %s display)", arg0))
420+
cmd.Flags().IntSliceVar(&pids, "pid", nil, fmt.Sprintf("Instance PID(s) to scale in (get from %s display --verbose)", arg0))
417421

418422
return cmd
419423
}
@@ -441,11 +445,13 @@ func newDisplay(state *cliState) *cobra.Command {
441445
}
442446

443447
func newStop(state *cliState) *cobra.Command {
448+
arg0 := playgroundCLIArg0()
449+
444450
var timeoutSec int
445451
cmd := &cobra.Command{
446452
Use: "stop",
447453
Short: "Stop a running playground",
448-
Example: "tiup playground-ng stop --tag my-cluster",
454+
Example: fmt.Sprintf("%s stop --tag my-cluster", arg0),
449455
RunE: func(cmd *cobra.Command, args []string) error {
450456
if timeoutSec <= 0 {
451457
timeoutSec = 60
@@ -567,7 +573,7 @@ func printDisplayFailureWarning(out io.Writer, err error) {
567573

568574
var lines []string
569575
if shouldSuggestPlaygroundNotRunning(err) {
570-
lines = append(lines, colorstr.Sprintf("[bold]Looks like no tiup playground-ng is running?[reset]"))
576+
lines = append(lines, colorstr.Sprintf("[bold]Looks like no %s is running?[reset]", playgroundCLIArg0()))
571577
}
572578
lines = append(lines, fmt.Sprintf("Error: %v", err))
573579

‎components/playground-ng/daemon_mode.go‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -149,17 +149,13 @@ func runBackgroundStarter(state *cliState) error {
149149
}
150150

151151
out := tuiv2output.Stdout.Get()
152-
colorstr.Fprintf(out, backgroundStarterReadyMessage(state.tag))
152+
colorstr.Fprintf(out, fmt.Sprintf("\n[dim]Cluster running in background ([bold]-d[reset][dim]).[reset]\n[dim]To stop: [bold]%s stop --tag %s[reset]\n", playgroundCLIArg0(), state.tag))
153153
return nil
154154
}
155155
}
156156
}
157157
}
158158

159-
func backgroundStarterReadyMessage(tag string) string {
160-
return fmt.Sprintf("\n[dim]Cluster running in background (-d).[reset]\n[dim]To stop: [reset]tiup playground-ng stop --tag %s\n", tag)
161-
}
162-
163159
func daemonEnv() []string {
164160
env := append([]string(nil), os.Environ()...)
165161

‎components/playground-ng/daemon_mode_test.go‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -87,10 +87,6 @@ func TestBuildDaemonArgs_FiltersEqualsFormsAndKeepsPositionals(t *testing.T) {
8787
}, got)
8888
}
8989

90-
func TestBackgroundStarterReadyMessage(t *testing.T) {
91-
require.Equal(t, "\n[dim]Cluster running in background (-d).[reset]\n[dim]To stop: [reset]tiup playground-ng stop --tag foo\n", backgroundStarterReadyMessage("foo"))
92-
}
93-
9490
func TestTailEventLog_ReplaysNewEventsAfterOffset(t *testing.T) {
9591
dir := t.TempDir()
9692
eventLogPath := filepath.Join(dir, "events.jsonl")

‎components/playground-ng/instances.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ func stopAll(out io.Writer, timeout time.Duration, state *cliState) error {
114114
return fmt.Errorf("cli state is nil")
115115
}
116116
if strings.TrimSpace(state.tag) != "" || strings.TrimSpace(state.tiupDataDir) != "" {
117-
return fmt.Errorf("stop-all does not accept --tag or TIUP_INSTANCE_DATA_DIR; use 'tiup playground-ng stop' instead")
117+
return fmt.Errorf("stop-all does not accept --tag or TIUP_INSTANCE_DATA_DIR; use '%s' instead", playgroundCLICommand("stop"))
118118
}
119119

120120
targets, err := listPlaygroundTargets(state.dataDir)

‎components/playground-ng/main.go‎

Lines changed: 63 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"time"
3232

3333
"github.com/fatih/color"
34+
cc "github.com/ivanpirog/coloredcobra"
3435
"github.com/pingcap/errors"
3536
"github.com/pingcap/tiup/components/playground-ng/proc"
3637
"github.com/pingcap/tiup/pkg/environment"
@@ -121,24 +122,33 @@ func execute(state *cliState) error {
121122
state = newCLIState()
122123
}
123124

125+
arg0 := playgroundCLIArg0()
126+
124127
rootCmd := &cobra.Command{
125128
Use: fmt.Sprintf("%s [version]", filepath.Base(os.Args[0])),
126-
Long: `Bootstrap a TiDB cluster in your local host, the latest release version will be chosen
127-
if you don't specified a version.
128-
129-
Examples:
130-
$ tiup playground-ng nightly # Start a TiDB nightly version local cluster
131-
$ tiup playground-ng v5.0.1 --db 3 --pd 3 --kv 3 # Start a local cluster with 10 nodes
132-
$ tiup playground-ng nightly --without-monitor # Start a local cluster and disable monitor system
133-
$ tiup playground-ng --pd.config ~/config/pd.toml # Start a local cluster with specified configuration file
134-
$ tiup playground-ng --db.binpath /xx/tidb-server # Start a local cluster with component binary path
135-
$ tiup playground-ng --tag xx # Start a local cluster with data dir named 'xx' and uncleaned after exit
136-
$ tiup playground-ng -d --tag xx # Start a local cluster in background (daemon mode)
137-
$ tiup playground-ng stop --tag xx # Stop the cluster started with --tag xx
138-
$ tiup playground-ng stop-all # Stop all running playground-ng instances
139-
$ tiup playground-ng ps # List all running playground-ng instances
140-
$ tiup playground-ng --mode tikv-slim # Start a local tikv only cluster (No TiDB or TiFlash Available)
141-
$ tiup playground-ng --mode tikv-slim --kv 3 --pd 3 # Start a local tikv only cluster with 6 nodes`,
129+
Long: colorstr.Sprintf(`>_ [bold]TiUP Playground[reset] [dim](ng)[reset]
130+
131+
Start and manage a TiDB cluster locally for development.
132+
133+
[bold]Examples:[reset]
134+
135+
[dim]Start a cluster using latest release version:[reset]
136+
[cyan]%[1]s[reset]
137+
138+
[dim]Start a nightly cluster:[reset]
139+
[cyan]%[1]s nightly[reset]
140+
141+
[dim]Start a TiKV-only cluster:[reset]
142+
[cyan]%[1]s --mode tikv-slim[reset]
143+
144+
[dim]Start a cluster and run in background:[reset]
145+
[cyan]%[1]s -d[reset]
146+
147+
[dim]Start a tagged cluster (data will not be cleaned after exit):[reset]
148+
[cyan]%[1]s --tag foo[reset]
149+
150+
[dim]Stop the specified cluster:[reset]
151+
[cyan]%[1]s stop --tag foo[reset]`, arg0),
142152
SilenceUsage: true,
143153
SilenceErrors: true,
144154
Version: version.NewTiUPVersion().String(),
@@ -395,10 +405,43 @@ Examples:
395405
color.NoColor = true
396406
}
397407

398-
tui.AddColorFunctionsForCobra()
399-
tui.BeautifyCobraUsageAndHelp(rootCmd)
408+
cobra.AddTemplateFunc("pgCmdLine", func(useLine string) string {
409+
return rewriteCobraUseLine(arg0, useLine)
410+
})
411+
cobra.AddTemplateFunc("pgCmdPath", func(commandPath string) string {
412+
return rewriteCobraCommandPath(arg0, commandPath)
413+
})
414+
415+
usageTpl := rootCmd.UsageTemplate()
416+
usageTpl = strings.ReplaceAll(usageTpl, "{{.UseLine}}", "{{pgCmdLine .UseLine}}")
417+
usageTpl = strings.ReplaceAll(usageTpl, "{{.CommandPath}}", "{{pgCmdPath .CommandPath}}")
418+
rootCmd.SetUsageTemplate(usageTpl)
419+
420+
cc.Init(&cc.Config{
421+
RootCmd: rootCmd,
422+
Headings: cc.Bold,
423+
Commands: cc.Cyan + cc.Bold,
424+
})
425+
426+
usageTpl = rootCmd.UsageTemplate()
427+
usageTpl = strings.ReplaceAll(usageTpl, "(CommandStyle .CommandPath)", "(CommandStyle (pgCmdPath .CommandPath))")
428+
rootCmd.SetUsageTemplate(usageTpl)
429+
430+
// Cobra's default version flag uses the command name derived from `Use`,
431+
// which is the binary name (e.g. tiup-playground-ng). For standalone runs we
432+
// want argv0 (e.g. bin/tiup-playground-ng); for TiUP component mode we want
433+
// `tiup playground-ng[:<ver>]`.
434+
rootCmd.InitDefaultHelpFlag()
435+
if f := rootCmd.Flags().Lookup("help"); f != nil {
436+
f.Usage = fmt.Sprintf("help for %s", arg0)
437+
}
438+
439+
rootCmd.InitDefaultVersionFlag()
440+
if f := rootCmd.Flags().Lookup("version"); f != nil {
441+
f.Usage = fmt.Sprintf("version for %s", arg0)
442+
}
400443

401-
rootCmd.Flags().StringVar(&state.options.ShOpt.Mode, "mode", "tidb", fmt.Sprintf("tiup playground-ng mode: '%s', '%s', '%s', '%s', '%s'", proc.ModeNormal, proc.ModeCSE, proc.ModeNextGen, proc.ModeDisAgg, proc.ModeTiKVSlim))
444+
rootCmd.Flags().StringVar(&state.options.ShOpt.Mode, "mode", "tidb", fmt.Sprintf("%s mode: '%s', '%s', '%s', '%s', '%s'", arg0, proc.ModeNormal, proc.ModeCSE, proc.ModeNextGen, proc.ModeDisAgg, proc.ModeTiKVSlim))
402445
rootCmd.Flags().StringVar(&state.options.ShOpt.PDMode, "pd.mode", "pd", "PD mode: 'pd', 'ms'")
403446
rootCmd.Flags().StringVar(&state.options.ShOpt.CSE.S3Endpoint, "cse.s3_endpoint", "http://127.0.0.1:9000",
404447
fmt.Sprintf("Object store URL for --mode=%s, --mode=%s, --mode=%s", proc.ModeCSE, proc.ModeDisAgg, proc.ModeNextGen))
@@ -930,7 +973,7 @@ var _ repository.DownloadProgress = (*repoDownloadProgress)(nil)
930973
var _ repository.DownloadProgressReporter = (*repoDownloadProgress)(nil)
931974

932975
func main() {
933-
tui.RegisterArg0("tiup playground-ng")
976+
tui.RegisterArg0(playgroundCLIArg0())
934977

935978
state := newCLIState()
936979

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ require (
2626
github.com/gogo/protobuf v1.3.2
2727
github.com/google/uuid v1.6.0
2828
github.com/gorilla/mux v1.8.0
29+
github.com/ivanpirog/coloredcobra v1.0.1
2930
github.com/jedib0t/go-pretty/v6 v6.6.7
3031
github.com/jeremywohl/flatten v1.0.1
3132
github.com/joomcode/errorx v1.1.0

0 commit comments

Comments
 (0)