Replace store.WithContext with normal ctx param

api/api_gomux.go

@@ -576,7 +576,7 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) {
 
 	// Query data store (instrumented)
 	rc.inst.Start("db")
-	entities, err := api.es.WithContext(ctx).ReadEntities(rc.entityType, q, f)
+	entities, err := api.es.ReadEntities(ctx, rc.entityType, q, f)
 	rc.inst.Stop("db")
 	if err != nil {
 		api.readError(rc, w, err)
@@ -650,7 +650,7 @@ func (api *API) postEntitiesHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Write new entities to data store
-	ids, err = api.es.WithContext(ctx).CreateEntities(rc.wo, entities)
+	ids, err = api.es.CreateEntities(ctx, rc.wo, entities)
 	rc.gm.Inc(metrics.Created, int64(len(ids)))
 
 reply:
@@ -717,7 +717,7 @@ func (api *API) putEntitiesHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Patch all entities matching query
-	entities, err = api.es.WithContext(ctx).UpdateEntities(rc.wo, q, patch)
+	entities, err = api.es.UpdateEntities(ctx, rc.wo, q, patch)
 	rc.gm.Val(metrics.UpdateBulk, int64(len(entities)))
 	rc.gm.Inc(metrics.Updated, int64(len(entities)))
 
@@ -764,7 +764,7 @@ func (api *API) deleteEntitiesHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Delete entities, returns the deleted entities
-	entities, err = api.es.WithContext(ctx).DeleteEntities(rc.wo, q)
+	entities, err = api.es.DeleteEntities(ctx, rc.wo, q)
 	rc.gm.Val(metrics.DeleteBulk, int64(len(entities)))
 	rc.gm.Inc(metrics.Deleted, int64(len(entities)))
 
@@ -805,7 +805,7 @@ func (api *API) getEntityHandler(w http.ResponseWriter, r *http.Request) {
 
 	// Read the entity by ID
 	q, _ := query.Translate("_id=" + rc.entityId)
-	entities, err := api.es.WithContext(ctx).ReadEntities(rc.entityType, q, f)
+	entities, err := api.es.ReadEntities(ctx, rc.entityType, q, f)
 	if err != nil {
 		api.readError(rc, w, err)
 		return
@@ -838,7 +838,7 @@ func (api *API) getLabelsHandler(w http.ResponseWriter, r *http.Request) {
 	rc.gm.Inc(metrics.ReadLabels, 1) // specific read type
 
 	q, _ := query.Translate("_id=" + rc.entityId)
-	entities, err := api.es.WithContext(ctx).ReadEntities(rc.entityType, q, etre.QueryFilter{})
+	entities, err := api.es.ReadEntities(ctx, rc.entityType, q, etre.QueryFilter{})
 	if err != nil {
 		api.readError(rc, w, err)
 		return
@@ -893,7 +893,7 @@ func (api *API) postEntityHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Create new entity
-	ids, err = api.es.WithContext(ctx).CreateEntities(rc.wo, []etre.Entity{newEntity})
+	ids, err = api.es.CreateEntities(ctx, rc.wo, []etre.Entity{newEntity})
 	if err == nil {
 		rc.gm.Inc(metrics.Created, 1)
 	}
@@ -947,7 +947,7 @@ func (api *API) putEntityHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Patch one entity by ID
-	entities, err = api.es.WithContext(ctx).UpdateEntities(rc.wo, q, patch)
+	entities, err = api.es.UpdateEntities(ctx, rc.wo, q, patch)
 	if err != nil {
 		goto reply
 	} else if len(entities) == 0 {
@@ -991,7 +991,7 @@ func (api *API) deleteEntityHandler(w http.ResponseWriter, r *http.Request) {
 	q, _ := query.Translate("_id=" + rc.entityId)
 
 	// Delete one entity by ID
-	entities, err = api.es.WithContext(ctx).DeleteEntities(rc.wo, q)
+	entities, err = api.es.DeleteEntities(ctx, rc.wo, q)
 	if err != nil {
 		goto reply
 	} else if len(entities) == 0 {
@@ -1041,7 +1041,7 @@ func (api *API) deleteLabelHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Delete label from entity
-	diff, err = api.es.WithContext(ctx).DeleteLabel(rc.wo, label)
+	diff, err = api.es.DeleteLabel(ctx, rc.wo, label)
 	if err != nil {
 		if err == etre.ErrEntityNotFound {
 			err = nil // delete is idempotent

api/api_test.go

@@ -198,11 +198,8 @@ func TestClientQueryTimeout(t *testing.T) {
 	// and plumbed all the way down to the entity.Store context
 	var gotCtx context.Context
 	store := mock.EntityStore{}
-	store.WithContextFunc = func(ctx context.Context) entity.Store {
+	store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 		gotCtx = ctx
-		return store
-	}
-	store.ReadEntitiesFunc = func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 		return testEntitiesWithObjectIDs[0:1], nil
 	}
 	server := setup(t, defaultConfig, store)
@@ -244,24 +241,25 @@ func TestContextPropagation(t *testing.T) {
 	// Make sure context values from the request are propagated all the way down to the entity.Store context
 	var gotCtx context.Context
 	store := mock.EntityStore{}
-	store.WithContextFunc = func(ctx context.Context) entity.Store {
-		gotCtx = ctx
-		return store
-	}
 	// We're going to test all operations, so we need to set all of these funcs
-	store.ReadEntitiesFunc = func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+	store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		gotCtx = ctx
 		return testEntitiesWithObjectIDs[0:1], nil
 	}
-	store.CreateEntitiesFunc = func(op entity.WriteOp, entities []etre.Entity) ([]string, error) {
+	store.CreateEntitiesFunc = func(ctx context.Context, op entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		gotCtx = ctx
 		return []string{testEntityIds[0]}, nil
 	}
-	store.UpdateEntitiesFunc = func(op entity.WriteOp, q query.Query, e etre.Entity) ([]etre.Entity, error) {
+	store.UpdateEntitiesFunc = func(ctx context.Context, op entity.WriteOp, q query.Query, e etre.Entity) ([]etre.Entity, error) {
+		gotCtx = ctx
 		return testEntitiesWithObjectIDs[0:1], nil
 	}
-	store.DeleteEntitiesFunc = func(op entity.WriteOp, q query.Query) ([]etre.Entity, error) {
+	store.DeleteEntitiesFunc = func(ctx context.Context, op entity.WriteOp, q query.Query) ([]etre.Entity, error) {
+		gotCtx = ctx
 		return testEntitiesWithObjectIDs[0:1], nil
 	}
-	store.DeleteLabelFunc = func(op entity.WriteOp, label string) (etre.Entity, error) {
+	store.DeleteLabelFunc = func(ctx context.Context, op entity.WriteOp, label string) (etre.Entity, error) {
+		gotCtx = ctx
 		return testEntitiesWithObjectIDs[0], nil
 	}
 

api/bulk_write_test.go

@@ -3,6 +3,7 @@
 package api_test
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"net/url"
@@ -32,7 +33,7 @@ func TestPostEntitiesOK(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotEntities []etre.Entity
 	store := mock.EntityStore{
-		CreateEntitiesFunc: func(wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		CreateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
 			gotWO = wo
 			gotEntities = entities
 			return []string{"id1", "id2"}, nil
@@ -95,7 +96,7 @@ func TestPostEntitiesErrors(t *testing.T) {
 	// Most importantly: CreateEntities() should _not_ be called.
 	created := false
 	store := mock.EntityStore{
-		CreateEntitiesFunc: func(wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		CreateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
 			created = true
 			return []string{"id1", "id2"}, nil
 		},
@@ -246,7 +247,7 @@ func TestPutEntitiesOK(t *testing.T) {
 	var gotQuery query.Query
 	var gotPatch etre.Entity
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			gotWO = wo
 			gotQuery = q
 			gotPatch = patch
@@ -330,7 +331,7 @@ func TestPutEntitiesErrors(t *testing.T) {
 	// metrics when any input is invalid. The UpdateEntities() should not be called.
 	updated := false
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			updated = true
 			return []etre.Entity{}, nil
 		},
@@ -496,7 +497,7 @@ func TestDeleteEntitiesOK(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotQuery query.Query
 	store := mock.EntityStore{
-		DeleteEntitiesFunc: func(wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
+		DeleteEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
 			gotWO = wo
 			gotQuery = q
 			return []etre.Entity{
@@ -575,7 +576,7 @@ func TestDeleteEntitiesErrors(t *testing.T) {
 	// metrics when any input is invalid. The DeleteEntities() should not be called.
 	deleted := false
 	store := mock.EntityStore{
-		DeleteEntitiesFunc: func(wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
+		DeleteEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
 			deleted = true
 			return []etre.Entity{}, nil
 		},

api/query_test.go

@@ -28,7 +28,7 @@ func TestQueryBasic(t *testing.T) {
 	var gotQuery query.Query
 	var gotFilter etre.QueryFilter
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			gotQuery = q
 			gotFilter = f
 			return testEntitiesWithObjectIDs, nil
@@ -171,7 +171,7 @@ func TestQueryNoMatches(t *testing.T) {
 	// is still 200 OK in this case because there's no error.
 	var gotQuery query.Query
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			gotQuery = q
 			return []etre.Entity{}, nil // no matching queries
 		},
@@ -220,7 +220,7 @@ func TestQueryErrorsDatabaseError(t *testing.T) {
 	// Test that GET /entities/:type?query=Q handles a database error correctly.
 	// Db errors (and only db errors return HTTP 503 "Service Unavailable".
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			return nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}
 		},
 	}
@@ -269,7 +269,7 @@ func TestQueryErrorsNoEntityType(t *testing.T) {
 	// You can run "../test/coverage -test.run TestQueryErrorsNoEntityType" and
 	// see that the handler is never called.
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			return nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}
 		},
 	}
@@ -310,7 +310,7 @@ func TestQueryErrorsTimeout(t *testing.T) {
 	// Test that GET /entities/:type?query=Q handles a database timeout correctly.
 	// Db errors (and only db errors return HTTP 503 "Service Unavailable".
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 			defer cancel()
 			<-ctx.Done()
@@ -357,7 +357,7 @@ func TestQueryErrorsTimeout(t *testing.T) {
 func TestResponseCompression(t *testing.T) {
 	// Stand up the server
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			return testEntitiesWithObjectIDs, nil
 		},
 	}

api/single_entity_read_test.go

@@ -3,6 +3,7 @@
 package api_test
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -31,7 +32,7 @@ func TestGetEntityBasic(t *testing.T) {
 	var gotQuery query.Query
 	var gotFilter etre.QueryFilter
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			gotQuery = q
 			gotFilter = f
 			return testEntitiesWithObjectIDs[0:1], nil
@@ -83,7 +84,7 @@ func TestGetEntityReturnLabels(t *testing.T) {
 	// that the URL param "labels=" is processed and passed along to the entity.Store.
 	var gotFilter etre.QueryFilter
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			gotFilter = f
 			return testEntitiesWithObjectIDs[0:1], nil
 		},
@@ -129,7 +130,7 @@ func TestGetEntityNotFound(t *testing.T) {
 	// We simulate this by making ReadEntities() below return an empty list which
 	// the real entity.Store() does when no entity exists with the given _id.
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			return []etre.Entity{}, nil
 		},
 	}
@@ -163,7 +164,7 @@ func TestGetEntityErrors(t *testing.T) {
 	read := false
 	var dbError error
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			read = true
 			return nil, dbError
 		},
@@ -257,7 +258,7 @@ func TestGetEntityLabels(t *testing.T) {
 	var gotQuery query.Query
 	var gotFilter etre.QueryFilter
 	store := mock.EntityStore{
-		ReadEntitiesFunc: func(entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
+		ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) {
 			gotQuery = q
 			gotFilter = f
 			return testEntitiesWithObjectIDs[0:1], nil

api/single_entity_write_test.go

@@ -3,6 +3,7 @@
 package api_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -33,7 +34,7 @@ func TestPostEntityOK(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotEntities []etre.Entity
 	store := mock.EntityStore{
-		CreateEntitiesFunc: func(wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		CreateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
 			gotWO = wo
 			gotEntities = entities
 			return []string{"id1"}, nil
@@ -96,7 +97,7 @@ func TestPostEntityDuplicate(t *testing.T) {
 	// Test that POST /entities/:type returns HTTP 403 Conflict on duplicate
 	// which we simulate by returning what entity.Store would:
 	store := mock.EntityStore{
-		CreateEntitiesFunc: func(wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		CreateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
 			// Real CreateEntities() always returns []string, not nil, because
 			// it supports partial writes
 			return []string{}, entity.DbError{Err: fmt.Errorf("dupe"), Type: "duplicate-entity"}
@@ -144,7 +145,7 @@ func TestPostEntityErrors(t *testing.T) {
 	// Test that POST /entities/:type returns an error for any issue
 	created := false
 	store := mock.EntityStore{
-		CreateEntitiesFunc: func(wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
+		CreateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, entities []etre.Entity) ([]string, error) {
 			created = true
 			return []string{"id1"}, nil
 		},
@@ -254,7 +255,7 @@ func TestPutEntityOK(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotQuery query.Query
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			gotWO = wo
 			gotQuery = q
 			diff := []etre.Entity{
@@ -329,7 +330,7 @@ func TestPutEntityDuplicate(t *testing.T) {
 	// Test that PUT /entities/:type/:id returns HTTP 403 Conflict on duplicate
 	// which we simulate by returning what entity.Store would:
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			return nil, entity.DbError{
 				Type:     "duplicate-entity", // the key to making this happen
 				EntityId: testEntityIds[0],
@@ -380,7 +381,7 @@ func TestPutEntityNotFound(t *testing.T) {
 	// Test that PUT /entities/:type/:id returns HTTP 404 when there's no entity
 	// with the given id. In this case, the entity.Store returns an empty diff:
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			return []etre.Entity{}, nil // no entities matched
 		},
 	}
@@ -419,7 +420,7 @@ func TestPutEntityErrors(t *testing.T) {
 	// Test that PUT /entities/:type/:id returns errors unless all inputs are correct
 	updated := false
 	store := mock.EntityStore{
-		UpdateEntitiesFunc: func(wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
+		UpdateEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query, patch etre.Entity) ([]etre.Entity, error) {
 			updated = true
 			return []etre.Entity{{"_id": testEntityId0, "_type": entityType, "_rev": int64(0)}}, nil
 		},
@@ -592,7 +593,7 @@ func TestDeleteEntityOK(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotQuery query.Query
 	store := mock.EntityStore{
-		DeleteEntitiesFunc: func(wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
+		DeleteEntitiesFunc: func(ctx context.Context, wo entity.WriteOp, q query.Query) ([]etre.Entity, error) {
 			gotWO = wo
 			gotQuery = q
 			diff := []etre.Entity{
@@ -671,7 +672,7 @@ func TestDeleteLabel(t *testing.T) {
 	var gotWO entity.WriteOp
 	var gotLabel string
 	store := mock.EntityStore{
-		DeleteLabelFunc: func(wo entity.WriteOp, label string) (etre.Entity, error) {
+		DeleteLabelFunc: func(ctx context.Context, wo entity.WriteOp, label string) (etre.Entity, error) {
 			gotWO = wo
 			gotLabel = label
 			return etre.Entity{"_id": testEntityId0, "_type": entityType, "_rev": int64(0), "foo": "oldVal"}, nil