Merge pull request #56 from benbjohnson/backup

Add backup/restore to pilosactl

client.go

@@ -1,14 +1,19 @@
 package pilosa
 
 import (
+	"archive/tar"
 	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
+	"log"
+	"math/rand"
 	"net/http"
 	"net/url"
 	"strconv"
+	"time"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/umbel/pilosa/internal"
@@ -38,6 +43,30 @@ func NewClient(host string) (*Client, error) {
 // Host returns the host the client was initialized with.
 func (c *Client) Host() string { return c.host }
 
+// SliceN returns the number of slices on a server.
+func (c *Client) SliceN() (uint64, error) {
+	// Execute request against the host.
+	u := url.URL{
+		Scheme: "http",
+		Host:   c.host,
+		Path:   "/slices/max",
+	}
+	resp, err := c.HTTPClient.Get(u.String())
+	if err != nil {
+		return 0, err
+	}
+	defer resp.Body.Close()
+
+	var rsp sliceMaxResponse
+	if resp.StatusCode != http.StatusOK {
+		return 0, fmt.Errorf("http: status=%d", resp.StatusCode)
+	} else if err := json.NewDecoder(resp.Body).Decode(&rsp); err != nil {
+		return 0, fmt.Errorf("json decode: %s", err)
+	}
+
+	return rsp.SliceMax, nil
+}
+
 // SliceNodes returns a list of nodes that own a slice.
 func (c *Client) SliceNodes(slice uint64) ([]*Node, error) {
 	// Execute request against the host.
@@ -140,6 +169,185 @@ func (c *Client) importNode(node *Node, buf []byte) error {
 	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 == "" {
+		return ErrDatabaseRequired
+	} else if frame == "" {
+		return ErrFrameRequired
+	}
+
+	// Create tar writer around writer.
+	tw := tar.NewWriter(w)
+
+	// Find the maximum number of slices.
+	sliceN, err := c.SliceN()
+	if err != nil {
+		return fmt.Errorf("slice n: %s", err)
+	}
+
+	// Backup every slice to the tar file.
+	for i := uint64(0); i <= sliceN; i++ {
+		if err := c.backupSliceTo(tw, db, frame, i); err != nil {
+			return err
+		}
+	}
+
+	// Close tar file.
+	if err := tw.Close(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// backupSliceTo backs up a single slice to tw.
+func (c *Client) backupSliceTo(tw *tar.Writer, db, frame string, slice uint64) error {
+	// Retrieve a list of nodes that own the slice.
+	nodes, err := c.SliceNodes(slice)
+	if err != nil {
+		return fmt.Errorf("slice nodes: %s", err)
+	}
+
+	// Try to backup slice from each one until successful.
+	var data []byte
+	for _, i := range rand.Perm(len(nodes)) {
+		buf, err := c.backupSliceNode(db, frame, slice, nodes[i])
+		if err == nil {
+			data = buf
+			break // backup successful
+		} else if err == ErrFragmentNotFound {
+			return nil // slice doesn't exist
+		} else if err != nil {
+			log.Println(err)
+			continue
+		}
+	}
+
+	// Return error if unable to backup from any slice.
+	if data == nil {
+		return fmt.Errorf("unable to backup slice %d", slice)
+	}
+
+	// Write slice file header.
+	if err := tw.WriteHeader(&tar.Header{
+		Name:    strconv.FormatUint(slice, 10),
+		Mode:    0666,
+		Size:    int64(len(data)),
+		ModTime: time.Now(),
+	}); err != nil {
+		return err
+	}
+
+	// Write buffer to file.
+	if _, err := tw.Write(data); err != nil {
+		return fmt.Errorf("write buffer: %s", err)
+	}
+
+	return nil
+}
+
+func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) ([]byte, error) {
+	u := url.URL{
+		Scheme: "http",
+		Host:   node.Host,
+		Path:   "/fragment/data",
+		RawQuery: url.Values{
+			"db":    {db},
+			"frame": {frame},
+			"slice": {strconv.FormatUint(slice, 10)},
+		}.Encode(),
+	}
+	resp, err := c.HTTPClient.Get(u.String())
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	// Return error if status is not OK.
+	if resp.StatusCode == http.StatusNotFound {
+		return nil, ErrFragmentNotFound
+	} else if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("unexpected backup status code: host=%s, code=%d", node.Host, resp.StatusCode)
+	}
+
+	return ioutil.ReadAll(resp.Body)
+}
+
+// RestoreFrom restores a frame from a backup file to an entire cluster.
+func (c *Client) RestoreFrom(r io.Reader, db, frame string) error {
+	if db == "" {
+		return ErrDatabaseRequired
+	} else if frame == "" {
+		return ErrFrameRequired
+	}
+
+	// Create tar reader around input.
+	tr := tar.NewReader(r)
+
+	// Process each file.
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			return nil
+		} else if err != nil {
+			return err
+		}
+
+		// Parse slice from entry name.
+		slice, err := strconv.ParseUint(hdr.Name, 10, 64)
+		if err != nil {
+			return fmt.Errorf("invalid backup entry: %s", hdr.Name)
+		}
+
+		// Read file into buffer.
+		var buf bytes.Buffer
+		if _, err := io.CopyN(&buf, tr, hdr.Size); err != nil {
+			return err
+		}
+
+		// Restore file to all nodes that own it.
+		if err := c.restoreSliceFrom(buf.Bytes(), db, frame, slice); err != nil {
+			return err
+		}
+	}
+}
+
+// restoreSliceFrom restores a single slice to all owning nodes.
+func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) error {
+	// Retrieve a list of nodes that own the slice.
+	nodes, err := c.SliceNodes(slice)
+	if err != nil {
+		return fmt.Errorf("slice nodes: %s", err)
+	}
+
+	// Restore slice to each owner.
+	for _, node := range nodes {
+		u := url.URL{
+			Scheme: "http",
+			Host:   node.Host,
+			Path:   "/fragment/data",
+			RawQuery: url.Values{
+				"db":    {db},
+				"frame": {frame},
+				"slice": {strconv.FormatUint(slice, 10)},
+			}.Encode(),
+		}
+		resp, err := c.HTTPClient.Post(u.String(), "application/octet-stream", bytes.NewReader(buf))
+		if err != nil {
+			return err
+		}
+		resp.Body.Close()
+
+		// Return error if response not OK.
+		if resp.StatusCode != http.StatusOK {
+			return fmt.Errorf("unexpected status code: host=%s, code=%d", node.Host, resp.StatusCode)
+		}
+	}
+
+	return nil
+}
+
 // Bit represents the location of a single bit.
 type Bit struct {
 	BitmapID  uint64

client_test.go

@@ -1,6 +1,7 @@
 package pilosa_test
 
 import (
+	"bytes"
 	"reflect"
 	"testing"
 
@@ -39,6 +40,51 @@ func TestClient_Import(t *testing.T) {
 	}
 }
 
+// Ensure client backup and restore a frame.
+func TestClient_BackupRestore(t *testing.T) {
+	idx := MustOpenIndex()
+	defer idx.Close()
+
+	idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBits(100, 1, 2, 3, SliceWidth-1)
+	idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(100, SliceWidth, SliceWidth+2)
+	idx.MustCreateFragmentIfNotExists("d", "f", 5).MustSetBits(100, (5*SliceWidth)+1)
+	idx.MustCreateFragmentIfNotExists("d", "f", 0).MustSetBits(200, 20000)
+
+	s := NewServer()
+	defer s.Close()
+	s.Handler.Host = s.Host()
+	s.Handler.Cluster = NewCluster(1)
+	s.Handler.Cluster.Nodes[0].Host = s.Host()
+	s.Handler.Index = idx.Index
+
+	c := MustNewClient(s.Host())
+
+	// Backup from frame.
+	var buf bytes.Buffer
+	if err := c.BackupTo(&buf, "d", "f"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Restore to a different frame.
+	if err := c.RestoreFrom(&buf, "x", "y"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify data.
+	if a := idx.Fragment("x", "y", 0).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{1, 2, 3, SliceWidth - 1}) {
+		t.Fatalf("unexpected bits(0): %+v", a)
+	}
+	if a := idx.Fragment("x", "y", 1).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{SliceWidth, SliceWidth + 2}) {
+		t.Fatalf("unexpected bits(0): %+v", a)
+	}
+	if a := idx.Fragment("x", "y", 5).Bitmap(100).Bits(); !reflect.DeepEqual(a, []uint64{(5 * SliceWidth) + 1}) {
+		t.Fatalf("unexpected bits(0): %+v", a)
+	}
+	if a := idx.Fragment("x", "y", 0).Bitmap(200).Bits(); !reflect.DeepEqual(a, []uint64{20000}) {
+		t.Fatalf("unexpected bits: %+v", a)
+	}
+}
+
 // Client represents a test wrapper for pilosa.Client.
 type Client struct {
 	*pilosa.Client

cmd/pilosactl/main.go

@@ -78,6 +78,10 @@ func (m *Main) Run() error {
 		cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
 	case "import":
 		cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
+	case "backup":
+		cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr)
+	case "restore":
+		cmd = NewRestoreCommand(m.Stdin, m.Stdout, m.Stderr)
 	default:
 		return ErrUnknownCommand
 	}
@@ -324,3 +328,175 @@ func (cmd *ImportCommand) parsePath(path string) ([]pilosa.Bit, error) {
 
 	return a, nil
 }
+
+// BackupCommand represents a command for backing up a frame.
+type BackupCommand struct {
+	// Destination host and port.
+	Host string
+
+	// Name of the database & frame to backup.
+	Database string
+	Frame    string
+
+	// Output file to write to.
+	Path string
+
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewBackupCommand returns a new instance of BackupCommand.
+func NewBackupCommand(stdin io.Reader, stdout, stderr io.Writer) *BackupCommand {
+	return &BackupCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *BackupCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(cmd.Stderr)
+	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 *BackupCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl backup -host HOST -d database -f frame -o PATH
+
+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 {
+	// Validate arguments.
+	if cmd.Path == "" {
+		return errors.New("output file required")
+	}
+
+	// Create a client to the server.
+	client, err := pilosa.NewClient(cmd.Host)
+	if err != nil {
+		return err
+	}
+
+	// Open output file.
+	f, err := os.Create(cmd.Path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// Begin streaming backup.
+	if err := client.BackupTo(f, cmd.Database, cmd.Frame); err != nil {
+		return err
+	}
+
+	// Sync & close file to ensure durability.
+	if err := f.Sync(); err != nil {
+		return err
+	} else if err = f.Close(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// RestoreCommand represents a command for restoring a frame from a backup.
+type RestoreCommand struct {
+	// Destination host and port.
+	Host string
+
+	// Name of the database & frame to backup.
+	Database string
+	Frame    string
+
+	// Import file to read from.
+	Path string
+
+	// Standard input/output
+	Stdin  io.Reader
+	Stdout io.Writer
+	Stderr io.Writer
+}
+
+// NewRestoreCommand returns a new instance of RestoreCommand.
+func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreCommand {
+	return &RestoreCommand{
+		Stdin:  stdin,
+		Stdout: stdout,
+		Stderr: stderr,
+	}
+}
+
+// ParseFlags parses command line flags from args.
+func (cmd *RestoreCommand) ParseFlags(args []string) error {
+	fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
+	fs.SetOutput(cmd.Stderr)
+	fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
+	fs.StringVar(&cmd.Database, "d", "", "database")
+	fs.StringVar(&cmd.Frame, "f", "", "frame")
+	if err := fs.Parse(args); err != nil {
+		return err
+	}
+
+	// Read input path from the args.
+	if fs.NArg() == 0 {
+		return errors.New("path required")
+	} else if fs.NArg() > 1 {
+		return errors.New("too many paths specified")
+	}
+	cmd.Path = fs.Arg(0)
+
+	return nil
+}
+
+// Usage returns the usage message to be printed.
+func (cmd *RestoreCommand) Usage() string {
+	return strings.TrimSpace(`
+usage: pilosactl restore -host HOST -d database -f frame PATH
+
+Restores a frame to the cluster from a backup file.
+`)
+}
+
+// Run executes the main program execution.
+func (cmd *RestoreCommand) Run() error {
+	// Validate arguments.
+	if cmd.Path == "" {
+		return errors.New("backup file required")
+	}
+
+	// Create a client to the server.
+	client, err := pilosa.NewClient(cmd.Host)
+	if err != nil {
+		return err
+	}
+
+	// Open backup file.
+	f, err := os.Open(cmd.Path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// Restore backup file to the cluster.
+	if err := client.RestoreFrom(f, cmd.Database, cmd.Frame); err != nil {
+		return err
+	}
+
+	return nil
+}

executor_test.go

@@ -128,7 +128,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) {
 		t.Fatalf("unexpected bitmap count: %d", n)
 	}
 
-	if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil); err != nil {
+	if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
 		t.Fatal(err)
 	} else {
 		if !res.(bool) {
@@ -139,7 +139,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) {
 	if n := f.Bitmap(11).Count(); n != 1 {
 		t.Fatalf("unexpected bitmap count: %d", n)
 	}
-	if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil); err != nil {
+	if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
 		t.Fatal(err)
 	} else {
 		if res.(bool) {

handler.go

@@ -152,7 +152,6 @@ func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) {
 }
 
 func (h *Handler) handleGetSliceMax(w http.ResponseWriter, r *http.Request) error {
-
 	sm := h.Index.SliceN()
 	if strings.Contains(r.Header.Get("Accept"), "application/x-protobuf") {
 		pb := &internal.SliceMaxResponse{
@@ -165,8 +164,11 @@ func (h *Handler) handleGetSliceMax(w http.ResponseWriter, r *http.Request) erro
 		}
 		return nil
 	}
-	resp := map[string]uint64{"SliceMax": sm}
-	return json.NewEncoder(w).Encode(resp)
+	return json.NewEncoder(w).Encode(sliceMaxResponse{SliceMax: sm})
+}
+
+type sliceMaxResponse struct {
+	SliceMax uint64 `json:"SliceMax"`
 }
 
 // readProfiles returns a list of profile objects by id.

pilosa.go

@@ -18,6 +18,9 @@ var (
 
 	// ErrFrameRequired is returned when no frame is specified.
 	ErrFrameRequired = errors.New("frame required")
+
+	// ErrFragmentNotFound is returned when a fragment does not exist.
+	ErrFragmentNotFound = errors.New("fragment not found")
 )
 
 // Version represents the current running version of Pilosa.