Merge pull request #1608 from travisturner/export-keys

add key translation to exports

Author: travisturner

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")
 	}
 

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]