diff --git a/api.go b/api.go index da337252d..88b759903 100644 --- a/api.go +++ b/api.go @@ -339,6 +339,18 @@ func (api *API) ExportCSV(_ context.Context, indexName string, fieldName string, return ErrClusterDoesNotOwnShard } + // Find index. + index := api.holder.Index(indexName) + if index == nil { + return newNotFoundError(ErrIndexNotFound) + } + + // Find field from the index. + field := index.Field(fieldName) + if field == nil { + return newNotFoundError(ErrFieldNotFound) + } + // Find the fragment. f := api.holder.fragment(indexName, fieldName, viewStandard, shard) if f == nil { @@ -348,13 +360,34 @@ func (api *API) ExportCSV(_ context.Context, indexName string, fieldName string, // Wrap writer with a CSV writer. cw := csv.NewWriter(w) + // Define the function to write each bit as a string, + // translating to keys where necessary. + fn := func(rowID, columnID uint64) error { + var rowStr string + var colStr string + var err error + + if field.keys() { + if rowStr, err = api.holder.translateFile.TranslateRowToString(index.Name(), field.Name(), rowID); err != nil { + return errors.Wrap(err, "translating row") + } + } else { + rowStr = strconv.FormatUint(rowID, 10) + } + + if index.Keys() { + if colStr, err = api.holder.translateFile.TranslateColumnToString(index.Name(), columnID); err != nil { + return errors.Wrap(err, "translating column") + } + } else { + colStr = strconv.FormatUint(columnID, 10) + } + + return cw.Write([]string{rowStr, colStr}) + } + // Iterate over each column. - if err := f.forEachBit(func(rowID, columnID uint64) error { - return cw.Write([]string{ - strconv.FormatUint(rowID, 10), - strconv.FormatUint(columnID, 10), - }) - }); err != nil { + if err := f.forEachBit(fn); err != nil { return errors.Wrap(err, "writing CSV") } diff --git a/http/client_test.go b/http/client_test.go index 87c433f05..131161d3c 100644 --- a/http/client_test.go +++ b/http/client_test.go @@ -15,6 +15,8 @@ package http_test import ( + "bufio" + "bytes" "context" "fmt" gohttp "net/http" @@ -168,6 +170,169 @@ func TestClient_MultiNode(t *testing.T) { } } +// Ensure client can export data. +func TestClient_Export(t *testing.T) { + cmd := test.MustRunCluster(t, 1)[0] + host := cmd.URL() + + cmd.MustCreateIndex(t, "keyed", pilosa.IndexOptions{Keys: true}) + cmd.MustCreateIndex(t, "unkeyed", pilosa.IndexOptions{Keys: false}) + + cmd.MustCreateField(t, "keyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys()) + cmd.MustCreateField(t, "keyed", "unkeyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000)) + cmd.MustCreateField(t, "unkeyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys()) + cmd.MustCreateField(t, "unkeyed", "unkeyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000)) + + c := MustNewClient(host, http.GetHTTPClient(nil)) + + data := []pilosa.Bit{ + {RowID: 1, ColumnID: 100, RowKey: "row1", ColumnKey: "col100"}, + {RowID: 1, ColumnID: 101, RowKey: "row1", ColumnKey: "col101"}, + {RowID: 1, ColumnID: 102, RowKey: "row1", ColumnKey: "col102"}, + {RowID: 1, ColumnID: 103, RowKey: "row1", ColumnKey: "col103"}, + {RowID: 2, ColumnID: 200, RowKey: "row2", ColumnKey: "col200"}, + {RowID: 2, ColumnID: 201, RowKey: "row2", ColumnKey: "col201"}, + {RowID: 2, ColumnID: 202, RowKey: "row2", ColumnKey: "col202"}, + {RowID: 2, ColumnID: 203, RowKey: "row2", ColumnKey: "col203"}, + } + + t.Run("Import unkeyed,unkeyedf", func(t *testing.T) { + // Populate data. + for _, bit := range data { + _, err := c.Query(context.Background(), "unkeyed", &pilosa.QueryRequest{ + Query: fmt.Sprintf(`Set(%d, unkeyedf=%d)`, bit.ColumnID, bit.RowID), + Remote: false, + }) + if err != nil { + t.Fatal(err) + } + } + + buf := bytes.NewBuffer(nil) + bw := bufio.NewWriter(buf) + + // Send export request. + if err := c.ExportCSV(context.Background(), "unkeyed", "unkeyedf", 0, bw); err != nil { + t.Fatal(err) + } + + got := buf.String() + + // Expected output. + exp := "" + for _, bit := range data { + exp += fmt.Sprintf("%d,%d\n", bit.RowID, bit.ColumnID) + } + + // Verify data. + if got != exp { + t.Fatalf("unexpected export data: %s", got) + } + }) + + t.Run("Import unkeyed,keyedf", func(t *testing.T) { + // Populate data. + for _, bit := range data { + _, err := c.Query(context.Background(), "unkeyed", &pilosa.QueryRequest{ + Query: fmt.Sprintf(`Set(%d, keyedf=%s)`, bit.ColumnID, bit.RowKey), + Remote: false, + }) + if err != nil { + t.Fatal(err) + } + } + + buf := bytes.NewBuffer(nil) + bw := bufio.NewWriter(buf) + + // Send export request. + if err := c.ExportCSV(context.Background(), "unkeyed", "keyedf", 0, bw); err != nil { + t.Fatal(err) + } + + got := buf.String() + + // Expected output. + exp := "" + for _, bit := range data { + exp += fmt.Sprintf("%s,%d\n", bit.RowKey, bit.ColumnID) + } + + // Verify data. + if got != exp { + t.Fatalf("unexpected export data: %s", got) + } + }) + + t.Run("Import keyed,unkeyedf", func(t *testing.T) { + // Populate data. + for _, bit := range data { + _, err := c.Query(context.Background(), "keyed", &pilosa.QueryRequest{ + Query: fmt.Sprintf(`Set("%s", unkeyedf=%d)`, bit.ColumnKey, bit.RowID), + Remote: false, + }) + if err != nil { + t.Fatal(err) + } + } + + buf := bytes.NewBuffer(nil) + bw := bufio.NewWriter(buf) + + // Send export request. + if err := c.ExportCSV(context.Background(), "keyed", "unkeyedf", 0, bw); err != nil { + t.Fatal(err) + } + + got := buf.String() + + // Expected output. + exp := "" + for _, bit := range data { + exp += fmt.Sprintf("%d,%s\n", bit.RowID, bit.ColumnKey) + } + + // Verify data. + if got != exp { + t.Fatalf("unexpected export data: %s", got) + } + }) + + t.Run("Import keyed,keyedf", func(t *testing.T) { + // Populate data. + for _, bit := range data { + _, err := c.Query(context.Background(), "keyed", &pilosa.QueryRequest{ + Query: fmt.Sprintf(`Set("%s", keyedf=%s)`, bit.ColumnKey, bit.RowKey), + Remote: false, + }) + if err != nil { + t.Fatal(err) + } + } + + buf := bytes.NewBuffer(nil) + bw := bufio.NewWriter(buf) + + // Send export request. + if err := c.ExportCSV(context.Background(), "keyed", "keyedf", 0, bw); err != nil { + t.Fatal(err) + } + + got := buf.String() + + // Expected output. + exp := "" + for _, bit := range data { + exp += fmt.Sprintf("%s,%s\n", bit.RowKey, bit.ColumnKey) + } + + // Verify data. + if got != exp { + t.Fatalf("unexpected export data: %s", got) + } + }) +} + // Ensure client can bulk import data. func TestClient_Import(t *testing.T) { cmd := test.MustRunCluster(t, 1)[0]