Merge pull request #107 from benbjohnson/export

Add export command to pilosactl

client.go

@@ -255,6 +255,77 @@ func (c *Client) importNode(node *Node, buf []byte) error {
 	return nil
 }
 
+// ExportCSV bulk exports data for a single slice from a host to CSV format.
+func (c *Client) ExportCSV(db, frame string, slice uint64, w io.Writer) error {
+	if db == "" {
+		return ErrDatabaseRequired
+	} else if frame == "" {
+		return ErrFrameRequired
+	}
+
+	// Retrieve a list of nodes that own the slice.
+	nodes, err := c.FragmentNodes(db, slice)
+	if err != nil {
+		return fmt.Errorf("slice nodes: %s", err)
+	}
+
+	// Attempt nodes in random order.
+	var e error
+	for _, i := range rand.Perm(len(nodes)) {
+		node := nodes[i]
+
+		if err := c.exportNodeCSV(node, db, frame, slice, w); err != nil {
+			e = fmt.Errorf("export node: host=%s, err=%s", node.Host, err)
+			continue
+		} else {
+			return nil
+		}
+	}
+
+	return e
+}
+
+// exportNode copies a CSV export from a node to w.
+func (c *Client) exportNodeCSV(node *Node, db, frame string, slice uint64, w io.Writer) error {
+	// Create URL.
+	u := url.URL{
+		Scheme: "http",
+		Host:   node.Host,
+		Path:   "/export",
+		RawQuery: url.Values{
+			"db":    {db},
+			"frame": {frame},
+			"slice": {strconv.FormatUint(slice, 10)},
+		}.Encode(),
+	}
+
+	// Generate HTTP request.
+	req, err := http.NewRequest("GET", u.String(), nil)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Accept", "text/csv")
+
+	// Execute request against the host.
+	resp, err := c.HTTPClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	// Validate status code.
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("invalid status: %d", resp.StatusCode)
+	}
+
+	// Copy body to writer.
+	if _, err := io.Copy(w, resp.Body); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // BackupTo backs up an entire frame from a cluster to w.
 func (c *Client) BackupTo(w io.Writer, db, frame string) error {
 	if db == "" {

cmd/pilosactl/main.go

@@ -80,6 +80,7 @@ The commands are:
 
 	config     prints the default configuration
 	import     imports data from a CSV file
+	export     exports data to a CSV file
 	backup     backs up a frame to an archive file
 	restore    restores a frame from an archive file
 	inspect    inspects fragment data files
@@ -109,6 +110,8 @@ func (m *Main) ParseFlags(args []string) error {
 		m.Cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "import":
 		m.Cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
+	case "export":
+		m.Cmd = NewExportCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "backup":
 		m.Cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "restore":
@@ -350,6 +353,118 @@ func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) {
 	return a, nil
 }
 
+// ExportCommand represents a command for bulk exporting data from a server.
+type ExportCommand struct {
+	// Remote host and port.
+	Host string
+
+	// Name of the database & frame to export from.
+	Database string
+	Frame    string
+
+	// Filename to export to.
+	Path string
+
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewExportCommand returns a new instance of ExportCommand.
+func NewExportCommand(stdin io.Reader, stdout, stderr io.Writer) *ExportCommand {
+	return &ExportCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *ExportCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(ioutil.Discard)
+	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
+	fs.StringVar(&cmd.Database, "d", "", "database")
+	fs.StringVar(&cmd.Frame, "f", "", "frame")
+	fs.StringVar(&cmd.Path, "o", "", "output file")
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Usage returns the usage message to be printed.
+func (cmd *ExportCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl export -host HOST -d database -f frame -o OUTFILE
+
+Bulk exports a fragment to a CSV file. If the OUTFILE is not specified then
+the output is written to STDOUT.
+
+The format of the CSV file is:
+
+	BITMAPID,PROFILEID
+
+The file does not contain any headers.
+`)
+}
+
+// Run executes the main program execution.
+func (cmd *ExportCommand) Run() error {
+	logger := log.New(cmd.Stderr, "", log.LstdFlags)
+
+	// Validate arguments.
+	if cmd.Database == "" {
+		return pilosa.ErrDatabaseRequired
+	} else if cmd.Frame == "" {
+		return pilosa.ErrFrameRequired
+	}
+
+	// Use output file, if specified.
+	// Otherwise use STDOUT.
+	var w io.Writer = cmd.Stdout
+	if cmd.Path != "" {
+		f, err := os.Create(cmd.Path)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+
+		w = f
+	}
+
+	// Create a client to the server.
+	client, err := pilosa.NewClient(cmd.Host)
+	if err != nil {
+		return err
+	}
+
+	// Determine slice count.
+	sliceN, err := client.SliceN()
+	if err != nil {
+		return err
+	}
+
+	// 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 {
+			return err
+		}
+	}
+
+	// Close writer, if applicable.
+	if w, ok := w.(io.Closer); ok {
+		if err := w.Close(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // BackupCommand represents a command for backing up a frame.
 type BackupCommand struct {
 	// Destination host and port.

fragment.go

@@ -431,6 +431,25 @@ func (f *Fragment) pos(bitmapID, profileID uint64) (uint64, error) {
 	return (bitmapID * SliceWidth) + (profileID % SliceWidth), nil
 }
 
+// ForEachBit executes fn for every bit set in the fragment.
+// Errors returned from fn are passed through.
+func (f *Fragment) ForEachBit(fn func(bitmapID, profileID uint64) error) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	var err error
+	f.storage.ForEach(func(i uint64) {
+		// Skip if an error has already occurred.
+		if err != nil {
+			return
+		}
+
+		// Invoke caller's function.
+		err = fn(i/SliceWidth, (f.slice*SliceWidth)+(i%SliceWidth))
+	})
+	return err
+}
+
 // Top returns the top bitmaps from the fragment.
 // If opt.Src is specified then only bitmaps which intersect src are returned.
 // If opt.FilterValues exist then the bitmap attribute specified by field is matched.

fragment_test.go

@@ -108,6 +108,35 @@ func TestFragment_Snapshot(t *testing.T) {
 	}
 }
 
+// Ensure a fragment can iterate over all bits in order.
+func TestFragment_ForEachBit(t *testing.T) {
+	f := MustOpenFragment("d", "f", 0)
+	defer f.Close()
+
+	// Set bits on the fragment.
+	if _, err := f.SetBit(100, 20, nil, 0); err != nil {
+		t.Fatal(err)
+	} else if _, err := f.SetBit(2, 38, nil, 0); err != nil {
+		t.Fatal(err)
+	} else if _, err := f.SetBit(2, 37, nil, 0); err != nil {
+		t.Fatal(err)
+	}
+
+	// Iterate over bits.
+	var result [][2]uint64
+	if err := f.ForEachBit(func(bitmapID, profileID uint64) error {
+		result = append(result, [2]uint64{bitmapID, profileID})
+		return nil
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify bits are correct.
+	if !reflect.DeepEqual(result, [][2]uint64{{2, 37}, {2, 38}, {100, 20}}) {
+		t.Fatalf("unexpected result: %#v", result)
+	}
+}
+
 // Ensure a fragment can return the top n results.
 func TestFragment_Top(t *testing.T) {
 	f := MustOpenFragment("d", "f", 0)

handler.go

@@ -1,6 +1,7 @@
 package pilosa
 
 import (
+	"encoding/csv"
 	"encoding/json"
 	"errors"
 	"expvar"
@@ -91,6 +92,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		default:
 			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
 		}
+	case "/export":
+		switch r.Method {
+		case "GET":
+			h.handleGetExport(w, r)
+		default:
+			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		}
 	case "/slices/max":
 		switch r.Method {
 		case "GET":
@@ -430,6 +438,58 @@ func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) {
 	w.Write(buf)
 }
 
+// handleGetExport handles /export requests.
+func (h *Handler) handleGetExport(w http.ResponseWriter, r *http.Request) {
+	switch r.Header.Get("Accept") {
+	case "text/csv":
+		h.handleGetExportCSV(w, r)
+	default:
+		http.Error(w, "Not acceptable", http.StatusNotAcceptable)
+	}
+}
+
+func (h *Handler) handleGetExportCSV(w http.ResponseWriter, r *http.Request) {
+	// Parse query parameters.
+	q := r.URL.Query()
+	db, frame := q.Get("db"), q.Get("frame")
+
+	slice, err := strconv.ParseUint(q.Get("slice"), 10, 64)
+	if err != nil {
+		http.Error(w, "invalid slice", http.StatusBadRequest)
+		return
+	}
+
+	// Validate that this handler owns the slice.
+	if !h.Cluster.OwnsFragment(h.Host, db, slice) {
+		mesg := fmt.Sprintf("host does not own slice %s-%s slice:%d", h.Host, db, slice)
+		http.Error(w, mesg, http.StatusPreconditionFailed)
+		return
+	}
+
+	// Find the fragment.
+	f := h.Index.Fragment(db, frame, slice)
+	if f == nil {
+		return
+	}
+
+	// Wrap writer with a CSV writer.
+	cw := csv.NewWriter(w)
+
+	// Iterate over each bit.
+	if err := f.ForEachBit(func(bitmapID, profileID uint64) error {
+		return cw.Write([]string{
+			strconv.FormatUint(bitmapID, 10),
+			strconv.FormatUint(profileID, 10),
+		})
+	}); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// Ensure data is flushed.
+	cw.Flush()
+}
+
 // handleGetFragmentNodes handles /fragment/nodes requests.
 func (h *Handler) handleGetFragmentNodes(w http.ResponseWriter, r *http.Request) {
 	q := r.URL.Query()