From f88a826b70712ba0328c27152f76d2487257deec Mon Sep 17 00:00:00 2001 From: Binaek Sarkar Date: Wed, 20 Jan 2021 15:03:30 +0530 Subject: [PATCH] Service management improvements. Closes #5 --- cmd/plugin.go | 15 ++++++------ cmd/service.go | 25 ++++++++++++-------- cmdconfig/cmd_flags.go | 8 +++++++ constants/colors.go | 7 ++++++ db/install.go | 7 +++--- db/interactive_client.go | 4 ++-- db/query.go | 32 +++++++++++++++++--------- db/running_info.go | 1 + db/start.go | 49 ++++++++++++++++++++++++++++------------ go.mod | 1 + go.sum | 1 + 11 files changed, 102 insertions(+), 48 deletions(-) create mode 100644 constants/colors.go diff --git a/cmd/plugin.go b/cmd/plugin.go index 87b8b0bd9b..d791a9673f 100644 --- a/cmd/plugin.go +++ b/cmd/plugin.go @@ -219,7 +219,6 @@ 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") @@ -227,13 +226,13 @@ func getPluginConnectionMap() (map[string][]string, error) { 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) diff --git a/cmd/service.go b/cmd/service.go index 48442020bc..de068f4c83 100644 --- a/cmd/service.go +++ b/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 diff --git a/cmdconfig/cmd_flags.go b/cmdconfig/cmd_flags.go index 8b1c23459d..e269291a19 100644 --- a/cmdconfig/cmd_flags.go +++ b/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 + } +} diff --git a/constants/colors.go b/constants/colors.go new file mode 100644 index 0000000000..ad354ce73d --- /dev/null +++ b/constants/colors.go @@ -0,0 +1,7 @@ +package constants + +import ( + "github.com/fatih/color" +) + +var Bold = color.New(color.Bold).SprintFunc() diff --git a/db/install.go b/db/install.go index 960d9afaae..ed538e9000 100644 --- a/db/install.go +++ b/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 diff --git a/db/interactive_client.go b/db/interactive_client.go index 2afa2e4dc8..ae40345984 100644 --- a/db/interactive_client.go +++ b/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) { diff --git a/db/query.go b/db/query.go index 3ea3dceaa6..11e9b5ed07 100644 --- a/db/query.go +++ b/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,25 +39,29 @@ 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) } @@ -61,14 +69,16 @@ func ExecuteQuery(queryString string) (*ResultStreamer, error) { 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) diff --git a/db/running_info.go b/db/running_info.go index 62400f3b98..6582ee014e 100644 --- a/db/running_info.go +++ b/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 diff --git a/db/start.go b/db/start.go index a7d63f1270..0d6cb8305b 100644 --- a/db/start.go +++ b/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. diff --git a/go.mod b/go.mod index 628c4c42d7..18de07cdf6 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index fb2ddde09b..0f016a6896 100644 --- a/go.sum +++ b/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=