Re-add "all" and support for multiple targets

internal/cmd/check.go

@@ -19,7 +19,6 @@ import (
 	"github.com/turbot/pipe-fittings/constants"
 	"github.com/turbot/pipe-fittings/contexthelpers"
 	"github.com/turbot/pipe-fittings/error_helpers"
-	"github.com/turbot/pipe-fittings/modconfig"
 	"github.com/turbot/pipe-fittings/statushooks"
 	"github.com/turbot/pipe-fittings/utils"
 	localcmdconfig "github.com/turbot/powerpipe/internal/cmdconfig"
@@ -37,10 +36,15 @@ var checkOutputMode = localconstants.CheckOutputModeText
 // generic command to handle benchmark and control execution
 func checkCmd[T controlinit.CheckTarget]() *cobra.Command {
 	typeName := utils.GetGenericTypeName[T]()
+	argsSupported := cobra.ExactArgs(1)
+	if typeName == "benchmark" {
+		argsSupported = cobra.MinimumNArgs(1)
+	}
+
 	cmd := &cobra.Command{
 		Use:              checkCmdUse(typeName),
 		TraverseChildren: true,
-		Args:             cobra.ExactArgs(1),
+		Args:             argsSupported,
 		Run:              runCheckCmd[T],
 		Short:            checkCmdShort(typeName),
 		Long:             checkCmdLong(typeName),
@@ -162,45 +166,46 @@ func runCheckCmd[T controlinit.CheckTarget](cmd *cobra.Command, args []string) {
 
 	// now filter the target
 	// get the execution trees
-	namedTree, err := getExecutionTree[T](ctx, initData)
+	trees, err := getExecutionTrees[T](ctx, initData)
 	error_helpers.FailOnError(err)
 
-	// execute controls synchronously (execute returns the number of alarms and errors)
-
 	// pull out useful properties
 	totalAlarms, totalErrors := 0, 0
 	defer func() {
 		// set the defined exit code after successful execution
 		exitCode = getExitCode(totalAlarms, totalErrors)
 	}()
 
-	err = executeTree(ctx, namedTree.tree, initData)
-	if err != nil {
-		totalErrors++
-		error_helpers.ShowError(ctx, err)
-		return
-	}
+	for _, namedTree := range trees {
+		// execute controls synchronously (execute returns the number of alarms and errors)
+		err = executeTree(ctx, namedTree.tree, initData)
+		if err != nil {
+			totalErrors++
+			error_helpers.ShowError(ctx, err)
+			return
+		}
 
-	// append the total number of alarms and errors for multiple runs
-	totalAlarms = namedTree.tree.Root.Summary.Status.Alarm
-	totalErrors = namedTree.tree.Root.Summary.Status.Error
+		// append the total number of alarms and errors for multiple runs
+		totalAlarms = namedTree.tree.Root.Summary.Status.Alarm
+		totalErrors = namedTree.tree.Root.Summary.Status.Error
 
-	err = publishSnapshot(ctx, namedTree.tree, viper.GetBool(constants.ArgShare), viper.GetBool(constants.ArgSnapshot))
-	if err != nil {
-		error_helpers.ShowError(ctx, err)
-		totalErrors++
-		return
-	}
-	if shouldPrintCheckTiming() {
-		display.PrintTiming(&localqueryresult.TimingMetadata{
-			Duration: time.Since(startTime),
-		})
-	}
+		err = publishSnapshot(ctx, namedTree.tree, viper.GetBool(constants.ArgShare), viper.GetBool(constants.ArgSnapshot))
+		if err != nil {
+			error_helpers.ShowError(ctx, err)
+			totalErrors++
+			return
+		}
+		if shouldPrintCheckTiming() {
+			display.PrintTiming(&localqueryresult.TimingMetadata{
+				Duration: time.Since(startTime),
+			})
+		}
 
-	err = exportExecutionTree(ctx, namedTree, initData, viper.GetStringSlice(constants.ArgExport))
-	if err != nil {
-		error_helpers.ShowError(ctx, err)
-		totalErrors++
+		err = exportExecutionTree(ctx, namedTree, initData, viper.GetStringSlice(constants.ArgExport))
+		if err != nil {
+			error_helpers.ShowError(ctx, err)
+			totalErrors++
+		}
 	}
 }
 
@@ -254,38 +259,36 @@ func publishSnapshot(ctx context.Context, executionTree *controlexecute.Executio
 	return nil
 }
 
-func getExecutionTree[T controlinit.CheckTarget](ctx context.Context, initData *controlinit.InitData[T]) (*namedExecutionTree, error) {
-	// todo kai needed???
+func getExecutionTrees[T controlinit.CheckTarget](ctx context.Context, initData *controlinit.InitData[T]) ([]*namedExecutionTree, error) {
+	var trees []*namedExecutionTree
 	if error_helpers.IsContextCanceled(ctx) {
 		return nil, ctx.Err()
 	}
 
-	target := initData.Target
-	executionTree, err := controlexecute.NewExecutionTree(ctx, initData.Workspace, initData.DefaultClient, initData.ControlFilter, target)
-	if err != nil {
-		return nil, sperr.WrapWithMessage(err, "could not create merged execution tree")
-	}
-
-	var name string
 	if initData.ExportManager.HasNamedExport(viper.GetStringSlice(constants.ArgExport)) {
-		name = fmt.Sprintf("check.%s", initData.Workspace.Mod.ShortName)
+		// if there is a named export - combine targets into a single tree
+		executionTree, err := controlexecute.NewExecutionTree(ctx, initData.Workspace, initData.DefaultClient, initData.ControlFilter, initData.Targets...)
+		if err != nil {
+			return nil, sperr.WrapWithMessage(err, "could not create merged execution tree")
+		}
+		name := fmt.Sprintf("check.%s", initData.Workspace.Mod.ShortName)
+		trees = append(trees, newNamedExecutionTree(name, executionTree))
 	} else {
-		name = getExportName(target.Name(), initData.Workspace.Mod.ShortName)
+		// otherwise return multiple trees
+		for _, target := range initData.Targets {
+			if error_helpers.IsContextCanceled(ctx) {
+				return nil, ctx.Err()
+			}
+			executionTree, err := controlexecute.NewExecutionTree(ctx, initData.Workspace, initData.DefaultClient, initData.ControlFilter, target)
+			if err != nil {
+				return nil, sperr.WrapWithMessage(err, "could not create execution tree for %s", target)
+			}
+
+			trees = append(trees, newNamedExecutionTree(target.Name(), executionTree))
+		}
 	}
 
-	return newNamedExecutionTree(name, executionTree), ctx.Err()
-
-}
-
-// getExportName resolves the base name of the target file
-func getExportName(targetName string, modShortName string) string {
-	parsedName, _ := modconfig.ParseResourceName(targetName)
-	if targetName == "all" {
-		// there will be no block type = manually construct name
-		return fmt.Sprintf("%s.%s", modShortName, parsedName.Name)
-	}
-	// default to just converting to valid resource name
-	return parsedName.ToFullNameWithMod(modShortName)
+	return trees, ctx.Err()
 }
 
 // get the exit code for successful check run

internal/cmd/dashboard.go

@@ -132,7 +132,9 @@ func dashboardRun(cmd *cobra.Command, args []string) {
 	initData.Result.DisplayMessages()
 
 	// so a dashboard name was specified - just call GenerateSnapshot
-	snap, err := dashboardexecute.GenerateSnapshot(ctx, initData.WorkspaceEvents, initData.Target, inputs)
+	target, err := initData.GetSingleTarget()
+	error_helpers.FailOnError(err)
+	snap, err := dashboardexecute.GenerateSnapshot(ctx, initData.WorkspaceEvents, target, inputs)
 	error_helpers.FailOnError(err)
 	// display the snapshot result (if needed)
 	displaySnapshot(snap)

internal/cmd/query.go

@@ -114,7 +114,7 @@ func queryRun(cmd *cobra.Command, args []string) {
 	// if there is a usage warning we display it
 	initData.Result.DisplayMessages()
 
-	// TODO check cancellation
+	// TODO KAI check cancellation
 	// start cancel handler to intercept interrupts and cancel the context
 	// NOTE: use the initData Cancel function to ensure any initialisation is cancelled if needed
 	//contexthelpers.StartCancelHandler(initData.Cancel)
@@ -125,7 +125,12 @@ func queryRun(cmd *cobra.Command, args []string) {
 	}
 
 	// execute query as a snapshot
-	snap, err := dashboardexecute.GenerateSnapshot(ctx, initData.WorkspaceEvents, initData.Target, nil)
+	target, err := initData.GetSingleTarget()
+	if err != nil {
+		exitCode = constants.ExitCodeInitializationFailed
+		error_helpers.FailOnError(err)
+	}
+	snap, err := dashboardexecute.GenerateSnapshot(ctx, initData.WorkspaceEvents, target, nil)
 	if err != nil {
 		exitCode = constants.ExitCodeSnapshotCreationFailed
 		error_helpers.FailOnError(err)

internal/cmdconfig/cmd_targets.go

@@ -0,0 +1,197 @@
+package cmdconfig
+
+import (
+	"fmt"
+	"golang.org/x/exp/maps"
+	"strings"
+
+	"github.com/spf13/viper"
+	"github.com/turbot/go-kit/helpers"
+	"github.com/turbot/pipe-fittings/constants"
+	"github.com/turbot/pipe-fittings/modconfig"
+	"github.com/turbot/pipe-fittings/workspace"
+	"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
+)
+
+func ResolveTargets[T modconfig.ModTreeItem](cmdArgs []string, w *workspace.Workspace) ([]modconfig.ModTreeItem, error) {
+	if len(cmdArgs) == 0 {
+		return nil, nil
+	}
+
+	// special case handling for `benchmark run all`
+	targets, err := handleAllArg[T](cmdArgs, w)
+	if err != nil {
+		return nil, err
+	}
+	if targets != nil {
+		return targets, nil
+	}
+	if len(cmdArgs) == 1 {
+		return resolveSingleTarget[T](cmdArgs[0], w)
+	}
+	// the only command which supports multiple targets is benchmark run
+	return resolveBenchmarkTargets[T](cmdArgs, w)
+}
+
+// resolveSingleTarget extracts a single target and (if present) query args from the command line parameters
+//   - if no resource type is specified in the name, it is added from the command type
+//   - validate the resource type specified in the name matches the command type
+//   - verify the resource exists in the workspace
+//   - if the command type is 'query', the target may be a query string rather than a resource name
+//     in this case, convert into a query and add to workspace (to allow for simple snapshot generation)
+func resolveSingleTarget[T modconfig.ModTreeItem](cmdArg string, w *workspace.Workspace) ([]modconfig.ModTreeItem, error) {
+
+	var target modconfig.ModTreeItem
+	var queryArgs *modconfig.QueryArgs
+	var err error
+	// so there are multiple targets  - this must be the benchmark command, so we do not expect any args
+	// now try to resolve targets
+	// NOTE: we only expect multiple targets for benchmarks which do not support query args
+	// however for resilience (and in case this changes), collect query args into a map
+
+	target, queryArgs, err = workspace.ResolveResourceAndArgsFromSQLString[T](cmdArg, w)
+	if err != nil {
+		return nil, err
+	}
+	if helpers.IsNil(target) {
+		return nil, fmt.Errorf("'%s' not found in %s (%s)", cmdArg, w.Mod.Name(), w.Path)
+	}
+	if queryArgs != nil {
+		return nil, sperr.New("benchmarks do not support query args")
+	}
+
+	// ok we managed to resolve
+
+	// now check if any args were specified on the command line using the --arg flag
+	// if so verify no args were passed in the resource invocation, e.g. query.my_query("val1","val1"
+	commandLineQueryArgs, err := getCommandLineQueryArgs()
+	if err != nil {
+		return nil, err
+	}
+
+	// so args were passed using --arg
+	if !commandLineQueryArgs.Empty() {
+		// verify no args were passed in the resource invocation, e.g. query.my_query("val1","val2")
+		if queryArgs != nil {
+			return nil, sperr.New("both command line args and query invocation args are set")
+		}
+	}
+	// set args for targer
+	if queryArgs != nil {
+		// if the target is a query provider set the args
+		// (if the target is a dashboard, which i snot a query provider,
+		// we read the args from viper separately and use to populate the inputs)
+		if qp, ok := any(target).(modconfig.QueryProvider); ok {
+			qp.SetArgs(queryArgs)
+		}
+	}
+	return []modconfig.ModTreeItem{target}, nil
+
+}
+
+func resolveBenchmarkTargets[T modconfig.ModTreeItem](cmdArgs []string, w *workspace.Workspace) ([]modconfig.ModTreeItem, error) {
+	var targets []modconfig.ModTreeItem
+	// so there are multiple targets  - this must be the benchmark command, so we do not expect any args
+	// verify T is Benchmark (should be enforced by Cobra)
+	var empty T
+	if _, isBenchmark := (any(empty)).(*modconfig.Benchmark); !isBenchmark {
+		return nil, sperr.New("multiple targets are only supported for benchmarks")
+	}
+
+	// now try to resolve targets
+	for _, cmdArg := range cmdArgs {
+		target, queryArgs, err := workspace.ResolveResourceAndArgsFromSQLString[T](cmdArg, w)
+		if err != nil {
+			return nil, err
+		}
+		if helpers.IsNil(target) {
+			return nil, fmt.Errorf("'%s' not found in %s (%s)", cmdArgs[0], w.Mod.Name(), w.Path)
+		}
+		if queryArgs != nil {
+			return nil, sperr.New("benchmarks do not support query args")
+		}
+		targets = append(targets, target)
+	}
+
+	return targets, nil
+}
+
+func handleAllArg[T modconfig.ModTreeItem](args []string, w *workspace.Workspace) ([]modconfig.ModTreeItem, error) {
+	isAll := len(args) == 1 && args[0] == "all"
+	if !isAll {
+		return nil, nil
+	}
+	var empty T
+	if _, isBenchmark := (any(empty)).(*modconfig.Benchmark); !isBenchmark {
+		return nil, nil
+	}
+
+	// if the arg is "all", we want to execute all _direct_ children of the Mod
+	// but NOT children which come from dependency mods
+	filter := workspace.ResourceFilter{
+		WherePredicate: func(item modconfig.HclResource) bool {
+			mti, ok := item.(modconfig.ModTreeItem)
+			if !ok {
+				return false
+			}
+			return mti.GetMod().ShortName == w.Mod.ShortName
+		},
+	}
+	targetsMap, err := workspace.FilterWorkspaceResourcesOfType[T](w, filter)
+	if err != nil {
+		return nil, err
+	}
+
+	targets := ToModTreeItemSlice(maps.Values(targetsMap))
+
+	// make a root item to hold the benchmarks
+	resolvedItem := modconfig.NewRootBenchmarkWithChildren(w.Mod, targets).(modconfig.ModTreeItem)
+
+	return []modconfig.ModTreeItem{resolvedItem}, nil
+
+}
+
+// build a QueryArgs from any args passed using the --args flag
+func getCommandLineQueryArgs() (*modconfig.QueryArgs, error) {
+	argTuples := viper.GetStringSlice(constants.ArgArg)
+	var res = modconfig.NewQueryArgs()
+
+	if argTuples == nil {
+		return res, nil
+	}
+
+	for _, argTuple := range argTuples {
+		parts := strings.Split(argTuple, "=")
+		switch len(parts) {
+		case 1:
+			// if there is no '=' this must be a positional arg
+			if err := res.AddPositionalArgVal(parts[0]); err != nil {
+				return nil, err
+			}
+
+		case 2:
+			argName := parts[0]
+			argValue := parts[1]
+
+			if err := res.SetNamedArgVal(argName, argValue); err != nil {
+				return nil, err
+			}
+		default:
+			return nil, sperr.New("invalid arg format: %s", argTuple)
+		}
+	}
+	// we should not have both positional and named args
+	if len(res.ArgMap) > 0 && len(res.ArgList) > 0 {
+		return nil, sperr.New("cannot mix positional and named args")
+	}
+	return res, nil
+
+}
+
+func ToModTreeItemSlice[T modconfig.ModTreeItem](items []T) []modconfig.ModTreeItem {
+	var res []modconfig.ModTreeItem
+	for _, item := range items {
+		res = append(res, item)
+	}
+	return res
+}