diff --git a/executor.go b/executor.go index 01eb7da61..f9b862422 100644 --- a/executor.go +++ b/executor.go @@ -51,6 +51,12 @@ func (e *Executor) Execute(db string, q *pql.Query, slices []uint64) (interface{ return nil, ErrDatabaseRequired } + // Ignore slices for set calls. + switch root := q.Root.(type) { + case *pql.Set: + return nil, e.executeSet(db, root) + } + // If slices aren't specified, then include all of them. if len(slices) == 0 { // Round up the number of slices. @@ -224,6 +230,27 @@ func (e *Executor) executeCount(db string, c *pql.Count, slices []uint64) (uint6 return n, nil } +// executeSet executes a set call. +func (e *Executor) executeSet(db string, c *pql.Set) error { + slice := c.ProfileID / SliceWidth + + for _, node := range e.Cluster.SliceNodes(slice) { + // Update locally if host matches. + if node.Host == e.Host { + f := e.Index().Fragment(db, c.Frame, slice) + f.Bitmap(c.ID).SetBit(c.ProfileID) + continue + } + + // Forward call to remote node otherwise. + if _, err := e.exec(node, db, &pql.Query{Root: c}, nil); err != nil { + // FIXME: Handle errors more gracefully. + return err + } + } + return nil +} + // exec executes a PQL query remotely for a set of slices on a node. func (e *Executor) exec(node *Node, db string, q *pql.Query, slices []uint64) (result interface{}, err error) { // Encode request object. @@ -287,6 +314,8 @@ func (e *Executor) exec(node *Node, db string, q *pql.Query, slices []uint64) (r return decodePairs(pb.GetPairs()), nil case *pql.Count: return pb.GetN(), nil + case *pql.Set: + return nil, nil default: panic(fmt.Sprintf("invalid node for remote exec: %T", q.Root)) } diff --git a/executor_test.go b/executor_test.go index f41eb9401..33d993e18 100644 --- a/executor_test.go +++ b/executor_test.go @@ -102,6 +102,20 @@ func TestExecutor_Execute_Count(t *testing.T) { } } +// Ensure a set query can be executed. +func TestExecutor_Execute_Set(t *testing.T) { + e := NewExecutor(NewCluster(1)) + + if _, err := e.Execute("d", MustParse(`set(id=10, frame=f, profile_id=1)`), nil); err != nil { + t.Fatal(err) + } + + f := e.Index().Fragment("d", "f", 0) + if n := f.Bitmap(10).Count(); n != 1 { + t.Fatalf("unexpected bitmap count: %d", n) + } +} + // Ensure a remote query can return a bitmap. func TestExecutor_Execute_Remote_Bitmap(t *testing.T) { c := NewCluster(2) @@ -171,6 +185,43 @@ func TestExecutor_Execute_Remote_Count(t *testing.T) { } } +// Ensure a remote query can set bits on multiple nodes. +func TestExecutor_Execute_Remote_Set(t *testing.T) { + c := NewCluster(2) + c.ReplicaN = 2 + + // Create secondary server and update second cluster node. + s := NewServer() + defer s.Close() + c.Nodes[1].Host = s.Host() + + // Mock secondary server's executor to verify arguments. + var remoteCalled bool + s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64) (interface{}, error) { + if db != `d` { + t.Fatalf("unexpected db: %s", db) + } else if query.String() != `set(id=10, frame=f, profile_id=2)` { + t.Fatalf("unexpected query: %s", query.String()) + } + remoteCalled = true + return nil, nil + } + + // Create local executor data. + e := NewExecutor(c) + if _, err := e.Execute("d", MustParse(`set(id=10, frame=f, profile_id=2)`), nil); err != nil { + t.Fatal(err) + } + + // Verify that one bit is set on both node's index. + if n := e.Index().Fragment("d", "f", 0).Bitmap(10).Count(); n != 1 { + t.Fatalf("unexpected local count: %d", n) + } + if !remoteCalled { + t.Fatalf("expected remote execution") + } +} + // Executor represents a test wrapper for pilosa.Executor. type Executor struct { *pilosa.Executor