Merge pull request #23 from benbjohnson/set

Add set() execution

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

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