diff --git a/cache.go b/cache.go index d3e2b6710..f0cd98603 100644 --- a/cache.go +++ b/cache.go @@ -271,7 +271,7 @@ func encodePair(p Pair) *internal.Pair { func decodePair(pb *internal.Pair) Pair { return Pair{ - ID: pb.Key, + ID: pb.Key, Count: pb.Count, } } diff --git a/executor.go b/executor.go index 9aff82cda..db67b9d0b 100644 --- a/executor.go +++ b/executor.go @@ -762,21 +762,31 @@ func (e *Executor) executeBulkSetBitmapAttrs(ctx context.Context, db string, cal // executeSetProfileAttrs executes a SetProfileAttrs() call. func (e *Executor) executeSetProfileAttrs(ctx context.Context, db string, c *pql.Call, opt *ExecOptions) error { - id, ok := c.Args["id"].(uint64) - if !ok { - return errors.New("SetProfileAttrs() id required") - } - - // Copy args and remove reserved fields. - attrs := pql.CopyArgs(c.Args) - delete(attrs, "id") - // Retrieve database. d := e.Index.DB(db) if d == nil { return ErrDatabaseNotFound } + var colName string + id, ok := c.Args["id"].(uint64) + if !ok { + // Retrieve columnLabel + columnLabel := d.columnLabel + col, ok := c.Args[columnLabel].(uint64) + if !ok { + return errors.New("SetProfileAttrs() id required") + } + id = col + colName = columnLabel + } else { + colName = "id" + } + + // Copy args and remove reserved fields. + attrs := pql.CopyArgs(c.Args) + delete(attrs, colName) + // Set attributes. if err := d.ProfileAttrStore().SetAttrs(id, attrs); err != nil { return err diff --git a/server/server_test.go b/server/server_test.go index 56ff9c109..0c717c16d 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -205,6 +205,40 @@ func TestMain_SetProfileAttrs(t *testing.T) { } } +// Ensure program can set profile attributes with columnLabel option. +func TestMain_SetProfileAttrsWithColumnOption(t *testing.T) { + m := MustRunMain() + defer m.Close() + + // Create frames. + client := m.Client() + if err := client.CreateDB(context.Background(), "d", pilosa.DBOptions{ColumnLabel: "col"}); err != nil && err != pilosa.ErrDatabaseExists { + t.Fatal(err) + } else if err := client.CreateFrame(context.Background(), "d", "x.n", pilosa.FrameOptions{}); err != nil { + t.Fatal(err) + } + + // Set bits on bitmap. + if _, err := m.Query("db=d", `SetBit(id=1, frame="x.n", col=100)`); err != nil { + t.Fatal(err) + } else if _, err := m.Query("db=d", `SetBit(id=1, frame="x.n", col=101)`); err != nil { + t.Fatal(err) + } + + // Set profile attributes. + if _, err := m.Query("db=d", `SetProfileAttrs(col=100, foo="bar")`); err != nil { + t.Fatal(err) + } + + // Query bitmap. + if res, err := m.Query("db=d&profiles=true", `Bitmap(id=1, frame="x.n")`); err != nil { + t.Fatal(err) + } else if res != `{"results":[{"attrs":{},"bits":[100,101]}],"profiles":[{"id":100,"attrs":{"foo":"bar"}}]}`+"\n" { + t.Fatalf("unexpected result: %s", res) + } + +} + // Ensure program can set bits on one cluster and then restore to a second cluster. func TestMain_FrameRestore(t *testing.T) { m0 := MustRunMain()