Maintain the order of execution when running multiple queries in batc…

…h mode. Closes #3728
turbot · Jan 2, 2024 · 085e933 · 085e933
1 parent c43a8d2
commit 085e933
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 69 deletions.
diff --git a/cmd/query.go b/cmd/query.go
@@ -220,67 +220,60 @@ func executeSnapshotQuery(initData *query.InitData, ctx context.Context) int {
 		}
 	}
 
-	// build ordered list of queries
-	// (ordered for testing repeatability)
-	var queryNames = utils.SortedMapKeys(initData.Queries)
-
-	if len(queryNames) > 0 {
-		for _, name := range queryNames {
-			resolvedQuery := initData.Queries[name]
-			// if a manual query is being run (i.e. not a named query), convert into a query and add to workspace
-			// this is to allow us to use existing dashboard execution code
-			queryProvider, existingResource := ensureSnapshotQueryResource(name, resolvedQuery, initData.Workspace)
-
-			// we need to pass the embedded initData to  GenerateSnapshot
-			baseInitData := &initData.InitData
-
-			// so a dashboard name was specified - just call GenerateSnapshot
-			snap, err := dashboardexecute.GenerateSnapshot(ctx, queryProvider.Name(), baseInitData, nil)
-			if err != nil {
-				exitCode = constants.ExitCodeSnapshotCreationFailed
-				error_helpers.FailOnError(err)
-			}
+	for _, resolvedQuery := range initData.Queries {
+		// if a manual query is being run (i.e. not a named query), convert into a query and add to workspace
+		// this is to allow us to use existing dashboard execution code
+		queryProvider, existingResource := ensureSnapshotQueryResource(resolvedQuery.Name, resolvedQuery, initData.Workspace)
+
+		// we need to pass the embedded initData to  GenerateSnapshot
+		baseInitData := &initData.InitData
+
+		// so a dashboard name was specified - just call GenerateSnapshot
+		snap, err := dashboardexecute.GenerateSnapshot(ctx, queryProvider.Name(), baseInitData, nil)
+		if err != nil {
+			exitCode = constants.ExitCodeSnapshotCreationFailed
+			error_helpers.FailOnError(err)
+		}
 
-			// set the filename root for the snapshot (in case needed)
-			if !existingResource {
-				snap.FileNameRoot = "query"
-			}
+		// set the filename root for the snapshot (in case needed)
+		if !existingResource {
+			snap.FileNameRoot = "query"
+		}
 
-			// display the result
-			switch viper.GetString(constants.ArgOutput) {
-			case constants.OutputFormatNone:
-				// do nothing
-			case constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort:
-				// if the format is snapshot, just dump it out
-				jsonOutput, err := json.MarshalIndent(snap, "", "  ")
-				if err != nil {
-					error_helpers.FailOnErrorWithMessage(err, "failed to display result as snapshot")
-				}
-				fmt.Println(string(jsonOutput))
-			default:
-				// otherwise convert the snapshot into a query result
-				result, err := snapshotToQueryResult(snap)
+		// display the result
+		switch viper.GetString(constants.ArgOutput) {
+		case constants.OutputFormatNone:
+			// do nothing
+		case constants.OutputFormatSnapshot, constants.OutputFormatSnapshotShort:
+			// if the format is snapshot, just dump it out
+			jsonOutput, err := json.MarshalIndent(snap, "", "  ")
+			if err != nil {
 				error_helpers.FailOnErrorWithMessage(err, "failed to display result as snapshot")
-				display.ShowOutput(ctx, result, display.WithTimingDisabled())
 			}
+			fmt.Println(string(jsonOutput))
+		default:
+			// otherwise convert the snapshot into a query result
+			result, err := snapshotToQueryResult(snap)
+			error_helpers.FailOnErrorWithMessage(err, "failed to display result as snapshot")
+			display.ShowOutput(ctx, result, display.WithTimingDisabled())
+		}
 
-			// share the snapshot if necessary
-			err = publishSnapshotIfNeeded(ctx, snap)
-			if err != nil {
-				exitCode = constants.ExitCodeSnapshotUploadFailed
-				error_helpers.FailOnErrorWithMessage(err, fmt.Sprintf("failed to publish snapshot to %s", viper.GetString(constants.ArgSnapshotLocation)))
-			}
+		// share the snapshot if necessary
+		err = publishSnapshotIfNeeded(ctx, snap)
+		if err != nil {
+			exitCode = constants.ExitCodeSnapshotUploadFailed
+			error_helpers.FailOnErrorWithMessage(err, fmt.Sprintf("failed to publish snapshot to %s", viper.GetString(constants.ArgSnapshotLocation)))
+		}
 
-			// export the result if necessary
-			exportArgs := viper.GetStringSlice(constants.ArgExport)
-			exportMsg, err := initData.ExportManager.DoExport(ctx, snap.FileNameRoot, snap, exportArgs)
-			error_helpers.FailOnErrorWithMessage(err, "failed to export snapshot")
-			// print the location where the file is exported
-			if len(exportMsg) > 0 && viper.GetBool(constants.ArgProgress) {
-				fmt.Printf("\n")
-				fmt.Println(strings.Join(exportMsg, "\n"))
-				fmt.Printf("\n")
-			}
+		// export the result if necessary
+		exportArgs := viper.GetStringSlice(constants.ArgExport)
+		exportMsg, err := initData.ExportManager.DoExport(ctx, snap.FileNameRoot, snap, exportArgs)
+		error_helpers.FailOnErrorWithMessage(err, "failed to export snapshot")
+		// print the location where the file is exported
+		if len(exportMsg) > 0 && viper.GetBool(constants.ArgProgress) {
+			fmt.Printf("\n")
+			fmt.Println(strings.Join(exportMsg, "\n"))
+			fmt.Printf("\n")
 		}
 	}
 	return 0

diff --git a/pkg/query/init_data.go b/pkg/query/init_data.go
@@ -23,7 +23,7 @@ type InitData struct {
 	cancelInitialisation context.CancelFunc
 	Loaded               chan struct{}
 	// map of query name to resolved query (key is the query text for command line queries)
-	Queries map[string]*modconfig.ResolvedQuery
+	Queries []*modconfig.ResolvedQuery
 }
 
 // NewInitData returns a new InitData object

diff --git a/pkg/query/queryexecute/execute.go b/pkg/query/queryexecute/execute.go
@@ -74,25 +74,22 @@ func executeQueries(ctx context.Context, initData *query.InitData) int {
 	// returned errors
 	failures := 0
 	t := time.Now()
-	// build ordered list of queries
-	// (ordered for testing repeatability)
-	var queryNames = utils.SortedMapKeys(initData.Queries)
+
 	var err error
 
-	for i, name := range queryNames {
-		q := initData.Queries[name]
+	for i, q := range initData.Queries {
 		// if executeQuery fails it returns err, else it returns the number of rows that returned errors while execution
 		if err, failures = executeQuery(ctx, initData.Client, q); err != nil {
 			failures++
-			error_helpers.ShowWarning(fmt.Sprintf("executeQueries: query %d of %d failed: %v", i+1, len(queryNames), error_helpers.DecodePgError(err)))
+			error_helpers.ShowWarning(fmt.Sprintf("executeQueries: query %d of %d failed: %v", i+1, len(initData.Queries), error_helpers.DecodePgError(err)))
 			// if timing flag is enabled, show the time taken for the query to fail
 			if cmdconfig.Viper().GetBool(constants.ArgTiming) {
 				display.DisplayErrorTiming(t)
 			}
 		}
 		// TODO move into display layer
 		// Only show the blank line between queries, not after the last one
-		if (i < len(queryNames)-1) && showBlankLineBetweenResults() {
+		if (i < len(initData.Queries)-1) && showBlankLineBetweenResults() {
 			fmt.Println()
 		}
 	}

diff --git a/pkg/steampipeconfig/modconfig/resolved_query.go b/pkg/steampipeconfig/modconfig/resolved_query.go
@@ -6,6 +6,7 @@ import (
 
 // ResolvedQuery contains the execute SQL, raw SQL and args string used to execute a query
 type ResolvedQuery struct {
+	Name        string
 	ExecuteSQL  string
 	RawSQL      string
 	Args        []any

diff --git a/pkg/workspace/workspace_queries.go b/pkg/workspace/workspace_queries.go
@@ -18,23 +18,23 @@ import (
 // GetQueriesFromArgs retrieves queries from args
 //
 // For each arg check if it is a named query or a file, before falling back to treating it as sql
-func (w *Workspace) GetQueriesFromArgs(args []string) (map[string]*modconfig.ResolvedQuery, error) {
+func (w *Workspace) GetQueriesFromArgs(args []string) ([]*modconfig.ResolvedQuery, error) {
 	utils.LogTime("execute.GetQueriesFromArgs start")
 	defer utils.LogTime("execute.GetQueriesFromArgs end")
 
-	var queries = make(map[string]*modconfig.ResolvedQuery)
-	for _, arg := range args {
+	var queries = make([]*modconfig.ResolvedQuery, len(args))
+	for idx, arg := range args {
 		resolvedQuery, queryProvider, err := w.ResolveQueryAndArgsFromSQLString(arg)
 		if err != nil {
 			return nil, err
 		}
 		if len(resolvedQuery.ExecuteSQL) > 0 {
 			// default name to the query text
-			queryName := resolvedQuery.ExecuteSQL
+			resolvedQuery.Name = resolvedQuery.ExecuteSQL
 			if queryProvider != nil {
-				queryName = queryProvider.Name()
+				resolvedQuery.Name = queryProvider.Name()
 			}
-			queries[queryName] = resolvedQuery
+			queries[idx] = resolvedQuery
 		}
 	}
 	return queries, nil