diff --git a/cmd/backup.go b/cmd/backup.go index 418f6a720..adc1cfcd8 100644 --- a/cmd/backup.go +++ b/cmd/backup.go @@ -3,40 +3,37 @@ package cmd import ( "context" "fmt" - "log" + "io" "os" "github.com/spf13/cobra" - "github.com/spf13/viper" "github.com/pilosa/pilosa/ctl" ) -var backuper = ctl.NewBackupCommand(os.Stdin, os.Stdout, os.Stderr) - -var backupCmd = &cobra.Command{ - Use: "backup", - Short: "Backup data from pilosa.", - Long: ` +func NewBackupCmd(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + backuper := ctl.NewBackupCommand(os.Stdin, os.Stdout, os.Stderr) + backupCmd := &cobra.Command{ + Use: "backup", + Short: "Backup data from pilosa.", + Long: ` Backs up the database and frame from across the cluster into a single file. `, - Run: func(cmd *cobra.Command, args []string) { - if err := backuper.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if err := backuper.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + flags := backupCmd.Flags() + flags.StringVarP(&backuper.Host, "host", "", "localhost:15000", "host:port of Pilosa.") + flags.StringVarP(&backuper.Database, "database", "d", "", "Pilosa database to backup into.") + flags.StringVarP(&backuper.Frame, "frame", "f", "", "Frame to backup into.") + flags.StringVarP(&backuper.Path, "output-file", "o", "", "File to write backup to - default stdout") + + return backupCmd } func init() { - backupCmd.Flags().StringVarP(&backuper.Host, "host", "", "localhost:15000", "host:port of Pilosa.") - backupCmd.Flags().StringVarP(&backuper.Database, "database", "d", "", "Pilosa database to backup into.") - backupCmd.Flags().StringVarP(&backuper.Frame, "frame", "f", "", "Frame to backup into.") - backupCmd.Flags().StringVarP(&backuper.Path, "output-file", "o", "", "File to write backup to - default stdout") - - err := viper.BindPFlags(backupCmd.Flags()) - if err != nil { - log.Fatalf("Error binding backup flags: %v", err) - } - - RootCmd.AddCommand(backupCmd) + subcommandFns["backup"] = NewBackupCmd } diff --git a/cmd/bench.go b/cmd/bench.go index 3e093d68e..a1d89637d 100644 --- a/cmd/bench.go +++ b/cmd/bench.go @@ -3,41 +3,38 @@ package cmd import ( "context" "fmt" - "log" + "io" "os" "github.com/spf13/cobra" - "github.com/spf13/viper" "github.com/pilosa/pilosa/ctl" ) -var bencher = ctl.NewBenchCommand(os.Stdin, os.Stdout, os.Stderr) - -var benchCmd = &cobra.Command{ - Use: "bench", - Short: "Benchmark operations.", - Long: ` +func NewBenchCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + bencher := ctl.NewBenchCommand(os.Stdin, os.Stdout, os.Stderr) + benchCmd := &cobra.Command{ + Use: "bench", + Short: "Benchmark operations.", + Long: ` Executes a benchmark for a given operation against the database. `, - Run: func(cmd *cobra.Command, args []string) { - if err := bencher.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if err := bencher.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + flags := benchCmd.Flags() + flags.StringVarP(&bencher.Host, "host", "", "localhost:15000", "host:port of Pilosa.") + flags.StringVarP(&bencher.Database, "database", "d", "", "Pilosa database to benchmark.") + flags.StringVarP(&bencher.Frame, "frame", "f", "", "Frame to benchmark.") + flags.StringVarP(&bencher.Op, "operation", "o", "set-bit", "Operation to perform: choose from [set-bit]") + flags.IntVarP(&bencher.N, "num", "n", 0, "Number of operations to perform.") + + return benchCmd } func init() { - benchCmd.Flags().StringVarP(&bencher.Host, "host", "", "localhost:15000", "host:port of Pilosa.") - benchCmd.Flags().StringVarP(&bencher.Database, "database", "d", "", "Pilosa database to benchmark.") - benchCmd.Flags().StringVarP(&bencher.Frame, "frame", "f", "", "Frame to benchmark.") - benchCmd.Flags().StringVarP(&bencher.Op, "operation", "o", "set-bit", "Operation to perform: choose from [set-bit]") - benchCmd.Flags().IntVarP(&bencher.N, "num", "n", 0, "Number of operations to perform.") - - err := viper.BindPFlags(benchCmd.Flags()) - if err != nil { - log.Fatalf("Error binding bench flags: %v", err) - } - - RootCmd.AddCommand(benchCmd) + subcommandFns["bench"] = NewBenchCommand } diff --git a/cmd/check.go b/cmd/check.go index 560130187..ae8a24c24 100644 --- a/cmd/check.go +++ b/cmd/check.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "io" "os" "github.com/spf13/cobra" @@ -10,26 +11,28 @@ import ( "github.com/pilosa/pilosa/ctl" ) -var checker = ctl.NewCheckCommand(os.Stdin, os.Stdout, os.Stderr) - -var checkCmd = &cobra.Command{ - Use: "check [path2]...", - Short: "Do a consistency check on a pilosa data file.", - Long: ` +func NewCheckCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + checker := ctl.NewCheckCommand(os.Stdin, os.Stdout, os.Stderr) + checkCmd := &cobra.Command{ + Use: "check [path2]...", + Short: "Do a consistency check on a pilosa data file.", + Long: ` Performs a consistency check on data files. `, - Run: func(cmd *cobra.Command, args []string) { - if len(args) == 0 { - fmt.Println("path required") - return - } - checker.Paths = args - if err := checker.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if len(args) == 0 { + fmt.Println("path required") + return + } + checker.Paths = args + if err := checker.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + return checkCmd } func init() { - RootCmd.AddCommand(checkCmd) + subcommandFns["check"] = NewCheckCommand } diff --git a/cmd/config.go b/cmd/config.go index 5fccbfc7d..88c8f2710 100644 --- a/cmd/config.go +++ b/cmd/config.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "io" "os" "github.com/spf13/cobra" @@ -10,20 +11,23 @@ import ( "github.com/pilosa/pilosa/ctl" ) -var conf = ctl.NewConfigCommand(os.Stdin, os.Stdout, os.Stderr) - -var confCmd = &cobra.Command{ - Use: "config", - Short: "Print the default configuration.", - Long: `config prints the default configuration to stdout +func NewConfigCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + conf := ctl.NewConfigCommand(os.Stdin, os.Stdout, os.Stderr) + confCmd := &cobra.Command{ + Use: "config", + Short: "Print the default configuration.", + Long: `config prints the default configuration to stdout `, - Run: func(cmd *cobra.Command, args []string) { - if err := conf.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if err := conf.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + + return confCmd } func init() { - RootCmd.AddCommand(confCmd) + subcommandFns["config"] = NewConfigCommand } diff --git a/cmd/export.go b/cmd/export.go index 06a76118d..e833fe1ee 100644 --- a/cmd/export.go +++ b/cmd/export.go @@ -3,21 +3,20 @@ package cmd import ( "context" "fmt" - "log" + "io" "os" "github.com/spf13/cobra" - "github.com/spf13/viper" "github.com/pilosa/pilosa/ctl" ) -var exporter = ctl.NewExportCommand(os.Stdin, os.Stdout, os.Stderr) - -var exportCmd = &cobra.Command{ - Use: "export", - Short: "Export data from pilosa.", - Long: ` +func NewExportCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + exporter := ctl.NewExportCommand(os.Stdin, os.Stdout, os.Stderr) + exportCmd := &cobra.Command{ + Use: "export", + Short: "Export data from pilosa.", + Long: ` Bulk exports a fragment to a CSV file. If the OUTFILE is not specified then the output is written to STDOUT. @@ -27,23 +26,22 @@ The format of the CSV file is: The file does not contain any headers. `, - Run: func(cmd *cobra.Command, args []string) { - if err := exporter.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if err := exporter.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + flags := exportCmd.Flags() + + flags.StringVarP(&exporter.Host, "host", "", "localhost:15000", "host:port of Pilosa.") + flags.StringVarP(&exporter.Database, "database", "d", "", "Pilosa database to export into.") + flags.StringVarP(&exporter.Frame, "frame", "f", "", "Frame to export into.") + flags.StringVarP(&exporter.Path, "output-file", "o", "", "File to write export to - default stdout") + + return exportCmd } func init() { - exportCmd.Flags().StringVarP(&exporter.Host, "host", "", "localhost:15000", "host:port of Pilosa.") - exportCmd.Flags().StringVarP(&exporter.Database, "database", "d", "", "Pilosa database to export into.") - exportCmd.Flags().StringVarP(&exporter.Frame, "frame", "f", "", "Frame to export into.") - exportCmd.Flags().StringVarP(&exporter.Path, "output-file", "o", "", "File to write export to - default stdout") - - err := viper.BindPFlags(exportCmd.Flags()) - if err != nil { - log.Fatalf("Error binding export flags: %v", err) - } - - RootCmd.AddCommand(exportCmd) + subcommandFns["export"] = NewExportCommand } diff --git a/cmd/inspect.go b/cmd/inspect.go index 666c59d5c..fea61f1a7 100644 --- a/cmd/inspect.go +++ b/cmd/inspect.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "io" "os" "github.com/spf13/cobra" @@ -10,29 +11,32 @@ import ( "github.com/pilosa/pilosa/ctl" ) -var inspecter = ctl.NewInspectCommand(os.Stdin, os.Stdout, os.Stderr) +func NewInspectCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + inspecter := ctl.NewInspectCommand(os.Stdin, os.Stdout, os.Stderr) -var inspectCmd = &cobra.Command{ - Use: "inspect", - Short: "Get stats on a pilosa data file.", - Long: ` + inspectCmd := &cobra.Command{ + Use: "inspect", + Short: "Get stats on a pilosa data file.", + Long: ` Inspects a data file and provides stats. `, - Run: func(cmd *cobra.Command, args []string) { - if len(args) == 0 { - fmt.Println("path required") - return - } else if len(args) > 1 { - fmt.Println("only one path allowed") - return - } - inspecter.Path = args[0] - if err := inspecter.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if len(args) == 0 { + fmt.Println("path required") + return + } else if len(args) > 1 { + fmt.Println("only one path allowed") + return + } + inspecter.Path = args[0] + if err := inspecter.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + return inspectCmd } func init() { - RootCmd.AddCommand(inspectCmd) + subcommandFns["inspect"] = NewInspectCommand } diff --git a/cmd/pilosa/main.go b/cmd/pilosa/main.go index efe06a0ac..76173db81 100644 --- a/cmd/pilosa/main.go +++ b/cmd/pilosa/main.go @@ -8,7 +8,7 @@ import ( ) func main() { - rootCmd := cmd.NewRootCmd(os.Stdin, os.Stdout, os.Stderr) + rootCmd := cmd.NewRootCommand(os.Stdin, os.Stdout, os.Stderr) if err := rootCmd.Execute(); err != nil { fmt.Println(err) os.Exit(1) diff --git a/cmd/restore.go b/cmd/restore.go index 399122ff3..afe5aade7 100644 --- a/cmd/restore.go +++ b/cmd/restore.go @@ -3,40 +3,38 @@ package cmd import ( "context" "fmt" - "log" + "io" "os" "github.com/spf13/cobra" - "github.com/spf13/viper" "github.com/pilosa/pilosa/ctl" ) -var restorer = ctl.NewRestoreCommand(os.Stdin, os.Stdout, os.Stderr) +func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + restorer := ctl.NewRestoreCommand(os.Stdin, os.Stdout, os.Stderr) -var restoreCmd = &cobra.Command{ - Use: "restore", - Short: "Restore data to pilosa from a backup file.", - Long: ` + restoreCmd := &cobra.Command{ + Use: "restore", + Short: "Restore data to pilosa from a backup file.", + Long: ` Restores a frame to the cluster from a backup file. `, - Run: func(cmd *cobra.Command, args []string) { - if err := restorer.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if err := restorer.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + flags := restoreCmd.Flags() + flags.StringVarP(&restorer.Host, "host", "", "localhost:15000", "host:port of Pilosa.") + flags.StringVarP(&restorer.Database, "database", "d", "", "Pilosa database to restore into.") + flags.StringVarP(&restorer.Frame, "frame", "f", "", "Frame to restore into.") + flags.StringVarP(&restorer.Path, "input-file", "i", "", "File to restore from.") + + return restoreCmd } func init() { - restoreCmd.Flags().StringVarP(&restorer.Host, "host", "", "localhost:15000", "host:port of Pilosa.") - restoreCmd.Flags().StringVarP(&restorer.Database, "database", "d", "", "Pilosa database to restore into.") - restoreCmd.Flags().StringVarP(&restorer.Frame, "frame", "f", "", "Frame to restore into.") - restoreCmd.Flags().StringVarP(&restorer.Path, "input-file", "i", "", "File to restore from.") - - err := viper.BindPFlags(restoreCmd.Flags()) - if err != nil { - log.Fatalf("Error binding restore flags: %v", err) - } - - RootCmd.AddCommand(restoreCmd) + subcommandFns["restore"] = NewRestoreCommand } diff --git a/cmd/server.go b/cmd/server.go index 2d686b2ea..58497744e 100644 --- a/cmd/server.go +++ b/cmd/server.go @@ -2,6 +2,7 @@ package cmd import ( "fmt" + "io" "os" "os/signal" "runtime/pprof" @@ -12,59 +13,58 @@ import ( "github.com/pilosa/pilosa/server" ) -var serve = server.NewCommand() - -var serveCmd = &cobra.Command{ - Use: "server", - Short: "Run Pilosa.", - Long: `pilosa server runs Pilosa. +func NewServeCmd(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { + serve := server.NewCommand() + serve.Stdin, serve.Stdout, serve.Stderr = stdin, stdout, stderr + serveCmd := &cobra.Command{ + Use: "server", + Short: "Run Pilosa.", + Long: `pilosa server runs Pilosa. It will load existing data from the configured directory, and start listening client connections on the configured port.`, - RunE: func(cmd *cobra.Command, args []string) error { - serve.Server.Handler.Version = Version - fmt.Fprintf(serve.Stderr, "Pilosa %s, build time %s\n", Version, BuildTime) + RunE: func(cmd *cobra.Command, args []string) error { + serve.Server.Handler.Version = Version + fmt.Fprintf(serve.Stderr, "Pilosa %s, build time %s\n", Version, BuildTime) - // Start CPU profiling. - if serve.CPUProfile != "" { - f, err := os.Create(serve.CPUProfile) - if err != nil { - return fmt.Errorf("create cpu profile: %v", err) + // Start CPU profiling. + if serve.CPUProfile != "" { + f, err := os.Create(serve.CPUProfile) + if err != nil { + return fmt.Errorf("create cpu profile: %v", err) + } + defer f.Close() + + fmt.Fprintln(serve.Stderr, "Starting cpu profile") + pprof.StartCPUProfile(f) + time.AfterFunc(serve.CPUTime, func() { + fmt.Fprintln(serve.Stderr, "Stopping cpu profile") + pprof.StopCPUProfile() + f.Close() + }) } - defer f.Close() - fmt.Fprintln(serve.Stderr, "Starting cpu profile") - pprof.StartCPUProfile(f) - time.AfterFunc(serve.CPUTime, func() { - fmt.Fprintln(serve.Stderr, "Stopping cpu profile") - pprof.StopCPUProfile() - f.Close() - }) - } + // Execute the program. + if err := serve.Run(); err != nil { + return err + } - // Execute the program. - if err := serve.Run(); err != nil { - return err - } + // First SIGKILL causes server to shut down gracefully. + c := make(chan os.Signal, 2) + signal.Notify(c, os.Interrupt) + sig := <-c + fmt.Fprintf(serve.Stderr, "Received %s; gracefully shutting down...\n", sig.String()) - // First SIGKILL causes server to shut down gracefully. - c := make(chan os.Signal, 2) - signal.Notify(c, os.Interrupt) - sig := <-c - fmt.Fprintf(serve.Stderr, "Received %s; gracefully shutting down...\n", sig.String()) + // Second signal causes a hard shutdown. + go func() { <-c; os.Exit(1) }() - // Second signal causes a hard shutdown. - go func() { <-c; os.Exit(1) }() - - if err := serve.Close(); err != nil { - return err - } - return nil - }, -} - -func init() { + if err := serve.Close(); err != nil { + return err + } + return nil + }, + } flags := serveCmd.Flags() flags.StringVarP(&serve.ConfigPath, "config", "c", "", "Configuration file to read from.") @@ -72,5 +72,9 @@ func init() { flags.StringVarP(&serve.CPUProfile, "cpu-profile", "", "", "Where to store CPU profile.") flags.DurationVarP(&serve.CPUTime, "cpu-time", "", 30*time.Second, "CPU profile duration.") - RootCmd.AddCommand(serveCmd) + return serveCmd +} + +func init() { + subcommandFns["server"] = NewServeCmd } diff --git a/cmd/sort.go b/cmd/sort.go index f9a1d682e..aafacf59e 100644 --- a/cmd/sort.go +++ b/cmd/sort.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "io" "os" "github.com/spf13/cobra" @@ -10,12 +11,14 @@ import ( "github.com/pilosa/pilosa/ctl" ) -var sorter = ctl.NewSortCommand(os.Stdin, os.Stdout, os.Stderr) +func NewSortCommand(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command { -var sortCmd = &cobra.Command{ - Use: "sort ", - Short: "Sort import data for optimal import performance.", - Long: ` + sorter := ctl.NewSortCommand(os.Stdin, os.Stdout, os.Stderr) + + sortCmd := &cobra.Command{ + Use: "sort ", + Short: "Sort import data for optimal import performance.", + Long: ` Sorts the import data at PATH into the optimal sort order for importing. The format of the CSV file is: @@ -24,21 +27,23 @@ The format of the CSV file is: The file should contain no headers. `, - Run: func(cmd *cobra.Command, args []string) { - if len(args) == 0 { - fmt.Println("path required") - return - } else if len(args) > 1 { - fmt.Println("only one path supported") - return - } - sorter.Path = args[0] - if err := sorter.Run(context.Background()); err != nil { - fmt.Println(err) - } - }, + Run: func(cmd *cobra.Command, args []string) { + if len(args) == 0 { + fmt.Println("path required") + return + } else if len(args) > 1 { + fmt.Println("only one path supported") + return + } + sorter.Path = args[0] + if err := sorter.Run(context.Background()); err != nil { + fmt.Println(err) + } + }, + } + return sortCmd } func init() { - RootCmd.AddCommand(sortCmd) + subcommandFns["sort"] = NewSortCommand }