From cf7d668b49c519b5fe6d567cc98201d54aa8284f Mon Sep 17 00:00:00 2001 From: Alan Bernstein Date: Sat, 30 Nov 2019 08:29:57 -0600 Subject: [PATCH] Support import column attrs in client --- client.go | 6 ++++ encoding/proto/proto.go | 28 +++++++++++++++ http/client.go | 50 +++++++++++++++++++++++++++ http/client_test.go | 75 +++++++++++++++++++++++++++++++++++++++++ internal/public.proto | 8 +++++ 5 files changed, 167 insertions(+) diff --git a/client.go b/client.go index f343f539e..03203b8a9 100644 --- a/client.go +++ b/client.go @@ -70,6 +70,7 @@ type InternalClient interface { SendMessage(ctx context.Context, uri *URI, msg []byte) error RetrieveShardFromURI(ctx context.Context, index, field, view string, shard uint64, uri URI) (io.ReadCloser, error) ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, req *ImportRoaringRequest) error + ImportColumnAttrs(ctx context.Context, uri *URI, index string, req *ImportColumnAttrsRequest) error } //=============== @@ -137,6 +138,11 @@ func (n nopInternalClient) ImportValue2(ctx context.Context, req *ImportValueReq func (n nopInternalClient) ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, req *ImportRoaringRequest) error { return nil } + +func (n nopInternalClient) ImportColumnAttrs(ctx context.Context, uri *URI, index string, req *ImportColumnAttrsRequest) error { + return nil +} + func (n nopInternalClient) EnsureIndex(ctx context.Context, name string, options IndexOptions) error { return nil } diff --git a/encoding/proto/proto.go b/encoding/proto/proto.go index d9e6665fb..322cca1d1 100644 --- a/encoding/proto/proto.go +++ b/encoding/proto/proto.go @@ -225,6 +225,14 @@ func (Serializer) Unmarshal(buf []byte, m pilosa.Message) error { } decodeImportRoaringRequest(msg, mt) return nil + case *pilosa.ImportColumnAttrsRequest: + msg := &internal.ImportColumnAttrsRequest{} + err := proto.Unmarshal(buf, msg) + if err != nil { + return errors.Wrap(err, "unmarshaling ImportColumnAttrsRequest") + } + decodeImportColumnAttrsRequest(msg, mt) + return nil case *pilosa.ImportResponse: msg := &internal.ImportResponse{} err := proto.Unmarshal(buf, msg) @@ -318,6 +326,8 @@ func encodeToProto(m pilosa.Message) proto.Message { return encodeImportValueRequest(mt) case *pilosa.ImportRoaringRequest: return encodeImportRoaringRequest(mt) + case *pilosa.ImportColumnAttrsRequest: + return encodeImportColumnAttrsRequest(mt) case *pilosa.ImportResponse: return encodeImportResponse(mt) case *pilosa.BlockDataRequest: @@ -395,6 +405,16 @@ func encodeImportRoaringRequest(m *pilosa.ImportRoaringRequest) *internal.Import } } +func encodeImportColumnAttrsRequest(m *pilosa.ImportColumnAttrsRequest) *internal.ImportColumnAttrsRequest { + return &internal.ImportColumnAttrsRequest{ + Index: m.Index, + Shard: m.Shard, + AttrKey: m.AttrKey, + AttrVals: m.AttrVals, + ColumnIDs: m.ColumnIDs, + } +} + func encodeQueryRequest(m *pilosa.QueryRequest) *internal.QueryRequest { r := &internal.QueryRequest{ Query: m.Query, @@ -1010,6 +1030,14 @@ func decodeImportRoaringRequest(pb *internal.ImportRoaringRequest, m *pilosa.Imp m.Views = views } +func decodeImportColumnAttrsRequest(pb *internal.ImportColumnAttrsRequest, m *pilosa.ImportColumnAttrsRequest) { + m.Index = pb.Index + m.Shard = pb.Shard + m.AttrKey = pb.AttrKey + m.AttrVals = pb.AttrVals + m.ColumnIDs = pb.ColumnIDs +} + func decodeImportResponse(pb *internal.ImportResponse, m *pilosa.ImportResponse) { m.Err = pb.Err } diff --git a/http/client.go b/http/client.go index 1b4a9e976..1c5de0197 100644 --- a/http/client.go +++ b/http/client.go @@ -696,6 +696,56 @@ func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pilosa.URI, ind return nil } +// ImportColumnAttrs does bulk import of column attrs +func (c *InternalClient) ImportColumnAttrs(ctx context.Context, uri *pilosa.URI, index string, req *pilosa.ImportColumnAttrsRequest) error { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ImportRoaring") + defer span.Finish() + + if index == "" { + return pilosa.ErrIndexRequired + } + if uri == nil { + uri = c.defaultURI + } + + url := fmt.Sprintf("%s/index/%s/import-column-attrs", uri, index) + + // Marshal data to protobuf. + data, err := c.serializer.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshal import-column-attrs request") + } + + // Generate HTTP request. + httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(data)) + if err != nil { + return errors.Wrap(err, "creating request") + } + httpReq.Header.Set("Content-Type", "application/x-protobuf") + httpReq.Header.Set("Accept", "application/x-protobuf") + httpReq.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Execute request against the host. + resp, err := c.executeRequest(httpReq.WithContext(ctx)) + if err != nil { + return err + } + defer resp.Body.Close() + + dec := json.NewDecoder(resp.Body) + rbody := &pilosa.ImportResponse{} + err = dec.Decode(rbody) + // Decode can return EOF when no error occurred. helpful! + if err != nil && err != io.EOF { + return errors.Wrap(err, "decoding response body") + } + if rbody.Err != "" { + return errors.Wrap(errors.New(rbody.Err), "importing roaring") + } + + return nil +} + // ExportCSV bulk exports data for a single shard from a host to CSV format. func (c *InternalClient) ExportCSV(ctx context.Context, index, field string, shard uint64, w io.Writer) error { span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ExportCSV") diff --git a/http/client_test.go b/http/client_test.go index bc889b10d..f8743555e 100644 --- a/http/client_test.go +++ b/http/client_test.go @@ -22,6 +22,7 @@ import ( "fmt" gohttp "net/http" "reflect" + "strconv" "testing" "time" @@ -394,6 +395,60 @@ func TestClient_Import(t *testing.T) { } } +// Ensure client can bulk import column attrs. +func TestClient_ImportColumnAttrs(t *testing.T) { + cluster := test.MustNewCluster(t, 2) + for _, c := range cluster { + c.Config.Cluster.ReplicaN = 2 + } + err := cluster.Start() + if err != nil { + t.Fatalf("starting cluster: %v", err) + } + defer cluster.Close() + + ctx := context.Background() + _, err = cluster[0].API.CreateIndex(ctx, "i", pilosa.IndexOptions{}) + if err != nil { + t.Fatalf("creating index: %v", err) + } + _, err = cluster[0].API.CreateField(ctx, "i", "f", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 100)) + if err != nil { + t.Fatalf("creating field: %v", err) + } + _, err = cluster[0].API.Query(ctx, &pilosa.QueryRequest{Index: "i", Query: "Set(0, f=0) Set(1, f=0) Set(2, f=0) Set(3, f=0) Set(4, f=0)"}) + if err != nil { + t.Fatalf("querying: %v", err) + } + + attrKey := "k" + // Send import request. + host := cluster[0].URL() + c := MustNewClient(host, http.GetHTTPClient(nil)) + colAttrsReq := makeImportColumnAttrsRequest("i", 0, attrKey) + if err := c.ImportColumnAttrs(ctx, &cluster[1].API.Node().URI, "i", colAttrsReq); err != nil { + t.Fatal(err) + } + + // Verify data. + pql := "Options(Row(f=0), columnAttrs=true)" + res, err := cluster[1].API.Query(ctx, &pilosa.QueryRequest{Index: "i", Query: pql}) + if err != nil { + t.Fatal(err) + } + if len(res.ColumnAttrSets) != 5 { + t.Fatal("incorrect number of column attrs set") + } + + for _, v := range res.ColumnAttrSets { + attrVal := attrFun(v.ID) + if attrVal != v.Attrs[attrKey] { + t.Fatal(err) + } + } + +} + // Ensure client can bulk import data. func TestClient_ImportRoaring(t *testing.T) { cluster := test.MustNewCluster(t, 2) @@ -1195,3 +1250,23 @@ func makeImportRoaringRequest(clear bool, viewData string) *pilosa.ImportRoaring }, } } + +func attrFun(id uint64) string { + return strconv.FormatInt(int64(id), 10) +} + +func makeImportColumnAttrsRequest(index string, shard int64, attrKey string) *pilosa.ImportColumnAttrsRequest { + colIDs := make([]uint64, 0, 5) + attrVals := make([]string, 0, 5) + for n := uint64(0); n < 5; n++ { + colIDs = append(colIDs, n) + attrVals = append(attrVals, attrFun(n)) + } + return &pilosa.ImportColumnAttrsRequest{ + Index: index, + Shard: shard, + AttrKey: attrKey, + ColumnIDs: colIDs, + AttrVals: attrVals, + } +} diff --git a/internal/public.proto b/internal/public.proto index ca6cf26d6..b8d569f85 100644 --- a/internal/public.proto +++ b/internal/public.proto @@ -130,3 +130,11 @@ message ImportRoaringRequest { bool Clear = 1; repeated ImportRoaringRequestView views = 2; } + +message ImportColumnAttrsRequest { + string Index = 1; + int64 Shard = 2; + string AttrKey = 3; + repeated string AttrVals = 4; + repeated uint64 ColumnIDs = 5; +} \ No newline at end of file