diff --git a/client.go b/client.go index e941dce04..e2a17d681 100644 --- a/client.go +++ b/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 == "" { diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 262a74f27..5ea31da21 100644 --- a/cmd/pilosactl/main.go +++ b/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. diff --git a/fragment.go b/fragment.go index 04010125e..c1f401b18 100644 --- a/fragment.go +++ b/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. diff --git a/fragment_test.go b/fragment_test.go index 2c1def646..339c83f56 100644 --- a/fragment_test.go +++ b/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) diff --git a/handler.go b/handler.go index fe4cecd85..c5a4607d4 100644 --- a/handler.go +++ b/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()