Merge pull request #141 from benbjohnson/ctx

Add context to pilosactl.
This commit is contained in:
Travis Turner 2016-11-16 22:00:15 -06:00 committed by GitHub
commit a17bb2a19e
2 changed files with 27 additions and 25 deletions

View file

@ -3,6 +3,7 @@ package main_test
import (
"bytes"
"encoding/json"
"context"
"fmt"
"io"
"io/ioutil"
@ -217,7 +218,7 @@ func TestMain_FrameRestore(t *testing.T) {
client, err := pilosa.NewClient(m2.Server.Host)
if err != nil {
t.Fatal(err)
} else if err := client.RestoreFrame(m0.Server.Host, "d", "f"); err != nil {
} else if err := client.RestoreFrame(context.Background(), m0.Server.Host, "d", "f"); err != nil {
t.Fatal(err)
}

View file

@ -2,6 +2,7 @@ package main
import (
"bufio"
"context"
"encoding/csv"
"errors"
"flag"
@ -96,7 +97,7 @@ Use the "-h" flag with any command for more information.
}
// Run executes the main program execution.
func (m *Main) Run() error { return m.Cmd.Run() }
func (m *Main) Run() error { return m.Cmd.Run(context.Background()) }
// ParseFlags parses command line flags from args.
func (m *Main) ParseFlags(args []string) error {
@ -149,7 +150,7 @@ func (m *Main) ParseFlags(args []string) error {
type Command interface {
Usage() string
ParseFlags(args []string) error
Run() error
Run(context.Context) error
}
// ConfigCommand represents a command for printing a default config.
@ -189,7 +190,7 @@ Prints the default configuration file to standard out.
}
// Run executes the main program execution.
func (cmd *ConfigCommand) Run() error {
func (cmd *ConfigCommand) Run(ctx context.Context) error {
fmt.Fprintln(cmd.Stdout, strings.TrimSpace(`
data-dir = "~/.pilosa"
host = "localhost:15000"
@ -276,7 +277,7 @@ The file should contain no headers.
}
// Run executes the main program execution.
func (cmd *ImportCommand) Run() error {
func (cmd *ImportCommand) Run(ctx context.Context) error {
logger := log.New(cmd.Stderr, "", log.LstdFlags)
// Validate arguments.
@ -300,7 +301,7 @@ func (cmd *ImportCommand) Run() error {
for _, path := range cmd.Paths {
// Parse path into bits.
logger.Printf("parsing: %s", path)
if err := cmd.importPath(path); err != nil {
if err := cmd.importPath(ctx, path); err != nil {
return err
}
}
@ -309,7 +310,7 @@ func (cmd *ImportCommand) Run() error {
}
// importPath parses a path into bits and imports it to the server.
func (cmd *ImportCommand) importPath(path string) error {
func (cmd *ImportCommand) importPath(ctx context.Context, path string) error {
a := make([]pilosa.Bit, 0, cmd.BufferSize)
// Open file for reading.
@ -356,7 +357,7 @@ func (cmd *ImportCommand) importPath(path string) error {
// If we've reached the buffer size then import bits.
if len(a) == cmd.BufferSize {
if err := cmd.importBits(a); err != nil {
if err := cmd.importBits(ctx, a); err != nil {
return err
}
a = a[:0]
@ -364,7 +365,7 @@ func (cmd *ImportCommand) importPath(path string) error {
}
// If there are still bits in the buffer then flush them.
if err := cmd.importBits(a); err != nil {
if err := cmd.importBits(ctx, a); err != nil {
return err
}
@ -372,7 +373,7 @@ func (cmd *ImportCommand) importPath(path string) error {
}
// importPath parses a path into bits and imports it to the server.
func (cmd *ImportCommand) importBits(bits []pilosa.Bit) error {
func (cmd *ImportCommand) importBits(ctx context.Context, bits []pilosa.Bit) error {
logger := log.New(cmd.Stderr, "", log.LstdFlags)
// Group bits by slice.
@ -382,7 +383,7 @@ func (cmd *ImportCommand) importBits(bits []pilosa.Bit) error {
// Parse path into bits.
for slice, bits := range bitsBySlice {
logger.Printf("importing slice: %d, n=%d", slice, len(bits))
if err := cmd.Client.Import(cmd.Database, cmd.Frame, slice, bits); err != nil {
if err := cmd.Client.Import(ctx, cmd.Database, cmd.Frame, slice, bits); err != nil {
return err
}
}
@ -449,7 +450,7 @@ The file does not contain any headers.
}
// Run executes the main program execution.
func (cmd *ExportCommand) Run() error {
func (cmd *ExportCommand) Run(ctx context.Context) error {
logger := log.New(cmd.Stderr, "", log.LstdFlags)
// Validate arguments.
@ -479,7 +480,7 @@ func (cmd *ExportCommand) Run() error {
}
// Determine slice count.
sliceN, err := client.SliceN()
sliceN, err := client.SliceN(ctx)
if err != nil {
return err
}
@ -487,7 +488,7 @@ func (cmd *ExportCommand) Run() error {
// Export each slice.
for slice := uint64(0); slice <= sliceN; slice++ {
logger.Printf("exporting slice: %d", slice)
if err := client.ExportCSV(cmd.Database, cmd.Frame, slice, w); err != nil {
if err := client.ExportCSV(ctx, cmd.Database, cmd.Frame, slice, w); err != nil {
return err
}
}
@ -557,7 +558,7 @@ The file should contain no headers.
}
// Run executes the main program execution.
func (cmd *SortCommand) Run() error {
func (cmd *SortCommand) Run(ctx context.Context) error {
// Open file for reading.
f, err := os.Open(cmd.Path)
if err != nil {
@ -660,7 +661,7 @@ Backs up the database and frame from across the cluster into a single file.
}
// Run executes the main program execution.
func (cmd *BackupCommand) Run() error {
func (cmd *BackupCommand) Run(ctx context.Context) error {
// Validate arguments.
if cmd.Path == "" {
return errors.New("output file required")
@ -680,7 +681,7 @@ func (cmd *BackupCommand) Run() error {
defer f.Close()
// Begin streaming backup.
if err := client.BackupTo(f, cmd.Database, cmd.Frame); err != nil {
if err := client.BackupTo(ctx, f, cmd.Database, cmd.Frame); err != nil {
return err
}
@ -753,7 +754,7 @@ Restores a frame to the cluster from a backup file.
}
// Run executes the main program execution.
func (cmd *RestoreCommand) Run() error {
func (cmd *RestoreCommand) Run(ctx context.Context) error {
// Validate arguments.
if cmd.Path == "" {
return errors.New("backup file required")
@ -773,7 +774,7 @@ func (cmd *RestoreCommand) Run() error {
defer f.Close()
// Restore backup file to the cluster.
if err := client.RestoreFrom(f, cmd.Database, cmd.Frame); err != nil {
if err := client.RestoreFrom(ctx, f, cmd.Database, cmd.Frame); err != nil {
return err
}
@ -830,7 +831,7 @@ Inspects a data file and provides stats.
}
// Run executes the main program execution.
func (cmd *InspectCommand) Run() error {
func (cmd *InspectCommand) Run(ctx context.Context) error {
// Open file handle.
f, err := os.Open(cmd.Path)
if err != nil {
@ -937,7 +938,7 @@ Performs a consistency check on data files.
}
// Run executes the main program execution.
func (cmd *CheckCommand) Run() error {
func (cmd *CheckCommand) Run(ctx context.Context) error {
for _, path := range cmd.Paths {
switch filepath.Ext(path) {
case "":
@ -1095,7 +1096,7 @@ The following operations are available:
}
// Run executes the main program execution.
func (cmd *BenchCommand) Run() error {
func (cmd *BenchCommand) Run(ctx context.Context) error {
// Create a client to the server.
client, err := pilosa.NewClient(cmd.Host)
if err != nil {
@ -1104,7 +1105,7 @@ func (cmd *BenchCommand) Run() error {
switch cmd.Op {
case "set-bit":
return cmd.runSetBit(client)
return cmd.runSetBit(ctx, client)
case "":
return errors.New("op required")
default:
@ -1113,7 +1114,7 @@ func (cmd *BenchCommand) Run() error {
}
// runSetBit executes a benchmark of random SetBit() operations.
func (cmd *BenchCommand) runSetBit(client *pilosa.Client) error {
func (cmd *BenchCommand) runSetBit(ctx context.Context, client *pilosa.Client) error {
if cmd.N == 0 {
return errors.New("operation count required")
} else if cmd.Database == "" {
@ -1134,7 +1135,7 @@ func (cmd *BenchCommand) runSetBit(client *pilosa.Client) error {
q := fmt.Sprintf(`SetBit(id=%d, frame="%s", profileID=%d)`, bitmapID, cmd.Frame, profileID)
if _, err := client.ExecuteQuery(cmd.Database, q, true); err != nil {
if _, err := client.ExecuteQuery(ctx, cmd.Database, q, true); err != nil {
return err
}
}