cli: SetupRootCommand, SetupPluginRootCommand: set context on root-co…

…mmands Add a context-argument to; - cli.SetupRootCommand - cli.SetupPluginRootCommand - cli-plugins/plugin.RunPlugin - cli-plugins/plugin.Run Signed-off-by: Sebastiaan van Stijn <[email protected]>
docker · Feb 21, 2024 · a4f985c · a4f985c
1 parent c3dfbd5
commit a4f985c
Show file tree

Hide file tree

Showing 9 changed files with 47 additions and 37 deletions.
diff --git a/cli-plugins/examples/helloworld/main.go b/cli-plugins/examples/helloworld/main.go
@@ -12,7 +12,8 @@ import (
 )
 
 func main() {
-	plugin.Run(func(dockerCli command.Cli) *cobra.Command {
+	ctx := context.Background()
+	plugin.Run(ctx, func(dockerCli command.Cli) *cobra.Command {
 		goodbye := &cobra.Command{
 			Use:   "goodbye",
 			Short: "Say Goodbye instead of Hello",

diff --git a/cli-plugins/plugin/plugin.go b/cli-plugins/plugin/plugin.go
@@ -27,21 +27,15 @@ import (
 var PersistentPreRunE func(*cobra.Command, []string) error
 
 // RunPlugin executes the specified plugin command
-func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager.Metadata) error {
-	tcmd := newPluginCommand(dockerCli, plugin, meta)
+func RunPlugin(ctx context.Context, dockerCli *command.DockerCli, plugin *cobra.Command, meta manager.Metadata) error {
+	tcmd := newPluginCommand(ctx, dockerCli, plugin, meta)
 
 	var persistentPreRunOnce sync.Once
 	PersistentPreRunE = func(cmd *cobra.Command, _ []string) error {
 		var err error
 		persistentPreRunOnce.Do(func() {
-			cmdContext := cmd.Context()
-			// TODO: revisit and make sure this check makes sense
-			// see: https://github.com/docker/cli/pull/4599#discussion_r1422487271
-			if cmdContext == nil {
-				cmdContext = context.TODO()
-			}
-			ctx, cancel := context.WithCancel(cmdContext)
-			cmd.SetContext(ctx)
+			ctx2, cancel := context.WithCancel(cmd.Context())
+			cmd.SetContext(ctx2)
 			// Set up the context to cancel based on signalling via CLI socket.
 			socket.ConnectAndWait(cancel)
 
@@ -65,16 +59,16 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager
 }
 
 // Run is the top-level entry point to the CLI plugin framework. It should be called from your plugin's `main()` function.
-func Run(makeCmd func(command.Cli) *cobra.Command, meta manager.Metadata) {
-	dockerCli, err := command.NewDockerCli()
+func Run(ctx context.Context, makeCmd func(command.Cli) *cobra.Command, meta manager.Metadata) {
+	dockerCli, err := command.NewDockerCli(command.WithBaseContext(ctx))
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}
 
 	plugin := makeCmd(dockerCli)
 
-	if err := RunPlugin(dockerCli, plugin, meta); err != nil {
+	if err := RunPlugin(ctx, dockerCli, plugin, meta); err != nil {
 		if sterr, ok := err.(cli.StatusError); ok {
 			if sterr.Status != "" {
 				fmt.Fprintln(dockerCli.Err(), sterr.Status)
@@ -123,7 +117,7 @@ func withPluginClientConn(name string) command.CLIOption {
 	})
 }
 
-func newPluginCommand(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager.Metadata) *cli.TopLevelCommand {
+func newPluginCommand(ctx context.Context, dockerCli *command.DockerCli, plugin *cobra.Command, meta manager.Metadata) *cli.TopLevelCommand {
 	name := plugin.Name()
 	fullname := manager.NamePrefix + name
 
@@ -144,7 +138,7 @@ func newPluginCommand(dockerCli *command.DockerCli, plugin *cobra.Command, meta
 			DisableDescriptions: true,
 		},
 	}
-	opts, _ := cli.SetupPluginRootCommand(cmd)
+	opts, _ := cli.SetupPluginRootCommand(ctx, cmd)
 
 	cmd.SetIn(dockerCli.In())
 	cmd.SetOut(dockerCli.Out())

diff --git a/cli/cobra.go b/cli/cobra.go
@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -22,7 +23,8 @@ import (
 
 // setupCommonRootCommand contains the setup common to
 // SetupRootCommand and SetupPluginRootCommand.
-func setupCommonRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *cobra.Command) {
+func setupCommonRootCommand(ctx context.Context, rootCmd *cobra.Command) (*cliflags.ClientOptions, *cobra.Command) {
+	rootCmd.SetContext(ctx)
 	opts := cliflags.NewClientOptions()
 	opts.InstallFlags(rootCmd.Flags())
 
@@ -74,14 +76,14 @@ func setupCommonRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *c
 
 // SetupRootCommand sets default usage, help, and error handling for the
 // root command.
-func SetupRootCommand(rootCmd *cobra.Command) (opts *cliflags.ClientOptions, helpCmd *cobra.Command) {
+func SetupRootCommand(ctx context.Context, rootCmd *cobra.Command) (opts *cliflags.ClientOptions, helpCmd *cobra.Command) {
 	rootCmd.SetVersionTemplate("Docker version {{.Version}}\n")
-	return setupCommonRootCommand(rootCmd)
+	return setupCommonRootCommand(ctx, rootCmd)
 }
 
 // SetupPluginRootCommand sets default usage, help and error handling for a plugin root command.
-func SetupPluginRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet) {
-	opts, _ := setupCommonRootCommand(rootCmd)
+func SetupPluginRootCommand(ctx context.Context, rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet) {
+	opts, _ := setupCommonRootCommand(ctx, rootCmd)
 	return opts, rootCmd.Flags()
 }
 

diff --git a/cmd/docker/builder_test.go b/cmd/docker/builder_test.go
@@ -98,7 +98,8 @@ echo '{"SchemaVersion":"0.1.0","Vendor":"Docker Inc.","Version":"v0.6.3","ShortD
 				dockerCli.ConfigFile().Aliases = map[string]string{"builder": "buildx"}
 			}
 
-			tcmd := newDockerCommand(dockerCli)
+			ctx := context.TODO()
+			tcmd := newDockerCommand(ctx, dockerCli)
 			tcmd.SetArgs([]string{"build", "."})
 
 			cmd, args, err := tcmd.HandleGlobalFlags()
@@ -144,7 +145,8 @@ func TestBuildkitDisabled(t *testing.T) {
 	assert.NilError(t, dockerCli.Initialize(flags.NewClientOptions()))
 	dockerCli.ConfigFile().CLIPluginsExtraDirs = []string{dir.Path()}
 
-	tcmd := newDockerCommand(dockerCli)
+	ctx := context.TODO()
+	tcmd := newDockerCommand(ctx, dockerCli)
 	tcmd.SetArgs([]string{"build", "."})
 
 	cmd, args, err := tcmd.HandleGlobalFlags()
@@ -179,7 +181,8 @@ func TestBuilderBroken(t *testing.T) {
 	assert.NilError(t, dockerCli.Initialize(flags.NewClientOptions()))
 	dockerCli.ConfigFile().CLIPluginsExtraDirs = []string{dir.Path()}
 
-	tcmd := newDockerCommand(dockerCli)
+	ctx := context.TODO()
+	tcmd := newDockerCommand(ctx, dockerCli)
 	tcmd.SetArgs([]string{"build", "."})
 
 	cmd, args, err := tcmd.HandleGlobalFlags()
@@ -216,7 +219,8 @@ func TestBuilderBrokenEnforced(t *testing.T) {
 	assert.NilError(t, dockerCli.Initialize(flags.NewClientOptions()))
 	dockerCli.ConfigFile().CLIPluginsExtraDirs = []string{dir.Path()}
 
-	tcmd := newDockerCommand(dockerCli)
+	ctx := context.TODO()
+	tcmd := newDockerCommand(ctx, dockerCli)
 	tcmd.SetArgs([]string{"build", "."})
 
 	cmd, args, err := tcmd.HandleGlobalFlags()

diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go
@@ -35,7 +35,7 @@ func main() {
 	}
 	logrus.SetOutput(dockerCli.Err())
 
-	if err := runDocker(dockerCli); err != nil {
+	if err := runDocker(ctx, dockerCli); err != nil {
 		if sterr, ok := err.(cli.StatusError); ok {
 			if sterr.Status != "" {
 				fmt.Fprintln(dockerCli.Err(), sterr.Status)
@@ -52,7 +52,7 @@ func main() {
 	}
 }
 
-func newDockerCommand(dockerCli *command.DockerCli) *cli.TopLevelCommand {
+func newDockerCommand(ctx context.Context, dockerCli *command.DockerCli) *cli.TopLevelCommand {
 	var (
 		opts    *cliflags.ClientOptions
 		helpCmd *cobra.Command
@@ -85,7 +85,7 @@ func newDockerCommand(dockerCli *command.DockerCli) *cli.TopLevelCommand {
 	cmd.SetOut(dockerCli.Out())
 	cmd.SetErr(dockerCli.Err())
 
-	opts, helpCmd = cli.SetupRootCommand(cmd)
+	opts, helpCmd = cli.SetupRootCommand(ctx, cmd)
 	_ = registerCompletionFuncForGlobalFlags(dockerCli.ContextStore(), cmd)
 	cmd.Flags().BoolP("version", "v", false, "Print version information and quit")
 	setFlagErrorFunc(dockerCli, cmd)
@@ -281,8 +281,8 @@ func tryPluginRun(dockerCli command.Cli, cmd *cobra.Command, subcommand string,
 	return nil
 }
 
-func runDocker(dockerCli *command.DockerCli) error {
-	tcmd := newDockerCommand(dockerCli)
+func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
+	tcmd := newDockerCommand(ctx, dockerCli)
 
 	cmd, args, err := tcmd.HandleGlobalFlags()
 	if err != nil {

diff --git a/cmd/docker/docker_test.go b/cmd/docker/docker_test.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"os"
 	"testing"
@@ -18,7 +19,8 @@ func TestClientDebugEnabled(t *testing.T) {
 
 	cli, err := command.NewDockerCli()
 	assert.NilError(t, err)
-	tcmd := newDockerCommand(cli)
+	ctx := context.TODO()
+	tcmd := newDockerCommand(ctx, cli)
 	tcmd.SetFlag("debug", "true")
 	cmd, _, err := tcmd.HandleGlobalFlags()
 	assert.NilError(t, err)
@@ -41,7 +43,8 @@ func runCliCommand(t *testing.T, r io.ReadCloser, w io.Writer, args ...string) e
 	}
 	cli, err := command.NewDockerCli(command.WithInputStream(r), command.WithCombinedStreams(w))
 	assert.NilError(t, err)
-	tcmd := newDockerCommand(cli)
+	ctx := context.TODO()
+	tcmd := newDockerCommand(ctx, cli)
 
 	tcmd.SetArgs(args)
 	cmd, _, err := tcmd.HandleGlobalFlags()

diff --git a/docs/generate/generate.go b/docs/generate/generate.go
@@ -6,6 +6,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"os"
 
@@ -37,7 +38,9 @@ func gen(opts *options) error {
 		Use:   "docker [OPTIONS] COMMAND [ARG...]",
 		Short: "The base command for the Docker CLI.",
 	}
-	clientOpts, _ := cli.SetupRootCommand(cmd)
+
+	ctx := context.Background()
+	clientOpts, _ := cli.SetupRootCommand(ctx, cmd)
 	if err := dockerCLI.Initialize(clientOpts); err != nil {
 		return err
 	}

diff --git a/e2e/cli-plugins/plugins/nopersistentprerun/main.go b/e2e/cli-plugins/plugins/nopersistentprerun/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/docker/cli/cli-plugins/manager"
@@ -10,8 +11,9 @@ import (
 )
 
 func main() {
-	plugin.Run(func(dockerCli command.Cli) *cobra.Command {
-		cmd := &cobra.Command{
+	ctx := context.Background()
+	plugin.Run(ctx, func(dockerCli command.Cli) *cobra.Command {
+		return &cobra.Command{
 			Use:   "nopersistentprerun",
 			Short: "Testing without PersistentPreRun hooks",
 			// PersistentPreRunE: Not specified, we need to test that it works in the absence of an explicit call
@@ -25,7 +27,6 @@ func main() {
 				return nil
 			},
 		}
-		return cmd
 	},
 		manager.Metadata{
 			SchemaVersion: "0.1.0",

diff --git a/e2e/cli-plugins/plugins/presocket/main.go b/e2e/cli-plugins/plugins/presocket/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
@@ -14,7 +15,8 @@ import (
 )
 
 func main() {
-	plugin.Run(RootCmd, manager.Metadata{
+	ctx := context.Background()
+	plugin.Run(ctx, RootCmd, manager.Metadata{
 		SchemaVersion: "0.1.0",
 		Vendor:        "Docker Inc.",
 		Version:       "test",