PoC for Subcommand Extension - Proposal 1 (Parser-dependent API)

ext/ext.go

@@ -26,6 +26,7 @@ const (
 	JSExtension ExtensionType = iota + 1
 	OutputExtension
 	SecretSourceExtension
+	SubcommandExtension
 )
 
 func (e ExtensionType) String() string {
@@ -37,6 +38,8 @@ func (e ExtensionType) String() string {
 		s = "output"
 	case SecretSourceExtension:
 		s = "secret-source"
+	case SubcommandExtension:
+		s = "subcommand"
 	}
 	return s
 }
@@ -161,4 +164,5 @@ func init() {
 	extensions[JSExtension] = make(map[string]*Extension)
 	extensions[OutputExtension] = make(map[string]*Extension)
 	extensions[SecretSourceExtension] = make(map[string]*Extension)
+	extensions[SubcommandExtension] = make(map[string]*Extension)
 }

internal/cmd/root.go

@@ -95,6 +95,10 @@ func newRootWithLauncher(gs *state.GlobalState, l *launcher) *rootCommand {
 		rootCmd.AddCommand(sc(gs))
 	}
 
+	for sc := range extensionSubcommands(gs, rootCmd.Commands()) {
+		rootCmd.AddCommand(sc)
+	}
+
 	c.cmd = rootCmd
 	return c
 }

internal/cmd/root_test.go

@@ -1,16 +1,24 @@
 package cmd
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
+	"go.k6.io/k6/cmd/state"
 	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/internal/cmd/tests"
+	"go.k6.io/k6/subcommand"
 )
 
 func TestRootCommandHelpDisplayCommands(t *testing.T) {
 	t.Parallel()
 
+	registerTestSubcommandExtensionsOnce.Do(func() {
+		registerTestSubcommandExtensions(t)
+	})
+
 	testCases := []struct {
 		name                  string
 		extraArgs             []string
@@ -70,6 +78,14 @@ func TestRootCommandHelpDisplayCommands(t *testing.T) {
 			name:               "should have version command",
 			wantStdoutContains: "  version     Show application version",
 		},
+		{
+			name:               "should have test-cmd-1 command",
+			wantStdoutContains: "  test-cmd-1  Test command 1",
+		},
+		{
+			name:               "should have test-cmd-2 command",
+			wantStdoutContains: "  test-cmd-2  Test command 2",
+		},
 	}
 
 	for _, tc := range testCases {
@@ -86,3 +102,25 @@ func TestRootCommandHelpDisplayCommands(t *testing.T) {
 		})
 	}
 }
+
+var registerTestSubcommandExtensionsOnce sync.Once //nolint:gochecknoglobals
+
+func registerTestSubcommandExtensions(t *testing.T) {
+	t.Helper()
+
+	subcommand.RegisterExtension("test-cmd-1", func(_ *state.GlobalState) *cobra.Command {
+		return &cobra.Command{
+			Use:   "test-cmd-1",
+			Short: "Test command 1",
+			Run:   func(_ *cobra.Command, _ []string) {},
+		}
+	})
+
+	subcommand.RegisterExtension("test-cmd-2", func(_ *state.GlobalState) *cobra.Command {
+		return &cobra.Command{
+			Use:   "test-cmd-2",
+			Short: "Test command 2",
+			Run:   func(_ *cobra.Command, _ []string) {},
+		}
+	})
+}

internal/cmd/subcommand.go

@@ -0,0 +1,55 @@
+package cmd
+
+import (
+	"iter"
+
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"go.k6.io/k6/cmd/state"
+	"go.k6.io/k6/ext"
+	"go.k6.io/k6/subcommand"
+)
+
+// extensionSubcommands returns an iterator over all registered subcommand extensions
+// that are not already defined in the given slice of commands.
+func extensionSubcommands(gs *state.GlobalState, defined []*cobra.Command) iter.Seq[*cobra.Command] {
+	already := make(map[string]struct{}, len(defined))
+	for _, cmd := range defined {
+		already[cmd.Name()] = struct{}{}
+	}
+
+	return func(yield func(*cobra.Command) bool) {
+		for _, extension := range ext.Get(ext.SubcommandExtension) {
+			if _, exists := already[extension.Name]; exists {
+				gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+					Warnf("subcommand already exists")
+				continue
+			}
+
+			already[extension.Name] = struct{}{}
+
+			if !yield(getCmdForExtension(extension, gs)) {
+				break
+			}
+		}
+	}
+}
+
+// getCmdForExtension gets a *cobra.Command for the given subcommand extension.
+func getCmdForExtension(extension *ext.Extension, gs *state.GlobalState) *cobra.Command {
+	ctor, ok := extension.Module.(subcommand.Constructor)
+	if !ok {
+		gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+			Fatalf("subcommand's constructor does not implement the subcommand.Constructor")
+	}
+
+	cmd := ctor(gs)
+
+	// Validate that the command's name matches the extension name.
+	if cmd.Name() != extension.Name {
+		gs.Logger.WithFields(logrus.Fields{"name": extension.Name, "path": extension.Path}).
+			Fatalf("subcommand's command name (%s) does not match the extension name (%s)", cmd.Name(), extension.Name)
+	}
+
+	return cmd
+}

subcommand/extension.go

@@ -0,0 +1,32 @@
+// Package subcommand provides functionality for registering k6 subcommand extensions.
+//
+// This package allows external modules to register new subcommands that will be
+// available in the k6 CLI. Subcommand extensions are registered during
+// package initialization and are called when the corresponding subcommand is invoked.
+package subcommand
+
+import (
+	"github.com/spf13/cobra"
+	"go.k6.io/k6/cmd/state"
+	"go.k6.io/k6/ext"
+)
+
+// Constructor is a function type that creates a new cobra.Command for a subcommand extension.
+// It receives a GlobalState instance that provides access to configuration, logging,
+// file system, and other shared k6 runtime state. The returned Command will be
+// integrated into k6's CLI as a subcommand.
+type Constructor func(*state.GlobalState) *cobra.Command
+
+// RegisterExtension registers a subcommand extension with the given name and constructor function.
+//
+// The name parameter specifies the subcommand name that users will invoke (e.g., "k6 <name>").
+// The constructor function will be called when k6 initializes to create the cobra.Command
+// instance for this subcommand.
+//
+// This function must be called during package initialization (typically in an init() function)
+// and will panic if a subcommand with the same name is already registered.
+//
+// The name parameter and the returned Command's Name() must match.
+func RegisterExtension(name string, c Constructor) {
+	ext.Register(name, ext.SubcommandExtension, c)
+}