github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/query/queryexecute/execute.go (about)

     1  package queryexecute
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/spf13/viper"
     9  	"github.com/turbot/steampipe/pkg/cmdconfig"
    10  	"github.com/turbot/steampipe/pkg/connection_sync"
    11  	"github.com/turbot/steampipe/pkg/constants"
    12  	"github.com/turbot/steampipe/pkg/contexthelpers"
    13  	"github.com/turbot/steampipe/pkg/db/db_common"
    14  	"github.com/turbot/steampipe/pkg/display"
    15  	"github.com/turbot/steampipe/pkg/error_helpers"
    16  	"github.com/turbot/steampipe/pkg/interactive"
    17  	"github.com/turbot/steampipe/pkg/query"
    18  	"github.com/turbot/steampipe/pkg/steampipeconfig/modconfig"
    19  	"github.com/turbot/steampipe/pkg/utils"
    20  )
    21  
    22  func RunInteractiveSession(ctx context.Context, initData *query.InitData) error {
    23  	utils.LogTime("execute.RunInteractiveSession start")
    24  	defer utils.LogTime("execute.RunInteractiveSession end")
    25  
    26  	// the db executor sends result data over resultsStreamer
    27  	result := interactive.RunInteractivePrompt(ctx, initData)
    28  
    29  	// print the data as it comes
    30  	for r := range result.Streamer.Results {
    31  		display.ShowOutput(ctx, r)
    32  		// signal to the resultStreamer that we are done with this chunk of the stream
    33  		result.Streamer.AllResultsRead()
    34  	}
    35  	return result.PromptErr
    36  }
    37  
    38  func RunBatchSession(ctx context.Context, initData *query.InitData) (int, error) {
    39  	// start cancel handler to intercept interrupts and cancel the context
    40  	// NOTE: use the initData Cancel function to ensure any initialisation is cancelled if needed
    41  	contexthelpers.StartCancelHandler(initData.Cancel)
    42  
    43  	// wait for init
    44  	<-initData.Loaded
    45  
    46  	if err := initData.Result.Error; err != nil {
    47  		return 0, err
    48  	}
    49  
    50  	// display any initialisation messages/warnings
    51  	initData.Result.DisplayMessages()
    52  
    53  	// if there is a custom search path, wait until the first connection of each plugin has loaded
    54  	if customSearchPath := initData.Client.GetCustomSearchPath(); customSearchPath != nil {
    55  		if err := connection_sync.WaitForSearchPathSchemas(ctx, initData.Client, customSearchPath); err != nil {
    56  			return 0, err
    57  		}
    58  	}
    59  
    60  	failures := 0
    61  	if len(initData.Queries) > 0 {
    62  		// if we have resolved any queries, run them
    63  		failures = executeQueries(ctx, initData)
    64  	}
    65  	// return the number of query failures and the number of rows that returned errors
    66  	return failures, nil
    67  }
    68  
    69  func executeQueries(ctx context.Context, initData *query.InitData) int {
    70  	utils.LogTime("queryexecute.executeQueries start")
    71  	defer utils.LogTime("queryexecute.executeQueries end")
    72  
    73  	// failures return the number of queries that failed and also the number of rows that
    74  	// returned errors
    75  	failures := 0
    76  	t := time.Now()
    77  
    78  	var err error
    79  
    80  	for i, q := range initData.Queries {
    81  		// if executeQuery fails it returns err, else it returns the number of rows that returned errors while execution
    82  		if err, failures = executeQuery(ctx, initData.Client, q); err != nil {
    83  			failures++
    84  			error_helpers.ShowWarning(fmt.Sprintf("executeQueries: query %d of %d failed: %v", i+1, len(initData.Queries), error_helpers.DecodePgError(err)))
    85  			// if timing flag is enabled, show the time taken for the query to fail
    86  			if cmdconfig.Viper().GetString(constants.ArgTiming) != constants.ArgOff {
    87  				display.DisplayErrorTiming(t)
    88  			}
    89  		}
    90  		// TODO move into display layer
    91  		// Only show the blank line between queries, not after the last one
    92  		if (i < len(initData.Queries)-1) && showBlankLineBetweenResults() {
    93  			fmt.Println()
    94  		}
    95  	}
    96  
    97  	return failures
    98  }
    99  
   100  func executeQuery(ctx context.Context, client db_common.Client, resolvedQuery *modconfig.ResolvedQuery) (error, int) {
   101  	utils.LogTime("query.execute.executeQuery start")
   102  	defer utils.LogTime("query.execute.executeQuery end")
   103  
   104  	// the db executor sends result data over resultsStreamer
   105  	resultsStreamer, err := db_common.ExecuteQuery(ctx, client, resolvedQuery.ExecuteSQL, resolvedQuery.Args...)
   106  	if err != nil {
   107  		return err, 0
   108  	}
   109  
   110  	rowErrors := 0 // get the number of rows that returned an error
   111  	// print the data as it comes
   112  	for r := range resultsStreamer.Results {
   113  		rowErrors = display.ShowOutput(ctx, r)
   114  		// signal to the resultStreamer that we are done with this result
   115  		resultsStreamer.AllResultsRead()
   116  	}
   117  	return nil, rowErrors
   118  }
   119  
   120  // if we are displaying csv with no header, do not include lines between the query results
   121  func showBlankLineBetweenResults() bool {
   122  	return !(viper.GetString(constants.ArgOutput) == "csv" && !viper.GetBool(constants.ArgHeader))
   123  }