github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/cmd/plugin_manager.go (about)

     1  package cmd
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"os"
     7  	"os/signal"
     8  	"strings"
     9  	"syscall"
    10  	"time"
    11  
    12  	"github.com/hashicorp/go-hclog"
    13  	"github.com/spf13/cobra"
    14  	"github.com/turbot/go-kit/helpers"
    15  	"github.com/turbot/go-kit/logging"
    16  	"github.com/turbot/go-kit/types"
    17  	sdklogging "github.com/turbot/steampipe-plugin-sdk/v5/logging"
    18  	"github.com/turbot/steampipe-plugin-sdk/v5/plugin"
    19  	"github.com/turbot/steampipe/pkg/cmdconfig"
    20  	"github.com/turbot/steampipe/pkg/connection"
    21  	"github.com/turbot/steampipe/pkg/constants"
    22  	"github.com/turbot/steampipe/pkg/filepaths"
    23  	"github.com/turbot/steampipe/pkg/pluginmanager_service"
    24  	"github.com/turbot/steampipe/pkg/steampipeconfig"
    25  )
    26  
    27  func pluginManagerCmd() *cobra.Command {
    28  	cmd := &cobra.Command{
    29  		Use:    "plugin-manager",
    30  		Run:    runPluginManagerCmd,
    31  		Hidden: true,
    32  	}
    33  	cmdconfig.OnCmd(cmd)
    34  	return cmd
    35  }
    36  
    37  func runPluginManagerCmd(cmd *cobra.Command, _ []string) {
    38  	var err error
    39  	defer func() {
    40  		if r := recover(); r != nil {
    41  			err = helpers.ToError(r)
    42  		}
    43  		if err != nil {
    44  			// write to stdout so the plugin manager can extract the error message
    45  			fmt.Println(fmt.Sprintf("%s%s", plugin.PluginStartupFailureMessage, err.Error()))
    46  		}
    47  		os.Exit(1)
    48  	}()
    49  
    50  	err = doRunPluginManager(cmd)
    51  }
    52  
    53  func doRunPluginManager(cmd *cobra.Command) error {
    54  	pluginManager, err := createPluginManager(cmd)
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	if shouldRunConnectionWatcher() {
    60  		log.Printf("[INFO] starting connection watcher")
    61  		connectionWatcher, err := connection.NewConnectionWatcher(pluginManager)
    62  		if err != nil {
    63  			return err
    64  		}
    65  
    66  		// close the connection watcher
    67  		defer connectionWatcher.Close()
    68  	}
    69  
    70  	log.Printf("[INFO] about to serve")
    71  	pluginManager.Serve()
    72  	return nil
    73  }
    74  
    75  func createPluginManager(cmd *cobra.Command) (*pluginmanager_service.PluginManager, error) {
    76  	ctx := cmd.Context()
    77  	logger := createPluginManagerLog()
    78  
    79  	log.Printf("[INFO] starting plugin manager")
    80  	// build config map
    81  	steampipeConfig, errorsAndWarnings := steampipeconfig.LoadConnectionConfig(ctx)
    82  	if errorsAndWarnings.GetError() != nil {
    83  		log.Printf("[WARN] failed to load connection config: %v", errorsAndWarnings.GetError())
    84  		return nil, errorsAndWarnings.Error
    85  	}
    86  
    87  	// add signal handler for sigpipe - this will be raised if we call displayWarning as stdout is piped
    88  	signalCh := make(chan os.Signal, 1)
    89  	signal.Notify(signalCh, syscall.SIGPIPE)
    90  	go func() {
    91  		for {
    92  			// swallow signal
    93  			<-signalCh
    94  		}
    95  	}()
    96  
    97  	// create a map of connections configs, excluding connections in error
    98  	configMap := connection.NewConnectionConfigMap(steampipeConfig.Connections)
    99  	log.Printf("[TRACE] loaded config map: %s", strings.Join(steampipeConfig.ConnectionNames(), ","))
   100  
   101  	pluginManager, err := pluginmanager_service.NewPluginManager(ctx, configMap, steampipeConfig.PluginsInstances, logger)
   102  	if err != nil {
   103  		log.Printf("[WARN] failed to create plugin manager: %s", err.Error())
   104  		return nil, err
   105  	}
   106  
   107  	return pluginManager, nil
   108  }
   109  
   110  func shouldRunConnectionWatcher() bool {
   111  	// if EnvConnectionWatcher is set, overwrite the value in DefaultConnectionOptions
   112  	if envStr, ok := os.LookupEnv(constants.EnvConnectionWatcher); ok {
   113  		if parsedEnv, err := types.ToBool(envStr); err == nil {
   114  			return parsedEnv
   115  		}
   116  	}
   117  	return true
   118  }
   119  
   120  func createPluginManagerLog() hclog.Logger {
   121  	// we use this logger to log from the plugin processes
   122  	// the plugin processes uses the `EscapeNewlineWriter` to map the '\n' byte to "\n" string literal
   123  	// this is to allow the plugin to send multiline log messages as a single log line.
   124  	//
   125  	// here we apply the reverse mapping to get back the original message
   126  	writer := sdklogging.NewUnescapeNewlineWriter(logging.NewRotatingLogWriter(filepaths.EnsureLogDir(), "plugin"))
   127  
   128  	logger := sdklogging.NewLogger(&hclog.LoggerOptions{
   129  		Output:     writer,
   130  		TimeFn:     func() time.Time { return time.Now().UTC() },
   131  		TimeFormat: "2006-01-02 15:04:05.000 UTC",
   132  	})
   133  	log.SetOutput(logger.StandardWriter(&hclog.StandardLoggerOptions{InferLevels: true}))
   134  	log.SetPrefix("")
   135  	log.SetFlags(0)
   136  	return logger
   137  }