Service management improvements. Closes #5

cmd/plugin.go

@@ -219,21 +219,20 @@ func runPluginRemoveCmd(cmd *cobra.Command, args []string) {
 
 // returns a map of pluginFullName -> []{connections using pluginFullName}
 func getPluginConnectionMap() (map[string][]string, error) {
- didWeStartService := false
  status, err := db.GetStatus()
  if err != nil {
  return nil, fmt.Errorf("Could not start steampipe service")
  }
 
  if status == nil {
  // the db service is not started - start it
- db.StartService()
- didWeStartService = true
- }
-
- if didWeStartService {
- // stop the database if we started it!
- defer db.StopDB(true)
+ db.StartService(db.InvokerPlugin)
+ defer func() {
+  status, _ := db.GetStatus()
+ if status.Invoker == db.InvokerPlugin {
+  db.StopDB(true)
+  }
+ }()
  }
 
  client, err := db.GetClient(true)

cmd/service.go

@@ -5,10 +5,10 @@ import (
  "fmt"
  "strings"
 
- "github.com/turbot/steampipe/cmdconfig"
-
  "github.com/spf13/cobra"
  "github.com/turbot/steampipe-plugin-sdk/logging"
+
+ "github.com/turbot/steampipe/cmdconfig"
  "github.com/turbot/steampipe/constants"
  "github.com/turbot/steampipe/db"
  "github.com/turbot/steampipe/utils"
@@ -55,11 +55,11 @@ connection from any Postgres compatible database client.`,
  cmdconfig.
  OnCmd(cmd).
  AddBoolFlag("background", "", true, "Run service in the background").
- AddBoolFlag("refresh", "", true, "Refresh connections").
- // TODO(nw) default to the configuration option?
- AddIntFlag("db-port", "", constants.DatabasePort, "Database service port. Chooses a free port by default.").
- // TODO(nw) should be validated to an enumerated list
- AddStringFlag("listen", "", "network", "Accept connections from: local (localhost only) or network (open)")
+ AddIntFlag("db-port", "", constants.DatabasePort, "Database service port.").
+ AddStringFlag("listen", "", string(db.ListenTypeNetwork), "Accept connections from: local (localhost only) or network (open)").
+ // Hidden flags for internal use
+ AddStringFlag("invoker", "", string(db.InvokerService), "Invoked by `service` or `query`", cmdconfig.FlagOptions.Hidden()).
+ AddBoolFlag("refresh", "", true, "Refresh connections on startup", cmdconfig.FlagOptions.Hidden())
 
  return cmd
 }
@@ -124,15 +124,20 @@ func runServiceStartCmd(cmd *cobra.Command, args []string) {
  }
 
  listen := db.StartListenType(cmdconfig.Viper().GetString("listen"))
-
  if err := listen.IsValid(); err != nil {
  utils.ShowError(err)
  return
  }
 
+ invoker := db.Invoker(cmdconfig.Viper().GetString("invoker"))
+ if err := invoker.IsValid(); err != nil {
+ utils.ShowError(err)
+ return
+ }
+
  db.EnsureDBInstalled()
 
- status, err := db.StartDB(cmdconfig.Viper().GetInt("db-port"), listen)
+ status, err := db.StartDB(cmdconfig.Viper().GetInt("db-port"), listen, invoker)
  if err != nil {
  utils.ShowError(err)
  return
@@ -189,7 +194,7 @@ to force a restart.
  return
  }
 
- status, err := db.StartDB(currentServiceStatus.Port, currentServiceStatus.ListenType)
+ status, err := db.StartDB(currentServiceStatus.Port, currentServiceStatus.ListenType, currentServiceStatus.Invoker)
  if err != nil {
  utils.ShowError(err)
  return

cmdconfig/cmd_flags.go

@@ -15,8 +15,10 @@ type flagOpt func(c *cobra.Command, name string, key string, v *viper.Viper)
 // FlagOptions :: shortcut for common flag options
 var FlagOptions = struct {
  Required func() flagOpt
+ Hidden func() flagOpt
 }{
  Required: requiredOpt,
+ Hidden: hiddenOpt,
 }
 
 // Helper function to mark a flag as required
@@ -29,3 +31,9 @@ func requiredOpt() flagOpt {
  c.Flag(name).Usage = fmt.Sprintf("%s %s", u, requiredColor("(required)"))
  }
 }
+
+func hiddenOpt() flagOpt {
+ return func(c *cobra.Command, name, key string, v *viper.Viper) {
+ c.Flag(name).Hidden = true
+ }
+}

constants/colors.go

@@ -0,0 +1,7 @@
+package constants
+
+import (
+ "github.com/fatih/color"
+)
+
+var Bold = color.New(color.Bold).SprintFunc()

db/install.go

@@ -78,7 +78,7 @@ func EnsureDBInstalled() {
 }
 
 func installSteampipeHub() error {
- StartService()
+ StartService(InvokerInstaller)
  rawClient, err := createDbClient()
  if err != nil {
  return err
@@ -105,11 +105,12 @@ func installSteampipeHub() error {
  return nil
 }
 
-func StartService() {
+// StartService :: invokes `steampipe service start --listen local --refresh=false --invoker query`
+func StartService(invoker Invoker) {
  log.Println("[TRACE] start service")
  // spawn a process to start the service, passing refresh=false to ensure we DO NOT refresh connections
  // (as we will do that ourselves)
- cmd := exec.Command(os.Args[0], "service", "start", "--listen", "local", "--refresh=false")
+ cmd := exec.Command(os.Args[0], "service", "start", "--listen", "local", "--refresh=false", "--invoker", string(invoker))
  cmd.Start()
  startedAt := time.Now()
  spinnerShown := false

db/interactive_client.go

@@ -41,10 +41,10 @@ func (c *InteractiveClient) close() {
 }
 
 // InteractiveQuery :: start an interactive prompt and return
-func (c *InteractiveClient) InteractiveQuery(resultsStreamer *ResultStreamer, stopServiceWhenDone bool) {
+func (c *InteractiveClient) InteractiveQuery(resultsStreamer *ResultStreamer, onCompleteCallback func()) {
  defer func() {
 
- shutdown(c.client, stopServiceWhenDone)
+ onCompleteCallback()
 
  r := recover()
  switch r.(type) {

db/query.go

@@ -6,12 +6,13 @@ import (
  "log"
 
  "github.com/turbot/steampipe-plugin-sdk/logging"
+
+ "github.com/turbot/steampipe/constants"
  "github.com/turbot/steampipe/utils"
 )
 
 // ExecuteQuery :: entry point for executing ad-hoc queries from outside the package
 func ExecuteQuery(queryString string) (*ResultStreamer, error) {
- var didWeStartService bool
  var err error
 
  logging.LogTime("db.ExecuteQuery start")
@@ -23,10 +24,13 @@ func ExecuteQuery(queryString string) (*ResultStreamer, error) {
  return nil, errors.New("could not retrieve service status")
  }
 
+ if status != nil && status.Invoker == InvokerQuery {
+ return nil, fmt.Errorf("You already have a %s session open. To run multiple sessions, first run %s", constants.Bold("steampipe query"), constants.Bold("steampipe service start"))
+ }
+
  if status == nil {
  // the db service is not started - start it
- StartService()
- didWeStartService = true
+ StartService(InvokerQuery)
  }
 
  client, err := GetClient(false)
@@ -35,40 +39,46 @@ func ExecuteQuery(queryString string) (*ResultStreamer, error) {
  // refresh connections
  if err = refreshConnections(client); err != nil {
  // shutdown the service if something went wrong!!!
- shutdown(client, didWeStartService)
+ shutdown(client)
  return nil, fmt.Errorf("failed to refresh connections: %v", err)
  }
 
  resultsStreamer := newQueryResults()
+
+ // this is a callback to close the db et-al. when things get done - no matter the mode
+ onComplete := func() { shutdown(client) }
+
  if queryString == "" {
  interactiveClient, err := newInteractiveClient(client)
  utils.FailOnErrorWithMessage(err, "interactive client failed to initialize")
 
  // start the interactive prompt in a go routine
- go interactiveClient.InteractiveQuery(resultsStreamer, didWeStartService)
+ go interactiveClient.InteractiveQuery(resultsStreamer, onComplete)
  } else {
  result, err := client.executeQuery(queryString)
  if err != nil {
  return nil, err
  }
  // send a single result to the streamer - this will close the channel afterwards
- // pass an onComplete callback function to shutdown the db
- onComplete := func() { shutdown(client, didWeStartService) }
+ // pass the onComplete callback function to shutdown the db
+ onComplete := func() { shutdown(client) }
  go resultsStreamer.streamSingleResult(result, onComplete)
  }
 
  logging.LogTime("db.ExecuteQuery end")
  return resultsStreamer, nil
 }
 
-func shutdown(client *Client, stopService bool) {
- log.Println("[TRACE] shutdown", stopService)
+func shutdown(client *Client) {
+ log.Println("[TRACE] shutdown")
  if client != nil {
  client.close()
  }
 
- // force stop
- if stopService {
+ status, _ := GetStatus()
+
+ // force stop if invoked by `query` and we are the last one
+ if status.Invoker == InvokerQuery {
  _, err := StopDB(true)
  if err != nil {
  utils.ShowError(err)

db/running_info.go

@@ -19,6 +19,7 @@ type RunningDBInstanceInfo struct {
  Port int
  Listen []string
  ListenType StartListenType
+ Invoker Invoker
  Password string
  User string
  Database string

db/start.go

@@ -19,6 +19,9 @@ type StartResult int
 // StartListenType :: pseudoEnum of network binding for postgres
 type StartListenType string
 
+// Invoker :: pseudoEnum for what starts the service
+type Invoker string
+
 const (
  // ServiceStarted :: StartResult - Service was started
  ServiceStarted StartResult = iota
@@ -29,29 +32,43 @@ const (
 )
 
 const (
- // NetworkListenType :: StartListenType - bind to all known interfaces
- NetworkListenType StartListenType = "network"
- // LocalListenType :: StartListenType - bind to localhost only
- LocalListenType = "local"
+ // ListenTypeNetwork :: StartListenType - bind to all known interfaces
+ ListenTypeNetwork StartListenType = "network"
+ // ListenTypeLocal :: StartListenType - bind to localhost only
+ ListenTypeLocal = "local"
+)
+
+const (
+ // InvokerService :: Invoker - when invoked by `service start`
+ InvokerService Invoker = "service"
+ // InvokerQuery :: Invoker - when invoked by `query`
+ InvokerQuery = "query"
+ // InvokerInstaller :: Invoker - when invoked by the `installer`
+ InvokerInstaller = "installer"
+ // InvokerPlugin :: Invoker - when invoked by the `pluginmanager`
+ InvokerPlugin = "plugin"
 )
 
 // IsValid :: validator for StartListenType known values
 func (slt StartListenType) IsValid() error {
  switch slt {
- case NetworkListenType, LocalListenType:
+ case ListenTypeNetwork, ListenTypeLocal:
  return nil
  }
- return fmt.Errorf("Invalid listen type. Can be one of '%v' or '%v'", NetworkListenType, LocalListenType)
+ return fmt.Errorf("Invalid listen type. Can be one of '%v' or '%v'", ListenTypeNetwork, ListenTypeLocal)
 }
 
-// StartDBWithDefaults :: start the database, if not already running
-// on `127.0.0.1:9193`
-func StartDBWithDefaults() (StartResult, error) {
- return StartDB(constants.DatabasePort, "local")
+// IsValid :: validator for Invoker known values
+func (slt Invoker) IsValid() error {
+ switch slt {
+ case InvokerService, InvokerQuery, InvokerInstaller, InvokerPlugin:
+ return nil
+ }
+ return fmt.Errorf("Invalid invoker. Can be one of '%v', '%v', '%v' or '%v'", InvokerService, InvokerQuery, InvokerInstaller, InvokerPlugin)
 }
 
 // StartDB :: start the database is not already running
-func StartDB(port int, listen StartListenType) (StartResult, error) {
+func StartDB(port int, listen StartListenType, invoker Invoker) (StartResult, error) {
  info, err := loadRunningInstanceInfo()
 
  if err != nil {
@@ -64,6 +81,9 @@ func StartDB(port int, listen StartListenType) (StartResult, error) {
  return ServiceFailedToStart, err
  }
  if processRunning {
+ if info.Invoker == InvokerQuery {
+ return ServiceAlreadyRunning, fmt.Errorf("You have a %s session open. Close this session before running %s", constants.Bold("steampipe query"), constants.Bold("steampipe service start"))
+ }
  return ServiceAlreadyRunning, nil
  }
  }
@@ -75,7 +95,7 @@ func StartDB(port int, listen StartListenType) (StartResult, error) {
 
  listenAddresses := "localhost"
 
- if listen == NetworkListenType {
+ if listen == ListenTypeNetwork {
  listenAddresses = "*"
  }
 
@@ -139,9 +159,10 @@ func StartDB(port int, listen StartListenType) (StartResult, error) {
  runningInfo.User = constants.DatabaseSuperUser
  runningInfo.Database = "postgres"
  runningInfo.ListenType = listen
+ runningInfo.Invoker = invoker
 
  runningInfo.Listen = []string{"localhost", "127.0.0.1"}
- if listen == NetworkListenType {
+ if listen == ListenTypeNetwork {
  addrs, _ := localAddresses()
  runningInfo.Listen = append(addrs, runningInfo.Listen...)
  }
@@ -164,7 +185,7 @@ func StartDB(port int, listen StartListenType) (StartResult, error) {
  // remove info file (if any)
  _ = removeRunningInstanceInfo()
  // try restarting
- return StartDB(port, listen)
+ return StartDB(port, listen, invoker)
  }
 
  // there was nothing to kill.

go.mod

@@ -14,6 +14,7 @@ require (
  github.com/fatih/color v1.7.0
  github.com/gertd/go-pluralize v0.1.7
  github.com/go-ole/go-ole v1.2.4 // indirect
+ github.com/google/uuid v1.1.2
  github.com/hashicorp/go-hclog v0.14.1
  github.com/hashicorp/go-plugin v1.3.0
  github.com/hashicorp/go-version v1.2.1

go.sum

@@ -186,6 +186,7 @@ github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXi
 github.com/google/pprof v0.0.0-20181206194817-3ea8567a2e57/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
 github.com/google/pprof v0.0.0-20190515194954-54271f7e092f/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
+github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=