Add context to Client, Executor, & Handler.

This changes the API of the types but does not alter the functionality
since only the `context.Background()` is currently being used. Adding
`Context` will help handle fault tolerance in the future by allowing
timeouts to be propagated across calls to different nodes.
This commit is contained in:
Ben Johnson 2016-11-10 13:37:00 -07:00
parent e92228d641
commit 8f78f854dd
9 changed files with 236 additions and 157 deletions

163
client.go
View file

@ -3,6 +3,7 @@ package pilosa
import (
"archive/tar"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
@ -45,14 +46,22 @@ func NewClient(host string) (*Client, error) {
func (c *Client) Host() string { return c.host }
// SliceN returns the number of slices on a server.
func (c *Client) SliceN() (uint64, error) {
func (c *Client) SliceN(ctx context.Context) (uint64, error) {
// Execute request against the host.
u := url.URL{
Scheme: "http",
Host: c.host,
Path: "/slices/max",
}
resp, err := c.HTTPClient.Get(u.String())
// Build request.
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return 0, err
}
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return 0, err
}
@ -69,14 +78,22 @@ func (c *Client) SliceN() (uint64, error) {
}
// Schema returns all database and frame schema information.
func (c *Client) Schema() ([]*DBInfo, error) {
func (c *Client) Schema(ctx context.Context) ([]*DBInfo, error) {
// Execute request against the host.
u := url.URL{
Scheme: "http",
Host: c.host,
Path: "/schema",
}
resp, err := c.HTTPClient.Get(u.String())
// Build request.
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return nil, err
}
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -92,7 +109,7 @@ func (c *Client) Schema() ([]*DBInfo, error) {
}
// FragmentNodes returns a list of nodes that own a slice.
func (c *Client) FragmentNodes(db string, slice uint64) ([]*Node, error) {
func (c *Client) FragmentNodes(ctx context.Context, db string, slice uint64) ([]*Node, error) {
// Execute request against the host.
u := url.URL{
Scheme: "http",
@ -100,7 +117,15 @@ func (c *Client) FragmentNodes(db string, slice uint64) ([]*Node, error) {
Path: "/fragment/nodes",
RawQuery: (url.Values{"db": {db}, "slice": {strconv.FormatUint(slice, 10)}}).Encode(),
}
resp, err := c.HTTPClient.Get(u.String())
// Build request.
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return nil, err
}
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -117,7 +142,7 @@ func (c *Client) FragmentNodes(db string, slice uint64) ([]*Node, error) {
}
// ExecuteQuery executes query against db on the server.
func (c *Client) ExecuteQuery(db, query string, allowRedirect bool) (result interface{}, err error) {
func (c *Client) ExecuteQuery(ctx context.Context, db, query string, allowRedirect bool) (result interface{}, err error) {
if db == "" {
return nil, ErrDatabaseRequired
} else if query == "" {
@ -145,7 +170,7 @@ func (c *Client) ExecuteQuery(db, query string, allowRedirect bool) (result inte
req.Header.Set("Accept", "application/x-protobuf")
// Execute request against the host.
resp, err := c.HTTPClient.Do(req)
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -170,7 +195,7 @@ func (c *Client) ExecuteQuery(db, query string, allowRedirect bool) (result inte
}
// Import bulk imports bits for a single slice to a host.
func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error {
func (c *Client) Import(ctx context.Context, db, frame string, slice uint64, bits []Bit) error {
if db == "" {
return ErrDatabaseRequired
} else if frame == "" {
@ -183,15 +208,14 @@ func (c *Client) Import(db, frame string, slice uint64, bits []Bit) error {
}
// Retrieve a list of nodes that own the slice.
nodes, err := c.FragmentNodes(db, slice)
nodes, err := c.FragmentNodes(ctx, db, slice)
if err != nil {
return fmt.Errorf("slice nodes: %s", err)
}
// Import to each node.
for _, node := range nodes {
if err := c.importNode(node, buf); err != nil {
if err := c.importNode(ctx, node, buf); err != nil {
return fmt.Errorf("import node: host=%s, err=%s", node.Host, err)
}
}
@ -219,7 +243,7 @@ func MarshalImportPayload(db, frame string, slice uint64, bits []Bit) ([]byte, e
}
// importNode sends a pre-marshaled import request to a node.
func (c *Client) importNode(node *Node, buf []byte) error {
func (c *Client) importNode(ctx context.Context, node *Node, buf []byte) error {
// Create URL & HTTP request.
u := url.URL{Scheme: "http", Host: node.Host, Path: "/import"}
req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
@ -231,7 +255,7 @@ func (c *Client) importNode(node *Node, buf []byte) error {
req.Header.Set("Accept", "application/x-protobuf")
// Execute request against the host.
resp, err := c.HTTPClient.Do(req)
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return err
}
@ -256,7 +280,7 @@ func (c *Client) importNode(node *Node, buf []byte) error {
}
// ExportCSV bulk exports data for a single slice from a host to CSV format.
func (c *Client) ExportCSV(db, frame string, slice uint64, w io.Writer) error {
func (c *Client) ExportCSV(ctx context.Context, db, frame string, slice uint64, w io.Writer) error {
if db == "" {
return ErrDatabaseRequired
} else if frame == "" {
@ -264,7 +288,7 @@ func (c *Client) ExportCSV(db, frame string, slice uint64, w io.Writer) error {
}
// Retrieve a list of nodes that own the slice.
nodes, err := c.FragmentNodes(db, slice)
nodes, err := c.FragmentNodes(ctx, db, slice)
if err != nil {
return fmt.Errorf("slice nodes: %s", err)
}
@ -274,7 +298,7 @@ func (c *Client) ExportCSV(db, frame string, slice uint64, w io.Writer) error {
for _, i := range rand.Perm(len(nodes)) {
node := nodes[i]
if err := c.exportNodeCSV(node, db, frame, slice, w); err != nil {
if err := c.exportNodeCSV(ctx, node, db, frame, slice, w); err != nil {
e = fmt.Errorf("export node: host=%s, err=%s", node.Host, err)
continue
} else {
@ -286,7 +310,7 @@ func (c *Client) ExportCSV(db, frame string, slice uint64, w io.Writer) error {
}
// exportNode copies a CSV export from a node to w.
func (c *Client) exportNodeCSV(node *Node, db, frame string, slice uint64, w io.Writer) error {
func (c *Client) exportNodeCSV(ctx context.Context, node *Node, db, frame string, slice uint64, w io.Writer) error {
// Create URL.
u := url.URL{
Scheme: "http",
@ -307,7 +331,7 @@ func (c *Client) exportNodeCSV(node *Node, db, frame string, slice uint64, w io.
req.Header.Set("Accept", "text/csv")
// Execute request against the host.
resp, err := c.HTTPClient.Do(req)
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return err
}
@ -327,7 +351,7 @@ func (c *Client) exportNodeCSV(node *Node, db, frame string, slice uint64, w io.
}
// BackupTo backs up an entire frame from a cluster to w.
func (c *Client) BackupTo(w io.Writer, db, frame string) error {
func (c *Client) BackupTo(ctx context.Context, w io.Writer, db, frame string) error {
if db == "" {
return ErrDatabaseRequired
} else if frame == "" {
@ -338,14 +362,14 @@ func (c *Client) BackupTo(w io.Writer, db, frame string) error {
tw := tar.NewWriter(w)
// Find the maximum number of slices.
sliceN, err := c.SliceN()
sliceN, err := c.SliceN(ctx)
if err != nil {
return fmt.Errorf("slice n: %s", err)
}
// Backup every slice to the tar file.
for i := uint64(0); i <= sliceN; i++ {
if err := c.backupSliceTo(tw, db, frame, i); err != nil {
if err := c.backupSliceTo(ctx, tw, db, frame, i); err != nil {
return err
}
}
@ -359,9 +383,9 @@ func (c *Client) BackupTo(w io.Writer, db, frame string) error {
}
// backupSliceTo backs up a single slice to tw.
func (c *Client) backupSliceTo(tw *tar.Writer, db, frame string, slice uint64) error {
func (c *Client) backupSliceTo(ctx context.Context, tw *tar.Writer, db, frame string, slice uint64) error {
// Return error if unable to backup from any slice.
r, err := c.BackupSlice(db, frame, slice)
r, err := c.BackupSlice(ctx, db, frame, slice)
if err != nil {
return fmt.Errorf("backup slice: slice=%d, err=%s", slice, err)
} else if r == nil {
@ -397,16 +421,16 @@ func (c *Client) backupSliceTo(tw *tar.Writer, db, frame string, slice uint64) e
// BackupSlice retrieves a streaming backup from a single slice.
// This function tries slice owners until one succeeds.
func (c *Client) BackupSlice(db, frame string, slice uint64) (io.ReadCloser, error) {
func (c *Client) BackupSlice(ctx context.Context, db, frame string, slice uint64) (io.ReadCloser, error) {
// Retrieve a list of nodes that own the slice.
nodes, err := c.FragmentNodes(db, slice)
nodes, err := c.FragmentNodes(ctx, db, slice)
if err != nil {
return nil, fmt.Errorf("slice nodes: %s", err)
}
// Try to backup slice from each one until successful.
for _, i := range rand.Perm(len(nodes)) {
r, err := c.backupSliceNode(db, frame, slice, nodes[i])
r, err := c.backupSliceNode(ctx, db, frame, slice, nodes[i])
if err == nil {
return r, nil // successfully attached
} else if err == ErrFragmentNotFound {
@ -420,7 +444,7 @@ func (c *Client) BackupSlice(db, frame string, slice uint64) (io.ReadCloser, err
return nil, fmt.Errorf("unable to connect to any owner")
}
func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) (io.ReadCloser, error) {
func (c *Client) backupSliceNode(ctx context.Context, db, frame string, slice uint64, node *Node) (io.ReadCloser, error) {
u := url.URL{
Scheme: "http",
Host: node.Host,
@ -431,7 +455,15 @@ func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) (io
"slice": {strconv.FormatUint(slice, 10)},
}.Encode(),
}
resp, err := c.HTTPClient.Get(u.String())
// Build request.
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return nil, err
}
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -449,7 +481,7 @@ func (c *Client) backupSliceNode(db, frame string, slice uint64, node *Node) (io
}
// RestoreFrom restores a frame from a backup file to an entire cluster.
func (c *Client) RestoreFrom(r io.Reader, db, frame string) error {
func (c *Client) RestoreFrom(ctx context.Context, r io.Reader, db, frame string) error {
if db == "" {
return ErrDatabaseRequired
} else if frame == "" {
@ -481,16 +513,16 @@ func (c *Client) RestoreFrom(r io.Reader, db, frame string) error {
}
// Restore file to all nodes that own it.
if err := c.restoreSliceFrom(buf.Bytes(), db, frame, slice); err != nil {
if err := c.restoreSliceFrom(ctx, buf.Bytes(), db, frame, slice); err != nil {
return err
}
}
}
// restoreSliceFrom restores a single slice to all owning nodes.
func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) error {
func (c *Client) restoreSliceFrom(ctx context.Context, buf []byte, db, frame string, slice uint64) error {
// Retrieve a list of nodes that own the slice.
nodes, err := c.FragmentNodes(db, slice)
nodes, err := c.FragmentNodes(ctx, db, slice)
if err != nil {
return fmt.Errorf("slice nodes: %s", err)
}
@ -507,7 +539,15 @@ func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) er
"slice": {strconv.FormatUint(slice, 10)},
}.Encode(),
}
resp, err := c.HTTPClient.Post(u.String(), "application/octet-stream", bytes.NewReader(buf))
// Build request.
req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/octet-stream")
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return err
}
@ -523,7 +563,7 @@ func (c *Client) restoreSliceFrom(buf []byte, db, frame string, slice uint64) er
}
// RestoreFrame restores an entire frame from a host in another cluster.
func (c *Client) RestoreFrame(host, db, frame string) error {
func (c *Client) RestoreFrame(ctx context.Context, host, db, frame string) error {
u := url.URL{
Scheme: "http",
Host: c.Host(),
@ -534,7 +574,16 @@ func (c *Client) RestoreFrame(host, db, frame string) error {
"frame": {frame},
}.Encode(),
}
resp, err := c.HTTPClient.Post(u.String(), "application/octet-stream", nil)
// Build request.
req, err := http.NewRequest("POST", u.String(), nil)
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/octet-stream")
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return err
}
@ -550,7 +599,7 @@ func (c *Client) RestoreFrame(host, db, frame string) error {
// FragmentBlocks returns a list of block checksums for a fragment on a host.
// Only returns blocks which contain data.
func (c *Client) FragmentBlocks(db, frame string, slice uint64) ([]FragmentBlock, error) {
func (c *Client) FragmentBlocks(ctx context.Context, db, frame string, slice uint64) ([]FragmentBlock, error) {
u := url.URL{
Scheme: "http",
Host: c.host,
@ -561,7 +610,15 @@ func (c *Client) FragmentBlocks(db, frame string, slice uint64) ([]FragmentBlock
"slice": {strconv.FormatUint(slice, 10)},
}.Encode(),
}
resp, err := c.HTTPClient.Get(u.String())
// Build request.
req, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return nil, err
}
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -585,7 +642,7 @@ func (c *Client) FragmentBlocks(db, frame string, slice uint64) ([]FragmentBlock
}
// BlockData returns bitmap/profile id pairs for a block.
func (c *Client) BlockData(db, frame string, slice uint64, block int) ([]uint64, []uint64, error) {
func (c *Client) BlockData(ctx context.Context, db, frame string, slice uint64, block int) ([]uint64, []uint64, error) {
buf, err := proto.Marshal(&internal.BlockDataRequest{
DB: proto.String(db),
Frame: proto.String(frame),
@ -605,7 +662,7 @@ func (c *Client) BlockData(db, frame string, slice uint64, block int) ([]uint64,
req.Header.Set("Content-Length", strconv.Itoa(len(buf)))
req.Header.Set("Accept", "application/protobuf")
resp, err := c.HTTPClient.Do(req)
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, nil, err
}
@ -631,7 +688,7 @@ func (c *Client) BlockData(db, frame string, slice uint64, block int) ([]uint64,
}
// ProfileAttrDiff returns data from differing blocks on a remote host.
func (c *Client) ProfileAttrDiff(db string, blks []AttrBlock) (map[uint64]map[string]interface{}, error) {
func (c *Client) ProfileAttrDiff(ctx context.Context, db string, blks []AttrBlock) (map[uint64]map[string]interface{}, error) {
u := url.URL{
Scheme: "http",
Host: c.host,
@ -645,8 +702,15 @@ func (c *Client) ProfileAttrDiff(db string, blks []AttrBlock) (map[uint64]map[st
return nil, err
}
// Send request.
resp, err := c.HTTPClient.Post(u.String(), "application/json", bytes.NewReader(buf))
// Build request.
req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}
@ -668,7 +732,7 @@ func (c *Client) ProfileAttrDiff(db string, blks []AttrBlock) (map[uint64]map[st
}
// BitmapAttrDiff returns data from differing blocks on a remote host.
func (c *Client) BitmapAttrDiff(db, frame string, blks []AttrBlock) (map[uint64]map[string]interface{}, error) {
func (c *Client) BitmapAttrDiff(ctx context.Context, db, frame string, blks []AttrBlock) (map[uint64]map[string]interface{}, error) {
u := url.URL{
Scheme: "http",
Host: c.host,
@ -682,8 +746,15 @@ func (c *Client) BitmapAttrDiff(db, frame string, blks []AttrBlock) (map[uint64]
return nil, err
}
// Send request.
resp, err := c.HTTPClient.Post(u.String(), "application/json", bytes.NewReader(buf))
// Build request.
req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
// Execute request.
resp, err := c.HTTPClient.Do(req.WithContext(ctx))
if err != nil {
return nil, err
}

View file

@ -2,6 +2,7 @@ package pilosa_test
import (
"bytes"
"context"
"reflect"
"testing"
@ -27,7 +28,7 @@ func TestClient_Import(t *testing.T) {
// Send import request.
c := MustNewClient(s.Host())
if err := c.Import("d", "f", 0, []pilosa.Bit{
if err := c.Import(context.Background(), "d", "f", 0, []pilosa.Bit{
{BitmapID: 0, ProfileID: 1},
{BitmapID: 0, ProfileID: 5},
{BitmapID: 200, ProfileID: 6},
@ -65,12 +66,12 @@ func TestClient_BackupRestore(t *testing.T) {
// Backup from frame.
var buf bytes.Buffer
if err := c.BackupTo(&buf, "d", "f"); err != nil {
if err := c.BackupTo(context.Background(), &buf, "d", "f"); err != nil {
t.Fatal(err)
}
// Restore to a different frame.
if err := c.RestoreFrom(&buf, "x", "y"); err != nil {
if err := c.RestoreFrom(context.Background(), &buf, "x", "y"); err != nil {
t.Fatal(err)
}
@ -110,7 +111,7 @@ func TestClient_FragmentBlocks(t *testing.T) {
// Retrieve blocks.
c := MustNewClient(s.Host())
blocks, err := c.FragmentBlocks("d", "f", 0)
blocks, err := c.FragmentBlocks(context.Background(), "d", "f", 0)
if err != nil {
t.Fatal(err)
} else if len(blocks) != 2 {

View file

@ -2,6 +2,7 @@ package pilosa
import (
"bytes"
"context"
"errors"
"fmt"
"io/ioutil"
@ -38,7 +39,7 @@ func NewExecutor() *Executor {
}
// Execute executes a PQL query.
func (e *Executor) Execute(db string, q *pql.Query, slices []uint64, opt *ExecOptions) ([]interface{}, error) {
func (e *Executor) Execute(ctx context.Context, db string, q *pql.Query, slices []uint64, opt *ExecOptions) ([]interface{}, error) {
// Verify that a database is set.
if db == "" {
return nil, ErrDatabaseRequired
@ -64,13 +65,13 @@ func (e *Executor) Execute(db string, q *pql.Query, slices []uint64, opt *ExecOp
// Optimize handling for bulk attribute insertion.
if hasOnlySetBitmapAttrs(q.Calls) {
return e.executeBulkSetBitmapAttrs(db, q.Calls, opt)
return e.executeBulkSetBitmapAttrs(ctx, db, q.Calls, opt)
}
// Execute each call serially.
results := make([]interface{}, 0, len(q.Calls))
for _, call := range q.Calls {
v, err := e.executeCall(db, call, slices, opt)
v, err := e.executeCall(ctx, db, call, slices, opt)
if err != nil {
return nil, err
}
@ -80,37 +81,37 @@ func (e *Executor) Execute(db string, q *pql.Query, slices []uint64, opt *ExecOp
}
// executeCall executes a call.
func (e *Executor) executeCall(db string, c pql.Call, slices []uint64, opt *ExecOptions) (interface{}, error) {
func (e *Executor) executeCall(ctx context.Context, db string, c pql.Call, slices []uint64, opt *ExecOptions) (interface{}, error) {
switch c := c.(type) {
case pql.BitmapCall:
return e.executeBitmapCall(db, c, slices, opt)
return e.executeBitmapCall(ctx, db, c, slices, opt)
case *pql.ClearBit:
return e.executeClearBit(db, c, opt)
return e.executeClearBit(ctx, db, c, opt)
case *pql.Count:
return e.executeCount(db, c, slices, opt)
return e.executeCount(ctx, db, c, slices, opt)
case *pql.Profile:
return e.executeProfile(db, c, opt)
return e.executeProfile(ctx, db, c, opt)
case *pql.SetBit:
return e.executeSetBit(db, c, opt)
return e.executeSetBit(ctx, db, c, opt)
case *pql.SetBitmapAttrs:
return nil, e.executeSetBitmapAttrs(db, c, opt)
return nil, e.executeSetBitmapAttrs(ctx, db, c, opt)
case *pql.SetProfileAttrs:
return nil, e.executeSetProfileAttrs(db, c, opt)
return nil, e.executeSetProfileAttrs(ctx, db, c, opt)
case *pql.TopN:
return e.executeTopN(db, c, slices, opt)
return e.executeTopN(ctx, db, c, slices, opt)
default:
panic("unreachable")
}
}
// executeBitmapCall executes a call that returns a bitmap.
func (e *Executor) executeBitmapCall(db string, c pql.BitmapCall, slices []uint64, opt *ExecOptions) (*Bitmap, error) {
func (e *Executor) executeBitmapCall(ctx context.Context, db string, c pql.BitmapCall, slices []uint64, opt *ExecOptions) (*Bitmap, error) {
other := NewBitmap()
for node, nodeSlices := range e.slicesByNode(db, slices) {
// Execute locally if the hostname matches.
if node.Host == e.Host {
for _, slice := range nodeSlices {
bm, err := e.executeBitmapCallSlice(db, c, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, c, slice)
if err != nil {
return nil, err
}
@ -120,7 +121,7 @@ func (e *Executor) executeBitmapCall(db string, c pql.BitmapCall, slices []uint6
}
// Otherwise execute remotely.
res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
res, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
if err != nil {
return nil, err
}
@ -143,18 +144,18 @@ func (e *Executor) executeBitmapCall(db string, c pql.BitmapCall, slices []uint6
}
// executeBitmapCallSlice executes a bitmap call for a single slice.
func (e *Executor) executeBitmapCallSlice(db string, c pql.BitmapCall, slice uint64) (*Bitmap, error) {
func (e *Executor) executeBitmapCallSlice(ctx context.Context, db string, c pql.BitmapCall, slice uint64) (*Bitmap, error) {
switch c := c.(type) {
case *pql.Bitmap:
return e.executeBitmapSlice(db, c, slice)
return e.executeBitmapSlice(ctx, db, c, slice)
case *pql.Difference:
return e.executeDifferenceSlice(db, c, slice)
return e.executeDifferenceSlice(ctx, db, c, slice)
case *pql.Intersect:
return e.executeIntersectSlice(db, c, slice)
return e.executeIntersectSlice(ctx, db, c, slice)
case *pql.Range:
return e.executeRangeSlice(db, c, slice)
return e.executeRangeSlice(ctx, db, c, slice)
case *pql.Union:
return e.executeUnionSlice(db, c, slice)
return e.executeUnionSlice(ctx, db, c, slice)
default:
panic("unreachable")
}
@ -163,9 +164,9 @@ func (e *Executor) executeBitmapCallSlice(db string, c pql.BitmapCall, slice uin
// executeTopN executes a TopN() call.
// This first performs the TopN() to determine the top results and then
// requeries to retrieve the full counts for each of the top results.
func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) {
func (e *Executor) executeTopN(ctx context.Context, db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) {
// Execute original query.
pairs, err := e.executeTopNSlices(db, c, slices, opt)
pairs, err := e.executeTopNSlices(ctx, db, c, slices, opt)
if err != nil {
return nil, err
}
@ -182,10 +183,10 @@ func (e *Executor) executeTopN(db string, c *pql.TopN, slices []uint64, opt *Exe
other.BitmapIDs = Pairs(pairs).Keys()
sort.Sort(uint64Slice(other.BitmapIDs))
return e.executeTopNSlices(db, &other, slices, opt)
return e.executeTopNSlices(ctx, db, &other, slices, opt)
}
func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) {
func (e *Executor) executeTopNSlices(ctx context.Context, db string, c *pql.TopN, slices []uint64, opt *ExecOptions) ([]Pair, error) {
slicesByNode := e.slicesByNode(db, slices)
type resp struct {
@ -198,13 +199,13 @@ func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, op
go func(node *Node, nodeSlices []uint64) {
// Execute locally if the hostname matches.
if node.Host == e.Host {
pairs, err := e.executeTopNSlicesLocal(db, c, nodeSlices)
pairs, err := e.executeTopNSlicesLocal(ctx, db, c, nodeSlices)
ch <- resp{pairs: pairs, err: err}
return
}
// Otherwise execute remotely.
res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
res, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
if err != nil {
ch <- resp{err: err}
return
@ -234,7 +235,7 @@ func (e *Executor) executeTopNSlices(db string, c *pql.TopN, slices []uint64, op
return results, nil
}
func (e *Executor) executeTopNSlicesLocal(db string, c *pql.TopN, slices []uint64) ([]Pair, error) {
func (e *Executor) executeTopNSlicesLocal(ctx context.Context, db string, c *pql.TopN, slices []uint64) ([]Pair, error) {
type resp struct {
pairs []Pair
err error
@ -244,7 +245,7 @@ func (e *Executor) executeTopNSlicesLocal(db string, c *pql.TopN, slices []uint6
// Execute TopN() in parallel across slices.
for _, slice := range slices {
go func(slice uint64) {
pairs, err := e.executeTopNSlice(db, c, slice)
pairs, err := e.executeTopNSlice(ctx, db, c, slice)
ch <- resp{pairs: pairs, err: err}
}(slice)
}
@ -263,11 +264,11 @@ func (e *Executor) executeTopNSlicesLocal(db string, c *pql.TopN, slices []uint6
}
// executeTopNSlice executes a TopN call for a single slice.
func (e *Executor) executeTopNSlice(db string, c *pql.TopN, slice uint64) ([]Pair, error) {
func (e *Executor) executeTopNSlice(ctx context.Context, db string, c *pql.TopN, slice uint64) ([]Pair, error) {
// Retrieve bitmap used to intersect.
var src *Bitmap
if c.Src != nil {
bm, err := e.executeBitmapCallSlice(db, c.Src, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, c.Src, slice)
if err != nil {
return nil, err
}
@ -295,10 +296,10 @@ func (e *Executor) executeTopNSlice(db string, c *pql.TopN, slice uint64) ([]Pai
}
// executeDifferenceSlice executes a difference() call for a local slice.
func (e *Executor) executeDifferenceSlice(db string, c *pql.Difference, slice uint64) (*Bitmap, error) {
func (e *Executor) executeDifferenceSlice(ctx context.Context, db string, c *pql.Difference, slice uint64) (*Bitmap, error) {
var other *Bitmap
for i, input := range c.Inputs {
bm, err := e.executeBitmapCallSlice(db, input, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, input, slice)
if err != nil {
return nil, err
}
@ -313,7 +314,7 @@ func (e *Executor) executeDifferenceSlice(db string, c *pql.Difference, slice ui
return other, nil
}
func (e *Executor) executeBitmapSlice(db string, c *pql.Bitmap, slice uint64) (*Bitmap, error) {
func (e *Executor) executeBitmapSlice(ctx context.Context, db string, c *pql.Bitmap, slice uint64) (*Bitmap, error) {
frame := c.Frame
if frame == "" {
frame = DefaultFrame
@ -327,10 +328,10 @@ func (e *Executor) executeBitmapSlice(db string, c *pql.Bitmap, slice uint64) (*
}
// executeIntersectSlice executes a intersect() call for a local slice.
func (e *Executor) executeIntersectSlice(db string, c *pql.Intersect, slice uint64) (*Bitmap, error) {
func (e *Executor) executeIntersectSlice(ctx context.Context, db string, c *pql.Intersect, slice uint64) (*Bitmap, error) {
var other *Bitmap
for i, input := range c.Inputs {
bm, err := e.executeBitmapCallSlice(db, input, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, input, slice)
if err != nil {
return nil, err
}
@ -346,7 +347,7 @@ func (e *Executor) executeIntersectSlice(db string, c *pql.Intersect, slice uint
}
// executeRangeSlice executes a range() call for a local slice.
func (e *Executor) executeRangeSlice(db string, c *pql.Range, slice uint64) (*Bitmap, error) {
func (e *Executor) executeRangeSlice(ctx context.Context, db string, c *pql.Range, slice uint64) (*Bitmap, error) {
frame := c.Frame
if frame == "" {
frame = DefaultFrame
@ -360,10 +361,10 @@ func (e *Executor) executeRangeSlice(db string, c *pql.Range, slice uint64) (*Bi
}
// executeUnionSlice executes a union() call for a local slice.
func (e *Executor) executeUnionSlice(db string, c *pql.Union, slice uint64) (*Bitmap, error) {
func (e *Executor) executeUnionSlice(ctx context.Context, db string, c *pql.Union, slice uint64) (*Bitmap, error) {
var other *Bitmap
for i, input := range c.Inputs {
bm, err := e.executeBitmapCallSlice(db, input, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, input, slice)
if err != nil {
return nil, err
}
@ -379,13 +380,13 @@ func (e *Executor) executeUnionSlice(db string, c *pql.Union, slice uint64) (*Bi
}
// executeCount executes a count() call.
func (e *Executor) executeCount(db string, c *pql.Count, slices []uint64, opt *ExecOptions) (uint64, error) {
func (e *Executor) executeCount(ctx context.Context, db string, c *pql.Count, slices []uint64, opt *ExecOptions) (uint64, error) {
var n uint64
for node, nodeSlices := range e.slicesByNode(db, slices) {
// Execute locally if the hostname matches.
if node.Host == e.Host {
for _, slice := range nodeSlices {
bm, err := e.executeBitmapCallSlice(db, c.Input, slice)
bm, err := e.executeBitmapCallSlice(ctx, db, c.Input, slice)
if err != nil {
return 0, err
}
@ -395,7 +396,7 @@ func (e *Executor) executeCount(db string, c *pql.Count, slices []uint64, opt *E
}
// Otherwise execute remotely.
res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
res, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nodeSlices, opt)
if err != nil {
return 0, err
}
@ -406,12 +407,12 @@ func (e *Executor) executeCount(db string, c *pql.Count, slices []uint64, opt *E
// executeProfile executes a Profile() call.
// This call only executes locally since the profile attibutes are stored locally.
func (e *Executor) executeProfile(db string, c *pql.Profile, opt *ExecOptions) (*Profile, error) {
func (e *Executor) executeProfile(ctx context.Context, db string, c *pql.Profile, opt *ExecOptions) (*Profile, error) {
panic("FIXME: impl: e.Index.ProfileAttr(c.ID)")
}
// executeClearBit executes a ClearBit() call.
func (e *Executor) executeClearBit(db string, c *pql.ClearBit, opt *ExecOptions) (bool, error) {
func (e *Executor) executeClearBit(ctx context.Context, db string, c *pql.ClearBit, opt *ExecOptions) (bool, error) {
slice := c.ProfileID / SliceWidth
ret := false
for _, node := range e.Cluster.FragmentNodes(db, slice) {
@ -436,7 +437,7 @@ func (e *Executor) executeClearBit(db string, c *pql.ClearBit, opt *ExecOptions)
}
// Forward call to remote node otherwise.
if res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt); err != nil {
if res, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt); err != nil {
return false, err
} else {
ret = res[0].(bool)
@ -446,7 +447,7 @@ func (e *Executor) executeClearBit(db string, c *pql.ClearBit, opt *ExecOptions)
}
// executeSetBit executes a SetBit() call.
func (e *Executor) executeSetBit(db string, c *pql.SetBit, opt *ExecOptions) (bool, error) {
func (e *Executor) executeSetBit(ctx context.Context, db string, c *pql.SetBit, opt *ExecOptions) (bool, error) {
slice := c.ProfileID / SliceWidth
ret := false
@ -473,7 +474,7 @@ func (e *Executor) executeSetBit(db string, c *pql.SetBit, opt *ExecOptions) (bo
}
// Forward call to remote node otherwise.
if res, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt); err != nil {
if res, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt); err != nil {
return false, err
} else {
ret = res[0].(bool)
@ -483,7 +484,7 @@ func (e *Executor) executeSetBit(db string, c *pql.SetBit, opt *ExecOptions) (bo
}
// executeSetBitmapAttrs executes a SetBitmapAttrs() call.
func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs, opt *ExecOptions) error {
func (e *Executor) executeSetBitmapAttrs(ctx context.Context, db string, c *pql.SetBitmapAttrs, opt *ExecOptions) error {
// Retrieve frame.
frame, err := e.Index.CreateFrameIfNotExists(db, c.Frame)
if err != nil {
@ -505,7 +506,7 @@ func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs, opt *
resp := make(chan error, len(nodes))
for _, node := range nodes {
go func(node *Node) {
_, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt)
_, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt)
resp <- err
}(node)
}
@ -521,7 +522,7 @@ func (e *Executor) executeSetBitmapAttrs(db string, c *pql.SetBitmapAttrs, opt *
}
// executeBulkSetBitmapAttrs executes a set of SetBitmapAttrs() calls.
func (e *Executor) executeBulkSetBitmapAttrs(db string, calls pql.Calls, opt *ExecOptions) ([]interface{}, error) {
func (e *Executor) executeBulkSetBitmapAttrs(ctx context.Context, db string, calls pql.Calls, opt *ExecOptions) ([]interface{}, error) {
// Collect attributes by frame/id.
m := make(map[string]map[uint64]map[string]interface{})
for _, call := range calls {
@ -569,7 +570,7 @@ func (e *Executor) executeBulkSetBitmapAttrs(db string, calls pql.Calls, opt *Ex
resp := make(chan error, len(nodes))
for _, node := range nodes {
go func(node *Node) {
_, err := e.exec(node, db, &pql.Query{Calls: calls}, nil, opt)
_, err := e.exec(ctx, node, db, &pql.Query{Calls: calls}, nil, opt)
resp <- err
}(node)
}
@ -586,7 +587,7 @@ func (e *Executor) executeBulkSetBitmapAttrs(db string, calls pql.Calls, opt *Ex
}
// executeSetProfileAttrs executes a SetProfileAttrs() call.
func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs, opt *ExecOptions) error {
func (e *Executor) executeSetProfileAttrs(ctx context.Context, db string, c *pql.SetProfileAttrs, opt *ExecOptions) error {
// Retrieve database.
d, err := e.Index.CreateDBIfNotExists(db)
if err != nil {
@ -608,7 +609,7 @@ func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs, opt
resp := make(chan error, len(nodes))
for _, node := range nodes {
go func(node *Node) {
_, err := e.exec(node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt)
_, err := e.exec(ctx, node, db, &pql.Query{Calls: pql.Calls{c}}, nil, opt)
resp <- err
}(node)
}
@ -624,7 +625,7 @@ func (e *Executor) executeSetProfileAttrs(db string, c *pql.SetProfileAttrs, opt
}
// 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, opt *ExecOptions) (results []interface{}, err error) {
func (e *Executor) exec(ctx context.Context, node *Node, db string, q *pql.Query, slices []uint64, opt *ExecOptions) (results []interface{}, err error) {
// Encode request object.
pbreq := &internal.QueryRequest{
DB: proto.String(db),

View file

@ -1,6 +1,7 @@
package pilosa_test
import (
"context"
"reflect"
"strings"
"testing"
@ -22,7 +23,7 @@ func TestExecutor_Execute_Bitmap(t *testing.T) {
}
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Bitmap(id=10, frame=f)`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Bitmap(id=10, frame=f)`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{3, SliceWidth + 1}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -42,7 +43,7 @@ func TestExecutor_Execute_Difference(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "general", 0).MustSetBits(11, 4)
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Difference(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Difference(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{1, 3}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -62,7 +63,7 @@ func TestExecutor_Execute_Intersect(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBits(11, SliceWidth+2)
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Intersect(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Intersect(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{1, SliceWidth + 2}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -81,7 +82,7 @@ func TestExecutor_Execute_Union(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "general", 1).MustSetBits(11, SliceWidth+2)
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Union(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Union(Bitmap(id=10), Bitmap(id=11))`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{0, 2, SliceWidth + 1, SliceWidth + 2}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -97,7 +98,7 @@ func TestExecutor_Execute_Count(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(10, SliceWidth+2)
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Count(Bitmap(id=10, frame=f))`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Count(Bitmap(id=10, frame=f))`), nil, nil); err != nil {
t.Fatal(err)
} else if res[0] != uint64(3) {
t.Fatalf("unexpected n: %d", res[0])
@ -115,7 +116,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) {
t.Fatalf("unexpected bitmap count: %d", n)
}
if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
t.Fatal(err)
} else {
if !res[0].(bool) {
@ -126,7 +127,7 @@ func TestExecutor_Execute_SetBit(t *testing.T) {
if n := f.Bitmap(11).Count(); n != 1 {
t.Fatalf("unexpected bitmap count: %d", n)
}
if res, err := e.Execute("d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`SetBit(id=11, frame=f, profileID=1)`), nil, nil); err != nil {
t.Fatal(err)
} else {
if res[0].(bool) {
@ -143,16 +144,16 @@ func TestExecutor_Execute_SetBitmapAttrs(t *testing.T) {
// Set two fields on f/10.
// Also set fields on other bitmaps and frames to test isolation.
e := NewExecutor(idx.Index, NewCluster(1))
if _, err := e.Execute("d", MustParse(`SetBitmapAttrs(id=10, frame=f, foo="bar")`), nil, nil); err != nil {
if _, err := e.Execute(context.Background(), "d", MustParse(`SetBitmapAttrs(id=10, frame=f, foo="bar")`), nil, nil); err != nil {
t.Fatal(err)
}
if _, err := e.Execute("d", MustParse(`SetBitmapAttrs(id=200, frame=f, YYY=1)`), nil, nil); err != nil {
if _, err := e.Execute(context.Background(), "d", MustParse(`SetBitmapAttrs(id=200, frame=f, YYY=1)`), nil, nil); err != nil {
t.Fatal(err)
}
if _, err := e.Execute("d", MustParse(`SetBitmapAttrs(id=10, frame=XXX, YYY=1)`), nil, nil); err != nil {
if _, err := e.Execute(context.Background(), "d", MustParse(`SetBitmapAttrs(id=10, frame=XXX, YYY=1)`), nil, nil); err != nil {
t.Fatal(err)
}
if _, err := e.Execute("d", MustParse(`SetBitmapAttrs(id=10, frame=f, baz=123, bat=true)`), nil, nil); err != nil {
if _, err := e.Execute(context.Background(), "d", MustParse(`SetBitmapAttrs(id=10, frame=f, baz=123, bat=true)`), nil, nil); err != nil {
t.Fatal(err)
}
@ -182,7 +183,7 @@ func TestExecutor_Execute_TopN(t *testing.T) {
// Execute query.
e := NewExecutor(idx.Index, NewCluster(1))
if result, err := e.Execute("d", MustParse(`TopN(frame=f, n=2)`), nil, nil); err != nil {
if result, err := e.Execute(context.Background(), "d", MustParse(`TopN(frame=f, n=2)`), nil, nil); err != nil {
t.Fatal(err)
} else if !reflect.DeepEqual(result[0], []pilosa.Pair{
{Key: 0, Count: 5},
@ -205,7 +206,7 @@ func TestExecutor_Execute_TopN_fill(t *testing.T) {
// Execute query.
e := NewExecutor(idx.Index, NewCluster(1))
if result, err := e.Execute("d", MustParse(`TopN(frame=f, n=1)`), nil, nil); err != nil {
if result, err := e.Execute(context.Background(), "d", MustParse(`TopN(frame=f, n=1)`), nil, nil); err != nil {
t.Fatal(err)
} else if !reflect.DeepEqual(result, []interface{}{[]pilosa.Pair{
{Key: 0, Count: 4},
@ -236,7 +237,7 @@ func TestExecutor_Execute_TopN_Src(t *testing.T) {
// Execute query.
e := NewExecutor(idx.Index, NewCluster(1))
if result, err := e.Execute("d", MustParse(`TopN(Bitmap(id=100, frame=other), frame=f, n=3)`), nil, nil); err != nil {
if result, err := e.Execute(context.Background(), "d", MustParse(`TopN(Bitmap(id=100, frame=other), frame=f, n=3)`), nil, nil); err != nil {
t.Fatal(err)
} else if !reflect.DeepEqual(result, []interface{}{[]pilosa.Pair{
{Key: 20, Count: 3},
@ -258,7 +259,7 @@ func TestExecutor_Execute_Range(t *testing.T) {
}
e := NewExecutor(idx.Index, NewCluster(1))
if res, err := e.Execute("d", MustParse(`Range(id=1, frame=f.t, start="2000-01-01T00:00", end="2000-01-01T01:00")`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Range(id=1, frame=f.t, start="2000-01-01T00:00", end="2000-01-01T01:00")`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{100}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -275,7 +276,7 @@ func TestExecutor_Execute_Remote_Bitmap(t *testing.T) {
c.Nodes[1].Host = s.Host()
// Mock secondary server's executor to verify arguments and return a bitmap.
s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
s.Handler.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
if db != `d` {
t.Fatalf("unexpected db: %s", db)
} else if query.String() != `Bitmap(id=10, frame=f)` {
@ -300,7 +301,7 @@ func TestExecutor_Execute_Remote_Bitmap(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(10, (1*SliceWidth)+1)
e := NewExecutor(idx.Index, c)
if res, err := e.Execute("d", MustParse(`Bitmap(id=10, frame=f)`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Bitmap(id=10, frame=f)`), nil, nil); err != nil {
t.Fatal(err)
} else if bits := res[0].(*pilosa.Bitmap).Bits(); !reflect.DeepEqual(bits, []uint64{1, 2, (1 * SliceWidth) + 1, 2*SliceWidth + 4}) {
t.Fatalf("unexpected bits: %+v", bits)
@ -317,7 +318,7 @@ func TestExecutor_Execute_Remote_Count(t *testing.T) {
c.Nodes[1].Host = s.Host()
// Mock secondary server's executor to return a count.
s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
s.Handler.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return []interface{}{uint64(10)}, nil
}
@ -328,7 +329,7 @@ func TestExecutor_Execute_Remote_Count(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "f", 1).MustSetBits(10, (1*SliceWidth)+2)
e := NewExecutor(idx.Index, c)
if res, err := e.Execute("d", MustParse(`Count(Bitmap(id=10, frame=f))`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`Count(Bitmap(id=10, frame=f))`), nil, nil); err != nil {
t.Fatal(err)
} else if res[0] != uint64(12) {
t.Fatalf("unexpected n: %d", res[0])
@ -347,7 +348,7 @@ func TestExecutor_Execute_Remote_SetBit(t *testing.T) {
// Mock secondary server's executor to verify arguments.
var remoteCalled bool
s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
s.Handler.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
if db != `d` {
t.Fatalf("unexpected db: %s", db)
} else if query.String() != `SetBit(id=10, frame=f, profileID=2)` {
@ -362,7 +363,7 @@ func TestExecutor_Execute_Remote_SetBit(t *testing.T) {
defer idx.Close()
e := NewExecutor(idx.Index, c)
if _, err := e.Execute("d", MustParse(`SetBit(id=10, frame=f, profileID=2)`), nil, nil); err != nil {
if _, err := e.Execute(context.Background(), "d", MustParse(`SetBit(id=10, frame=f, profileID=2)`), nil, nil); err != nil {
t.Fatal(err)
}
@ -386,7 +387,7 @@ func TestExecutor_Execute_Remote_TopN(t *testing.T) {
// Mock secondary server's executor to verify arguments and return a bitmap.
var remoteExecN int
s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
s.Handler.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
if db != `d` {
t.Fatalf("unexpected db: %s", db)
} else if !reflect.DeepEqual(slices, []uint64{0, 2, 4, 6}) {
@ -424,7 +425,7 @@ func TestExecutor_Execute_Remote_TopN(t *testing.T) {
idx.MustCreateFragmentIfNotExists("d", "f", 3).MustSetBits(30, (3*SliceWidth)+2)
e := NewExecutor(idx.Index, c)
if res, err := e.Execute("d", MustParse(`TopN(frame=f, n=3)`), nil, nil); err != nil {
if res, err := e.Execute(context.Background(), "d", MustParse(`TopN(frame=f, n=3)`), nil, nil); err != nil {
t.Fatal(err)
} else if !reflect.DeepEqual(res, []interface{}{[]pilosa.Pair{
{Key: 0, Count: 5},

View file

@ -4,6 +4,7 @@ import (
"archive/tar"
"bufio"
"bytes"
"context"
"crypto/sha1"
"encoding/binary"
"errors"
@ -1278,7 +1279,7 @@ func (s *FragmentSyncer) SyncFragment() error {
if err != nil {
return err
}
blocks, err := client.FragmentBlocks(s.Fragment.DB(), s.Fragment.Frame(), s.Fragment.Slice())
blocks, err := client.FragmentBlocks(context.Background(), s.Fragment.DB(), s.Fragment.Frame(), s.Fragment.Slice())
if err != nil && err != ErrFragmentNotFound {
return err
}
@ -1359,7 +1360,7 @@ func (s *FragmentSyncer) syncBlock(id int) error {
}
clients = append(clients, client)
bitmapIDs, profileIDs, err := client.BlockData(f.DB(), f.Frame(), f.Slice(), id)
bitmapIDs, profileIDs, err := client.BlockData(context.Background(), f.DB(), f.Frame(), f.Slice(), id)
if err != nil {
return err
}
@ -1405,7 +1406,7 @@ func (s *FragmentSyncer) syncBlock(id int) error {
}
// Execute query.
_, err := clients[i].ExecuteQuery(f.DB(), buf.String(), false)
_, err := clients[i].ExecuteQuery(context.Background(), f.DB(), buf.String(), false)
if err != nil {
return err
}

View file

@ -1,6 +1,7 @@
package pilosa
import (
"context"
"encoding/csv"
"encoding/json"
"errors"
@ -31,7 +32,7 @@ type Handler struct {
// The execution engine for running queries.
Executor interface {
Execute(db string, query *pql.Query, slices []uint64, opt *ExecOptions) ([]interface{}, error)
Execute(context context.Context, db string, query *pql.Query, slices []uint64, opt *ExecOptions) ([]interface{}, error)
}
// The version to report on the /version endpoint.
@ -215,7 +216,7 @@ func (h *Handler) handlePostQuery(w http.ResponseWriter, r *http.Request) {
}
// Execute the query.
results, err := h.Executor.Execute(req.DB, q, req.Slices, opt)
results, err := h.Executor.Execute(r.Context(), req.DB, q, req.Slices, opt)
resp := &QueryResponse{Results: results, Err: err}
// Fill profile attributes if requested.
@ -570,7 +571,7 @@ func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) {
}
// Find the correct fragment.
h.logger().Println("Go Import:", db, frame, slice)
h.logger().Println("importing:", db, frame, slice)
f, err := h.Index.CreateFragmentIfNotExists(db, frame, slice)
if err != nil {
h.logger().Printf("fragment error: db=%s, frame=%s, slice=%d, err=%s", db, frame, slice, err)
@ -815,7 +816,7 @@ func (h *Handler) handlePostFrameRestore(w http.ResponseWriter, r *http.Request)
}
// Determine the maximum number of slices.
sliceN, err := client.SliceN()
sliceN, err := client.SliceN(r.Context())
if err != nil {
http.Error(w, "cannot determine remote slice count: "+err.Error(), http.StatusInternalServerError)
return
@ -836,18 +837,18 @@ func (h *Handler) handlePostFrameRestore(w http.ResponseWriter, r *http.Request)
}
// Stream backup from remote node.
r, err := client.BackupSlice(db, frame, slice)
rd, err := client.BackupSlice(r.Context(), db, frame, slice)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
} else if r == nil {
} else if rd == nil {
continue // slice doesn't exist
}
// Restore to local frame and always close reader.
if err := func() error {
defer r.Close()
if _, err := f.ReadFrom(r); err != nil {
defer rd.Close()
if _, err := f.ReadFrom(rd); err != nil {
return err
}
return nil

View file

@ -2,6 +2,7 @@ package pilosa_test
import (
"bytes"
"context"
"encoding/json"
"errors"
"io"
@ -56,7 +57,7 @@ func TestHandler_Schema(t *testing.T) {
// Ensure the handler can accept URL arguments.
func TestHandler_Query_Args_URL(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
if db != "db0" {
t.Fatalf("unexpected db: %s", db)
} else if query.String() != `Count(Bitmap(id=100))` {
@ -79,7 +80,7 @@ func TestHandler_Query_Args_URL(t *testing.T) {
// Ensure the handler can accept arguments via protobufs.
func TestHandler_Query_Args_Protobuf(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
if db != "db0" {
t.Fatalf("unexpected db: %s", db)
} else if query.String() != `Count(Bitmap(id=100))` {
@ -125,7 +126,7 @@ func TestHandler_Query_Args_Err(t *testing.T) {
// Ensure the handler can execute a query with a uint64 response as JSON.
func TestHandler_Query_Uint64_JSON(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return []interface{}{uint64(100)}, nil
}
@ -141,7 +142,7 @@ func TestHandler_Query_Uint64_JSON(t *testing.T) {
// Ensure the handler can execute a query with a uint64 response as protobufs.
func TestHandler_Query_Uint64_Protobuf(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return []interface{}{uint64(100)}, nil
}
@ -164,7 +165,7 @@ func TestHandler_Query_Uint64_Protobuf(t *testing.T) {
// Ensure the handler can execute a query that returns a bitmap as JSON.
func TestHandler_Query_Bitmap_JSON(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
bm := pilosa.NewBitmap(1, 3, 66, pilosa.SliceWidth+1)
bm.Attrs = map[string]interface{}{"a": "b", "c": 1, "d": true}
return []interface{}{bm}, nil
@ -196,7 +197,7 @@ func TestHandler_Query_Bitmap_Profiles_JSON(t *testing.T) {
h := NewHandler()
h.Index = idx.Index
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
bm := pilosa.NewBitmap(1, 3, 66, pilosa.SliceWidth+1)
bm.Attrs = map[string]interface{}{"a": "b", "c": 1, "d": true}
return []interface{}{bm}, nil
@ -214,7 +215,7 @@ func TestHandler_Query_Bitmap_Profiles_JSON(t *testing.T) {
// Ensure the handler can execute a query that returns a bitmap as protobuf.
func TestHandler_Query_Bitmap_Protobuf(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
bm := pilosa.NewBitmap(1, pilosa.SliceWidth+1)
bm.Attrs = map[string]interface{}{"a": "b", "c": int64(1), "d": true}
return []interface{}{bm}, nil
@ -259,7 +260,7 @@ func TestHandler_Query_Bitmap_Profiles_Protobuf(t *testing.T) {
h := NewHandler()
h.Index = idx.Index
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
bm := pilosa.NewBitmap(1, pilosa.SliceWidth+1)
bm.Attrs = map[string]interface{}{"a": "b", "c": int64(1), "d": true}
return []interface{}{bm}, nil
@ -314,7 +315,7 @@ func TestHandler_Query_Bitmap_Profiles_Protobuf(t *testing.T) {
// Ensure the handler can execute a query that returns pairs as JSON.
func TestHandler_Query_Pairs_JSON(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return []interface{}{[]pilosa.Pair{
{Key: 1, Count: 2},
{Key: 3, Count: 4},
@ -333,7 +334,7 @@ func TestHandler_Query_Pairs_JSON(t *testing.T) {
// Ensure the handler can execute a query that returns pairs as protobuf.
func TestHandler_Query_Pairs_Protobuf(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return []interface{}{[]pilosa.Pair{
{Key: 1, Count: 2},
{Key: 3, Count: 4},
@ -359,7 +360,7 @@ func TestHandler_Query_Pairs_Protobuf(t *testing.T) {
// Ensure the handler can return an error as JSON.
func TestHandler_Query_Err_JSON(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return nil, errors.New("marker")
}
@ -375,7 +376,7 @@ func TestHandler_Query_Err_JSON(t *testing.T) {
// Ensure the handler can return an error as protobuf.
func TestHandler_Query_Err_Protobuf(t *testing.T) {
h := NewHandler()
h.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
h.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return nil, errors.New("marker")
}
@ -654,13 +655,13 @@ func NewHandler() *Handler {
// HandlerExecutor is a mock implementing pilosa.Handler.Executor.
type HandlerExecutor struct {
cluster *pilosa.Cluster
ExecuteFn func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error)
ExecuteFn func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error)
}
func (c *HandlerExecutor) Cluster() *pilosa.Cluster { return c.cluster }
func (c *HandlerExecutor) Execute(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return c.ExecuteFn(db, query, slices, opt)
func (c *HandlerExecutor) Execute(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
return c.ExecuteFn(ctx, db, query, slices, opt)
}
// Server represents a test wrapper for httptest.Server.

View file

@ -1,6 +1,7 @@
package pilosa
import (
"context"
"errors"
"fmt"
"io"
@ -401,7 +402,7 @@ func (s *IndexSyncer) syncDatabase(db string) error {
// Retrieve attributes from differing blocks.
// Skip update and recomputation if no attributes have changed.
m, err := client.ProfileAttrDiff(db, blks)
m, err := client.ProfileAttrDiff(context.Background(), db, blks)
if err != nil {
return err
} else if len(m) == 0 {
@ -446,7 +447,7 @@ func (s *IndexSyncer) syncFrame(db, name string) error {
// Retrieve attributes from differing blocks.
// Skip update and recomputation if no attributes have changed.
m, err := client.BitmapAttrDiff(db, name, blks)
m, err := client.BitmapAttrDiff(context.Background(), db, name, blks)
if err != nil {
return err
} else if len(m) == 0 {

View file

@ -2,6 +2,7 @@ package pilosa_test
import (
"bytes"
"context"
"io/ioutil"
"os"
"reflect"
@ -58,12 +59,12 @@ func TestIndexSyncer_SyncIndex(t *testing.T) {
s := NewServer()
defer s.Close()
s.Handler.Index = idx1.Index
s.Handler.Executor.ExecuteFn = func(db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
s.Handler.Executor.ExecuteFn = func(ctx context.Context, db string, query *pql.Query, slices []uint64, opt *pilosa.ExecOptions) ([]interface{}, error) {
e := pilosa.NewExecutor()
e.Index = idx1.Index
e.Host = cluster.Nodes[1].Host
e.Cluster = cluster
return e.Execute(db, query, slices, opt)
return e.Execute(ctx, db, query, slices, opt)
}
// Mock 2-node, fully replicated cluster.