Validate checkout output and export formats before execution. Closes #2619. Add support for using sps as snapshot output/export format. Closes #2623

cmd/check.go

@@ -14,14 +14,12 @@ import (
 	"github.com/turbot/steampipe/pkg/cmdconfig"
 	"github.com/turbot/steampipe/pkg/constants"
 	"github.com/turbot/steampipe/pkg/contexthelpers"
+	"github.com/turbot/steampipe/pkg/control"
 	"github.com/turbot/steampipe/pkg/control/controldisplay"
 	"github.com/turbot/steampipe/pkg/control/controlexecute"
 	"github.com/turbot/steampipe/pkg/control/controlstatus"
 	"github.com/turbot/steampipe/pkg/display"
 	"github.com/turbot/steampipe/pkg/error_helpers"
-	"github.com/turbot/steampipe/pkg/initialisation"
-	"github.com/turbot/steampipe/pkg/interactive"
-	"github.com/turbot/steampipe/pkg/statushooks"
 	"github.com/turbot/steampipe/pkg/utils"
 	"github.com/turbot/steampipe/pkg/workspace"
 )
@@ -91,7 +89,6 @@ You may specify one or more benchmarks or controls to run (separated by a space)
 
 func runCheckCmd(cmd *cobra.Command, args []string) {
 	utils.LogTime("runCheckCmd start")
-	initData := &initialisation.InitData{}
 
 	// setup a cancel context and start cancel handler
 	ctx, cancel := context.WithCancel(cmd.Context())
@@ -105,8 +102,6 @@ func runCheckCmd(cmd *cobra.Command, args []string) {
 			error_helpers.ShowError(ctx, helpers.ToError(r))
 			exitCode = constants.ExitCodeUnknownErrorPanic
 		}
-
-		initData.Cleanup(ctx)
 	}()
 
 	// verify we have an argument
@@ -121,8 +116,10 @@ func runCheckCmd(cmd *cobra.Command, args []string) {
 	}
 
 	// initialise
-	initData = initialiseCheck(ctx)
+	initData := control.NewInitData(ctx)
 	error_helpers.FailOnError(initData.Result.Error)
+	defer initData.Cleanup(ctx)
+
 	// if there is a usage warning we display it
 	initData.Result.DisplayMessages()
 
@@ -151,7 +148,7 @@ func runCheckCmd(cmd *cobra.Command, args []string) {
 
 		// execute controls synchronously (execute returns the number of failures)
 		failures += executionTree.Execute(ctx)
-		err = displayControlResults(ctx, executionTree)
+		err = displayControlResults(ctx, executionTree, initData.OutputFormatter)
 		error_helpers.FailOnError(err)
 
 		exportArgs := viper.GetStringSlice(constants.ArgExport)
@@ -201,78 +198,6 @@ func validateCheckArgs(ctx context.Context, cmd *cobra.Command, args []string) b
 	return true
 }
 
