From 983379928c665a7a9c91fbe39df4d80c6ad931a6 Mon Sep 17 00:00:00 2001
From: jaffee <jaffee@umbel.com>
Date: Wed, 9 Nov 2016 12:14:37 -0600
Subject: [PATCH] add cluster creation to pilosactl

---
 cmd/pilosa-benchmark-agent/main.go |   2 +-
 cmd/pilosactl/main.go              | 111 ++++++++++++++++++++++++++++-
 creator/creator.go                 |   2 -
 3 files changed, 110 insertions(+), 5 deletions(-)

diff --git a/cmd/pilosa-benchmark-agent/main.go b/cmd/pilosa-benchmark-agent/main.go
index 2dad3b38a..1ac76af68 100644
--- a/cmd/pilosa-benchmark-agent/main.go
+++ b/cmd/pilosa-benchmark-agent/main.go
@@ -110,6 +110,6 @@ func (m *Main) Run() error {
 		return fmt.Errorf("in m.Run initialization: %v", err)
 	}
 	res := sbm.Run(m.AgentNum)
-	fmt.Println(res)
+	fmt.Fprintln(m.Stdout, res)
 	return nil
 }
diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go
index 7b3d8199f..da24f762d 100644
--- a/cmd/pilosactl/main.go
+++ b/cmd/pilosactl/main.go
@@ -12,6 +12,7 @@ import (
 	"log"
 	"math/rand"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"sort"
 	"strconv"
@@ -21,8 +22,9 @@ import (
 	"time"
 	"unsafe"
 
-	"github.com/pilosa/pilosa"
-	"github.com/pilosa/pilosa/roaring"
+	"github.com/umbel/pilosa"
+	"github.com/umbel/pilosa/creator"
+	"github.com/umbel/pilosa/roaring"
 )
 
 var (
@@ -91,6 +93,7 @@ The commands are:
 	inspect    inspects fragment data files
 	check      performs a consistency check of data files
 	bench      benchmarks operations
+  create     create pilosa clusters
 
 Use the "-h" flag with any command for more information.
 `)
@@ -130,6 +133,8 @@ func (m *Main) ParseFlags(args []string) error {
 		m.Cmd = NewCheckCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "bench":
 		m.Cmd = NewBenchCommand(m.Stdin, m.Stdout, m.Stderr)
+	case "create":
+		m.Cmd = NewCreateCommand(m.Stdin, m.Stdout, m.Stderr)
 	default:
 		return ErrUnknownCommand
 	}
@@ -1147,6 +1152,108 @@ func (cmd *BenchCommand) runSetBit(ctx context.Context, client *pilosa.Client) e
 	return nil
 }
 
+// CreateCommand represents a command for creating a pilosa cluster.
+type CreateCommand struct {
+	// Type can be AWS, local, etc.
+	Type string
+
+	// ServerN is the number of pilosa hosts in the cluster
+	ServerN int
+
+	// ReplicaN is the replication number for the cluster
+	ReplicaN int
+
+	// run is used internally by local cluster to signal that the cluster should be run and not exit
+	run bool
+
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewCreateCommand returns a new instance of CreateCommand.
+func NewCreateCommand(stdin io.Reader, stdout, stderr io.Writer) *CreateCommand {
+	return &CreateCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *CreateCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(ioutil.Discard)
+	fs.StringVar(&cmd.Type, "type", "local", "Type of cluster - local, AWS, etc.")
+	fs.IntVar(&cmd.ServerN, "serverN", 3, "Number of hosts in cluster")
+	fs.IntVar(&cmd.ReplicaN, "replicaN", 1, "Replication factor for cluster")
+	fs.BoolVar(&cmd.run, "run", false, "run, don't exit")
+
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+	return nil
+}
+
+// Usage returns the usage message to be printed.
+func (cmd *CreateCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl create [args]
+
+Creates a cluster based on the arguments.
+
+The following flags are allowed:
+
+	-type 
+		type of cluster - local, AWS, etc.
+
+	-serverN
+		number of hosts in cluster
+
+	-replicaN
+		replication factor for cluster
+`)
+}
+
+// Run executes the main program execution.
+func (cmd *CreateCommand) Run() error {
+	var clus creator.Cluster
+	switch cmd.Type {
+	case "local":
+		var err error
+		if cmd.run {
+			clus, err = creator.NewLocalCluster(cmd.ReplicaN, cmd.ServerN)
+			if err != nil {
+				return fmt.Errorf("running create command: %v", err)
+			}
+			fmt.Fprintln(cmd.Stdout, strings.Join(clus.Hosts(), ","))
+			select {}
+		}
+		args := append(os.Args, "-run")
+		subcmd := exec.Command(args[0], args[1:]...)
+		pipeR, pipeW := io.Pipe()
+		subcmd.Stdout = pipeW
+		subcmd.Stderr = cmd.Stderr
+		scanner := bufio.NewScanner(pipeR)
+		err = subcmd.Start()
+		if err != nil {
+			return fmt.Errorf("error kicking off local cluster: %v", err)
+		}
+		scanner.Scan()
+		fmt.Fprintln(cmd.Stdout, scanner.Text())
+		subcmd.Stdout = ioutil.Discard
+		pipeR.Close()
+
+	case "AWS":
+		return fmt.Errorf("AWS cluster type is not yet implemented")
+	default:
+		return fmt.Errorf("Unknown cluster type %v", cmd.Type)
+	}
+
+	return nil
+}
+
 // readCSVRow reads a bitmap/profile pair from a CSV row.
 func readCSVRow(r *csv.Reader) (bitmapID, profileID uint64, err error) {
 	// Read CSV row.
diff --git a/creator/creator.go b/creator/creator.go
index 73e4f993a..f117bf208 100644
--- a/creator/creator.go
+++ b/creator/creator.go
@@ -4,7 +4,6 @@ package creator
 import (
 	"fmt"
 	"io/ioutil"
-	"log"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -77,7 +76,6 @@ func NewLocalCluster(replicaN, serverN int) (Cluster, error) {
 
 	// Open all servers.
 	for _, s := range servers {
-		log.Printf("opening : %v", s)
 		if err := s.Open(); err != nil {
 			return localCluster, err
 		}