-func initialiseCheck(ctx context.Context) *initialisation.InitData {
-	statushooks.SetStatus(ctx, "Initializing...")
-	defer statushooks.Done(ctx)
-
-	// load the workspace
-	w, err := interactive.LoadWorkspacePromptingForVariables(ctx)
-	error_helpers.FailOnErrorWithMessage(err, "failed to load workspace")
-
-	initData := initialisation.NewInitData(w).Init(ctx, constants.InvokerCheck)
-	if initData.Result.Error != nil {
-		return initData
-	}
-
-	// control specific init
-	if !w.ModfileExists() {
-		initData.Result.Error = workspace.ErrorNoModDefinition
-	}
-
-	if viper.GetString(constants.ArgOutput) == constants.OutputFormatNone {
-		// set progress to false
-		viper.Set(constants.ArgProgress, false)
-	}
-	// set color schema
-	err = initialiseCheckColorScheme()
-	if err != nil {
-		initData.Result.Error = err
-		return initData
-	}
-
-	if len(initData.Workspace.GetResourceMaps().Controls) == 0 {
-		initData.Result.AddWarnings("no controls found in current workspace")
-	}
-
-	if err := controldisplay.EnsureTemplates(); err != nil {
-		initData.Result.Error = err
-		return initData
-	}
-
-	if len(viper.GetStringSlice(constants.ArgExport)) > 0 {
-		registerCheckExporters(initData)
-	}
-
-	return initData
-}
-
-// register exporters for each of the supported check formats
-func registerCheckExporters(initData *initialisation.InitData) {
-	exporters, err := controldisplay.GetExporters()
-	error_helpers.FailOnErrorWithMessage(err, "failed to load exporters")
-
-	// register all exporters
-	initData.RegisterExporters(exporters...)
-}
-
-func initialiseCheckColorScheme() error {
-	theme := viper.GetString(constants.ArgTheme)
-	if !viper.GetBool(constants.ConfigKeyIsTerminalTTY) {
-		// enforce plain output for non-terminals
-		theme = "plain"
-	}
-	themeDef, ok := controldisplay.ColorSchemes[theme]
-	if !ok {
-		return fmt.Errorf("invalid theme '%s'", theme)
-	}
-	scheme, err := controldisplay.NewControlColorScheme(themeDef)
-	if err != nil {
-		return err
-	}
-	controldisplay.ControlColors = scheme
-	return nil
-}
-
 func printTiming(args []string, durations []time.Duration) {
 	headers := []string{"", "Duration"}
 	var rows [][]string
@@ -292,13 +217,7 @@ func shouldPrintTiming() bool {
 		(outputFormat == constants.OutputFormatText || outputFormat == constants.OutputFormatBrief)
 }
 
-func displayControlResults(ctx context.Context, executionTree *controlexecute.ExecutionTree) error {
-	output := viper.GetString(constants.ArgOutput)
-	formatter, err := parseOutputArg(output)
-	if err != nil {
-		fmt.Println(err)
-		return err
-	}
+func displayControlResults(ctx context.Context, executionTree *controlexecute.ExecutionTree, formatter controldisplay.Formatter) error {
 	reader, err := formatter.Format(ctx, executionTree)
 	if err != nil {
 		return err
@@ -306,13 +225,3 @@ func displayControlResults(ctx context.Context, executionTree *controlexecute.Ex
 	_, err = io.Copy(os.Stdout, reader)
 	return err
 }
-
-// parseOutputArg parses the --output flag value and returns the Formatter that can format the data
-func parseOutputArg(arg string) (formatter controldisplay.Formatter, err error) {
-	formatResolver, err := controldisplay.NewFormatResolver()
-	if err != nil {
-		return nil, err
-	}
-
-	return formatResolver.GetFormatter(arg)
-}

cmd/dashboard.go

@@ -23,7 +23,6 @@ import (
 	"github.com/turbot/steampipe/pkg/error_helpers"
 	"github.com/turbot/steampipe/pkg/export"
 	"github.com/turbot/steampipe/pkg/initialisation"
-	"github.com/turbot/steampipe/pkg/interactive"
 	"github.com/turbot/steampipe/pkg/statushooks"
 	"github.com/turbot/steampipe/pkg/steampipeconfig/modconfig"
 	"github.com/turbot/steampipe/pkg/utils"
@@ -86,7 +85,6 @@ func runDashboardCmd(cmd *cobra.Command, args []string) {
 			}
 		}
 		setExitCodeForDashboardError(err)
-
 	}()
 
 	// first check whether a dashboard name has been passed as an arg
@@ -103,7 +101,7 @@ func runDashboardCmd(cmd *cobra.Command, args []string) {
 		inputs, err := collectInputs()
 		error_helpers.FailOnError(err)
 
-		// run just this dashboard
+		// run just this dashboard - this handles all initialisation
 		err = runSingleDashboard(dashboardCtx, dashboardName, inputs)
 		error_helpers.FailOnError(err)
 
@@ -189,10 +187,10 @@ func validateDashboardArgs(ctx context.Context, args []string) (string, error) {
 		}
 	}
 
-	validOutputFormats := []string{constants.OutputFormatSnapshot, constants.OutputFormatNone}
+	validOutputFormats := []string{constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort, constants.OutputFormatNone}
 	output := viper.GetString(constants.ArgOutput)
 	if !helpers.StringSliceContains(validOutputFormats, output) {
-		return "", fmt.Errorf("invalid output format '%s', must be one of %s", output, strings.Join(validOutputFormats, ","))
+		return "", fmt.Errorf("invalid output format: '%s', must be one of [%s]", output, strings.Join(validOutputFormats, ", "))
 	}
 
 	return dashboardName, nil
@@ -200,7 +198,7 @@ func validateDashboardArgs(ctx context.Context, args []string) (string, error) {
 
 func displaySnapshot(snapshot *dashboardtypes.SteampipeSnapshot) {
 	switch viper.GetString(constants.ArgOutput) {
-	case constants.OutputFormatSnapshot:
+	case constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort:
 		// just display result
 		snapshotText, err := json.MarshalIndent(snapshot, "", "  ")
 		error_helpers.FailOnError(err)
@@ -210,27 +208,40 @@ func displaySnapshot(snapshot *dashboardtypes.SteampipeSnapshot) {
 
 func initDashboard(ctx context.Context) *initialisation.InitData {
 	dashboardserver.OutputWait(ctx, "Loading Workspace")
-	w, err := interactive.LoadWorkspacePromptingForVariables(ctx)
-	if err != nil {
-		return initialisation.NewErrorInitData(fmt.Errorf("failed to load workspace: %s", err.Error()))
-	}
 
 	// initialise
-	initData := getInitData(ctx, w)
+	initData := getInitData(ctx)
+	if initData.Result.Error != nil {
+		return initData
+	}
 
 	// there must be a mod-file
-	if !w.ModfileExists() {
+	if !initData.Workspace.ModfileExists() {
 		initData.Result.Error = workspace.ErrorNoModDefinition
 	}
 
 	return initData
 }
 
-func getInitData(ctx context.Context, w *workspace.Workspace) *initialisation.InitData {
-	initData := initialisation.NewInitData(w).
-		RegisterExporters(dashboardExporters()...).
+func getInitData(ctx context.Context) *initialisation.InitData {
+	w, err := workspace.LoadWorkspacePromptingForVariables(ctx)
+	if err != nil {
+		return initialisation.NewErrorInitData(fmt.Errorf("failed to load workspace: %s", err.Error()))
+	}
+
+	i := initialisation.NewInitData(w).
 		Init(ctx, constants.InvokerDashboard)
-	return initData
+
+	if len(viper.GetStringSlice(constants.ArgExport)) > 0 {
+		i.RegisterExporters(dashboardExporters()...)
+		// validate required export formats
+		if err := i.ExportManager.ValidateExportFormat(viper.GetStringSlice(constants.ArgExport)); err != nil {
+			i.Result.Error = err
+			return i
+		}
+	}
+
+	return i
 }
 
 func dashboardExporters() []export.Exporter {
@@ -238,22 +249,17 @@ func dashboardExporters() []export.Exporter {
 }
 
 func runSingleDashboard(ctx context.Context, targetName string, inputs map[string]interface{}) error {
-	w, err := interactive.LoadWorkspacePromptingForVariables(ctx)
-	error_helpers.FailOnErrorWithMessage(err, "failed to load workspace")
-
-	// targetName must be a named resource
-	// parse the name to verify
-	if err := verifyNamedResource(targetName, w); err != nil {
-		return err
-	}
-
-	initData := getInitData(ctx, w)
-
+	initData := getInitData(ctx)
 	// shutdown the service on exit
 	defer initData.Cleanup(ctx)
 	if err := initData.Result.Error; err != nil {
 		return initData.Result.Error
 	}
+	// targetName must be a named resource
+	// parse the name to verify
+	if err := verifyNamedResource(targetName, initData.Workspace); err != nil {
+		return err
+	}
 
 	// if there is a usage warning we display it
 	initData.Result.DisplayMessages()

cmd/query.go

@@ -20,7 +20,6 @@ import (
 	"github.com/turbot/steampipe/pkg/dashboard/dashboardtypes"
 	"github.com/turbot/steampipe/pkg/display"
 	"github.com/turbot/steampipe/pkg/error_helpers"
-	"github.com/turbot/steampipe/pkg/interactive"
 	"github.com/turbot/steampipe/pkg/query"
 	"github.com/turbot/steampipe/pkg/query/queryexecute"
 	"github.com/turbot/steampipe/pkg/query/queryresult"
@@ -125,15 +124,10 @@ func runQueryCmd(cmd *cobra.Command, args []string) {
 	// set config to indicate whether we are running an interactive query
 	viper.Set(constants.ConfigKeyInteractive, interactiveMode)
 
-	// load the workspace
-	w, err := interactive.LoadWorkspacePromptingForVariables(ctx)
-	error_helpers.FailOnErrorWithMessage(err, "failed to load workspace")
-
-	// so we have loaded a workspace - be sure to close it
-	defer w.Close()
-
 	// start the initializer
-	initData := query.NewInitData(ctx, w, args)
+	initData := query.NewInitData(ctx, args)
+	error_helpers.FailOnError(initData.Result.Error)
+	defer initData.Cleanup(ctx)
 
 	switch {
 	case interactiveMode:
@@ -141,7 +135,7 @@ func runQueryCmd(cmd *cobra.Command, args []string) {
 	case snapshotRequired():
 		// if we are either outputting snapshot format, or sharing the results as a snapshot, execute the query
 		// as a dashboard
-		exitCode = executeSnapshotQuery(initData, w, ctx)
+		exitCode = executeSnapshotQuery(initData, ctx)
 	default:
 		// NOTE: disable any status updates - we do not want 'loading' output from any queries
 		ctx = statushooks.DisableStatusHooks(ctx)
@@ -163,19 +157,16 @@ func validateQueryArgs(ctx context.Context, args []string) error {
 		return err
 	}
 
-	validOutputFormats := []string{constants.OutputFormatLine, constants.OutputFormatCSV, constants.OutputFormatTable, constants.OutputFormatJSON, constants.OutputFormatSnapshot, constants.OutputFormatNone}
+	validOutputFormats := []string{constants.OutputFormatLine, constants.OutputFormatCSV, constants.OutputFormatTable, constants.OutputFormatJSON, constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort, constants.OutputFormatNone}
 	output := viper.GetString(constants.ArgOutput)
 	if !helpers.StringSliceContains(validOutputFormats, output) {
-		return fmt.Errorf("invalid output format '%s', must be one of %s", output, strings.Join(validOutputFormats, ","))
+		return fmt.Errorf("invalid output format: '%s', must be one of [%s]", output, strings.Join(validOutputFormats, ", "))
 	}
 
 	return nil
 }
 
-func executeSnapshotQuery(initData *query.InitData, w *workspace.Workspace, ctx context.Context) int {
-	// ensure we close client
-	defer initData.Cleanup(ctx)
-
+func executeSnapshotQuery(initData *query.InitData, ctx context.Context) int {
 	// 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)
@@ -197,7 +188,7 @@ func executeSnapshotQuery(initData *query.InitData, w *workspace.Workspace, ctx
 			// this is to allow us to use existing dashboard execution code
 
 			// build query name and title
-			targetName := ensureQueryResource(name, query, i, len(queryNames), w)
+			targetName := ensureQueryResource(name, query, i, len(queryNames), initData.Workspace)
 
 			// we need to pass the embedded initData to  GenerateSnapshot
 			baseInitData := &initData.InitData
@@ -210,7 +201,7 @@ func executeSnapshotQuery(initData *query.InitData, w *workspace.Workspace, ctx
 			switch viper.GetString(constants.ArgOutput) {
 			case constants.OutputFormatNone:
 				// do nothing
-			case constants.OutputFormatSnapshot:
+			case constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort:
 				// if the format is snapshot, just dump it out
 				jsonOutput, err := json.MarshalIndent(snap, "", "  ")
 				if err != nil {
@@ -301,16 +292,17 @@ func ensureQueryResource(name string, query string, queryIdx, queryCount int, w
 }
 
 func snapshotRequired() bool {
+	SnapshotFormatNames := []string{constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort}
 	// if a snapshot exporter is specified return true
 	for _, e := range viper.GetStringSlice(constants.ArgExport) {
-		if e == constants.OutputFormatSnapshot || path.Ext(e) == constants.SnapshotExtension {
+		if helpers.StringSliceContains(SnapshotFormatNames, e) || path.Ext(e) == constants.SnapshotExtension {
 			return true
 		}
 	}
 	// if share/snapshot args are set or output is snapshot, return true
 	return viper.IsSet(constants.ArgShare) ||
 		viper.IsSet(constants.ArgSnapshot) ||
-		viper.GetString(constants.ArgOutput) == constants.OutputFormatSnapshot
+		helpers.StringSliceContains(SnapshotFormatNames, viper.GetString(constants.ArgOutput))
 
 }
 

pkg/constants/output_format.go

@@ -1,12 +1,13 @@
 package constants
 
 const (
-	OutputFormatCSV      = "csv"
-	OutputFormatJSON     = "json"
-	OutputFormatTable    = "table"
-	OutputFormatLine     = "line"
-	OutputFormatNone     = "none"
-	OutputFormatText     = "text"
-	OutputFormatBrief    = "brief"
-	OutputFormatSnapshot = "snapshot"
+	OutputFormatCSV           = "csv"
+	OutputFormatJSON          = "json"
+	OutputFormatTable         = "table"
+	OutputFormatLine          = "line"
+	OutputFormatNone          = "none"
+	OutputFormatText          = "text"
+	OutputFormatBrief         = "brief"
+	OutputFormatSnapshot      = "snapshot"
+	OutputFormatSnapshotShort = "sps"
 )

pkg/control/controldisplay/format_resolver.go

@@ -56,7 +56,7 @@ func (r *FormatResolver) GetFormatter(arg string) (Formatter, error) {
 		return formatter, nil
 	}
 
-	return nil, fmt.Errorf("could not resolve formatter for %s", arg)
+	return nil, fmt.Errorf(" invalid output format: '%s'", arg)
 }
 
 func (r *FormatResolver) registerFormatter(f Formatter) error {

pkg/control/controldisplay/formatter_snapshot.go

@@ -40,5 +40,5 @@ func (f SnapshotFormatter) Name() string {
 }
 
 func (f *SnapshotFormatter) Alias() string {
-	return ""
+	return "sps"
 }

pkg/control/init_data.go

@@ -0,0 +1,123 @@
+package control
+
+import (
+	"context"
+	"fmt"
+	"github.com/turbot/steampipe/pkg/control/controldisplay"
+	"github.com/turbot/steampipe/pkg/error_helpers"
+	"github.com/turbot/steampipe/pkg/statushooks"
+
+	"github.com/spf13/viper"
+	"github.com/turbot/steampipe/pkg/constants"
+	"github.com/turbot/steampipe/pkg/initialisation"
+	"github.com/turbot/steampipe/pkg/workspace"
+)
+
+type InitData struct {
+	initialisation.InitData
+	OutputFormatter controldisplay.Formatter
+}
+
+// NewInitData returns a new InitData object
+// It also starts an asynchronous population of the object
+// InitData.Done closes after asynchronous initialization completes
+func NewInitData(ctx context.Context) *InitData {
+
+	statushooks.SetStatus(ctx, "Initializing...")
+	defer statushooks.Done(ctx)
+
+	// load the workspace
+	w, err := workspace.LoadWorkspacePromptingForVariables(ctx)
+	if err != nil {
+		return &InitData{
+			InitData: *initialisation.NewErrorInitData(fmt.Errorf("failed to load workspace: %s", err.Error())),
+		}
+	}
+
+	i := &InitData{
+		InitData: *initialisation.NewInitData(w).Init(ctx, constants.InvokerCheck),
+	}
+	if i.Result.Error != nil {
+		return i
+	}
+
+	if !w.ModfileExists() {
+		i.Result.Error = workspace.ErrorNoModDefinition
+	}
+
+	if viper.GetString(constants.ArgOutput) == constants.OutputFormatNone {
+		// set progress to false
+		viper.Set(constants.ArgProgress, false)
+	}
+	// set color schema
+	err = initialiseCheckColorScheme()
+	if err != nil {
+		i.Result.Error = err
+		return i
+	}
+
+	if len(i.Workspace.GetResourceMaps().Controls) == 0 {
+		i.Result.AddWarnings("no controls found in current workspace")
+	}
+
+	if err := controldisplay.EnsureTemplates(); err != nil {
+		i.Result.Error = err
+		return i
+	}
+
+	if len(viper.GetStringSlice(constants.ArgExport)) > 0 {
+		i.registerCheckExporters()
+		// validate required export formats
+		if err := i.ExportManager.ValidateExportFormat(viper.GetStringSlice(constants.ArgExport)); err != nil {
+			i.Result.Error = err
+			return i
+		}
+	}
+
+	output := viper.GetString(constants.ArgOutput)
+	formatter, err := parseOutputArg(output)
+	if err != nil {
+		i.Result.Error = err
+		return i
+	}
+	i.OutputFormatter = formatter
+
+	return i
+}
+
+// register exporters for each of the supported check formats
+func (initData *InitData) registerCheckExporters() {
+	exporters, err := controldisplay.GetExporters()
+	error_helpers.FailOnErrorWithMessage(err, "failed to load exporters")
+
+	// register all exporters
+	initData.RegisterExporters(exporters...)
+}
+
+// parseOutputArg parses the --output flag value and returns the Formatter that can format the data
+func parseOutputArg(arg string) (formatter controldisplay.Formatter, err error) {
+	formatResolver, err := controldisplay.NewFormatResolver()
+	if err != nil {
+		return nil, err
+	}
+
+	return formatResolver.GetFormatter(arg)
+}
+
+func initialiseCheckColorScheme() error {
+	theme := viper.GetString(constants.ArgTheme)
+	if !viper.GetBool(constants.ConfigKeyIsTerminalTTY) {
+		// enforce plain output for non-terminals
+		theme = "plain"
+	}
+	themeDef, ok := controldisplay.ColorSchemes[theme]
+	if !ok {
+		return fmt.Errorf("invalid theme '%s'", theme)
+	}
+	scheme, err := controldisplay.NewControlColorScheme(themeDef)
+	if err != nil {
+		return err
+	}
+	controldisplay.ControlColors = scheme
+	return nil
+}

pkg/export/manager.go

@@ -3,10 +3,12 @@ package export
 import (
 	"context"
 	"fmt"
-	"github.com/turbot/steampipe/pkg/error_helpers"
 	"golang.org/x/exp/maps"
 	"path"
 	"strings"
+
+	"github.com/turbot/steampipe/pkg/error_helpers"
+	"github.com/turbot/steampipe/pkg/utils"
 )
 
 type Manager struct {
@@ -148,3 +150,17 @@ func (m *Manager) DoExport(ctx context.Context, targetName string, source Export
 	}
 	return error_helpers.CombineErrors(errors...)
 }
+
+func (m *Manager) ValidateExportFormat(exports []string) error {
+	var invalidFormats []string
+	for _, export := range exports {
+		if _, err := m.getExportTarget(export, "dummy_target_name"); err != nil {
+			invalidFormats = append(invalidFormats, export)
+		}
+	}
+	if invalidCount := len(invalidFormats); invalidCount > 0 {
+		return fmt.Errorf("invalid export %s: '%s'", utils.Pluralize("format", invalidCount), strings.Join(invalidFormats, "','"))
+	}
+	return nil
+
+}

pkg/export/snapshot_exporter.go

@@ -33,6 +33,10 @@ func (e *SnapshotExporter) FileExtension() string {
 	return constants.SnapshotExtension
 }
 
-func (e SnapshotExporter) Name() string {
+func (e *SnapshotExporter) Name() string {
 	return constants.OutputFormatSnapshot
 }
+
+func (*SnapshotExporter) Alias() string {
+	return "sps"
+}

pkg/initialisation/init_data.go

@@ -171,7 +171,6 @@ func (i *InitData) Cleanup(ctx context.Context) {
 	if i.Client != nil {
 		i.Client.Close(ctx)
 	}
-
 	if i.ShutdownTelemetry != nil {
 		i.ShutdownTelemetry()
 	}

pkg/interactive/load_workspace.go

@@ -1,39 +0,0 @@
-package interactive
-
-import (
-	"context"
-	"log"
-
-	"github.com/spf13/viper"
-	"github.com/turbot/steampipe/pkg/constants"
-	"github.com/turbot/steampipe/pkg/statushooks"
-	"github.com/turbot/steampipe/pkg/steampipeconfig/modconfig"
-	"github.com/turbot/steampipe/pkg/workspace"
-)
-
-func LoadWorkspacePromptingForVariables(ctx context.Context) (*workspace.Workspace, error) {
-	workspacePath := viper.GetString(constants.ArgModLocation)
-
-	w, err := workspace.Load(ctx, workspacePath)
-	if err == nil {
-		return w, nil
-	}
-	missingVariablesError, ok := err.(modconfig.MissingVariableError)
-	// if there was an error which is NOT a MissingVariableError, return it
-	if !ok {
-		return nil, err
-	}
-	// if interactive input is disabled, return the missing variables error
-	if !viper.GetBool(constants.ArgInput) {
-		return nil, missingVariablesError
-	}
-	// so we have missing variables - prompt for them
-	// first hide spinner if it is there
-	statushooks.Done(ctx)
-	if err := PromptForMissingVariables(ctx, missingVariablesError.MissingVariables, workspacePath); err != nil {
-		log.Printf("[TRACE] Interactive variables prompting returned error %v", err)
-		return nil, err
-	}
-	// ok we should have all variables now - reload workspace
-	return workspace.Load(ctx, workspacePath)
-}

pkg/query/init_data.go

@@ -2,12 +2,11 @@ package query
 
 import (
 	"context"
-
+	"fmt"
 	"github.com/spf13/viper"
 	"github.com/turbot/steampipe/pkg/constants"
 	"github.com/turbot/steampipe/pkg/export"
 	"github.com/turbot/steampipe/pkg/initialisation"
-	"github.com/turbot/steampipe/pkg/statushooks"
 	"github.com/turbot/steampipe/pkg/workspace"
 )
 
@@ -22,13 +21,29 @@ type InitData struct {
 // NewInitData returns a new InitData object
 // It also starts an asynchronous population of the object
 // InitData.Done closes after asynchronous initialization completes
-func NewInitData(ctx context.Context, w *workspace.Workspace, args []string) *InitData {
+func NewInitData(ctx context.Context, args []string) *InitData {
+	// load the workspace
+	w, err := workspace.LoadWorkspacePromptingForVariables(ctx)
+	if err != nil {
+		return &InitData{
+			InitData: *initialisation.NewErrorInitData(fmt.Errorf("failed to load workspace: %s", err.Error())),
+		}
+	}
+
 	i := &InitData{
 		InitData: *initialisation.NewInitData(w),
 		Loaded:   make(chan struct{}),
 	}
 
-	i.RegisterExporters(queryExporters()...)
+	if len(viper.GetStringSlice(constants.ArgExport)) > 0 {
+		i.RegisterExporters(queryExporters()...)
+
+		// validate required export formats
+		if err := i.ExportManager.ValidateExportFormat(viper.GetStringSlice(constants.ArgExport)); err != nil {
+			i.Result.Error = err
+			return i
+		}
+	}
 
 	go i.init(ctx, w, args)
 
@@ -90,8 +105,4 @@ func (i *InitData) init(ctx context.Context, w *workspace.Workspace, args []stri
 	// and call base init
 	i.InitData.Init(ctx, constants.InvokerQuery)
 
-	if len(args) > 0 {
-		// disable status hooks for batch mode
-		ctx = statushooks.DisableStatusHooks(ctx)
-	}
 }

pkg/query/queryexecute/execute.go

@@ -34,9 +34,6 @@ func RunInteractiveSession(ctx context.Context, initData *query.InitData) {
 }
 
 func RunBatchSession(ctx context.Context, initData *query.InitData) int {
-	// ensure we close client
-	defer initData.Cleanup(ctx)
-
 	// 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)

pkg/workspace/load_workspace.go

@@ -1,17 +1,46 @@
-package interactive
+package workspace
 
 import (
 	"context"
 	"fmt"
+	"log"
 
 	"github.com/hashicorp/terraform/terraform"
 	"github.com/spf13/viper"
 	"github.com/turbot/steampipe/pkg/constants"
+	"github.com/turbot/steampipe/pkg/statushooks"
 	"github.com/turbot/steampipe/pkg/steampipeconfig/inputvars"
 	"github.com/turbot/steampipe/pkg/steampipeconfig/modconfig"
 )
 
-func PromptForMissingVariables(ctx context.Context, missingVariables []*modconfig.Variable, workspacePath string) error {
+func LoadWorkspacePromptingForVariables(ctx context.Context) (*Workspace, error) {
+	workspacePath := viper.GetString(constants.ArgModLocation)
+
+	w, err := Load(ctx, workspacePath)
+	if err == nil {
+		return w, nil
+	}
+	missingVariablesError, ok := err.(modconfig.MissingVariableError)
+	// if there was an error which is NOT a MissingVariableError, return it
+	if !ok {
+		return nil, err
+	}
+	// if interactive input is disabled, return the missing variables error
+	if !viper.GetBool(constants.ArgInput) {
+		return nil, missingVariablesError
+	}
+	// so we have missing variables - prompt for them
+	// first hide spinner if it is there
+	statushooks.Done(ctx)
+	if err := promptForMissingVariables(ctx, missingVariablesError.MissingVariables, workspacePath); err != nil {
+		log.Printf("[TRACE] Interactive variables prompting returned error %v", err)
+		return nil, err
+	}
+	// ok we should have all variables now - reload workspace
+	return Load(ctx, workspacePath)
+}
+
+func promptForMissingVariables(ctx context.Context, missingVariables []*modconfig.Variable, workspacePath string) error {
 	fmt.Println()
 	fmt.Println("Variables defined with no value set.")
 	for _, v := range missingVariables {