diff --git a/api/api_gomux.go b/api/api_gomux.go index bc9f134..7d7ffe3 100644 --- a/api/api_gomux.go +++ b/api/api_gomux.go @@ -8,6 +8,7 @@ import ( "context" "encoding/json" "fmt" + "io" "log" "math/rand" "net/http" @@ -576,26 +577,96 @@ func (api *API) getEntitiesHandler(w http.ResponseWriter, r *http.Request) { // Query data store (instrumented) rc.inst.Start("db") - entities, err := api.es.ReadEntities(ctx, rc.entityType, q, f) + entities := api.es.StreamEntities(ctx, rc.entityType, q, f) rc.inst.Stop("db") - if err != nil { + + rc.inst.Start("encode-response") + + // FinalWriter is where we will write the data. If we're using gzip compression, this will be the gzip writer. If not, it will just be the original response writer. + var finalWriter io.Writer + finalWriter = w + + // gzip writer, only initialized if client accepts gzip encoding. + var gzw *gzip.Writer + + // encoder is the JSON encoder writing to finalWriter. We initialize it after we know whether we're using gzip or not. + var encoder *json.Encoder + + // Number of non-error records sent to the client + count := 0 + for e := range entities { + + // Handle errors returned by the database + if e.Err != nil { + api.readError(rc, w, e.Err) + return + } + // Handle context timeouts + if err := ctx.Err(); err != nil { + api.readError(rc, w, err) + return + } + + // Initialize gzip writer and JSON encoder on the first record, after we know there is data to return to the client. + // We also check the client's Accept-Encoding header to see if gzip is supported. + if count == 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { + w.Header().Set("Content-Encoding", "gzip") + gzw = gzip.NewWriter(w) + defer gzw.Close() + finalWriter = gzw + } + + // Start the JSON array if this is the first record, or put a comma separator if not the first record + // We don't do this before the for loop because we don't want to write the opening bracket if there is a database + // error in the first record returned. + if count == 0 { + finalWriter.Write([]byte("[")) + encoder = json.NewEncoder(finalWriter) + } else { + finalWriter.Write([]byte(",")) + } + + // Write the record and handle the error. + // encoder.Encode will add a newline, which is fine (nice for human readable output) + err = encoder.Encode(e.Entity) + if err != nil { + // api.readError will mangle the response, but if the encoder failed then the response is already mangled and there's not much else we can do. + api.readError(rc, w, err) + log.Println("ERROR: Read error while encoding response: ", err) + return + } + + // Count the record + count++ + + // If we're compressing, flush every 100 records to ensure data is sent to the client in a timely manner and not buffered in memory. + // Otherwise, the gzip flusher may buffer all the data instead of chunking. + if gzw != nil && count%100 == 0 { + gzw.Flush() + } + } + + // One final check of the context timeout. If the store closed the channel because the context was canceled, + // we want to return an error instead of partial results. + if err := ctx.Err(); err != nil { api.readError(rc, w, err) return } - rc.gm.Val(metrics.ReadMatch, int64(len(entities))) - // Success: return matching entities (possibly empty list) - rc.inst.Start("encode-response") - if len(entities) > 0 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { - // use gzip compression if we have data to send back and the client accepts gzip - w.Header().Set("Content-Encoding", "gzip") - gzw := gzip.NewWriter(w) - defer gzw.Close() - json.NewEncoder(gzw).Encode(entities) + // Clean up the JSON array + if count == 0 { + // Never saw any data. Write an empty array. + finalWriter.Write([]byte("[]")) } else { - // no compression - json.NewEncoder(w).Encode(entities) + // We wrote some data, now we need to close the array + finalWriter.Write([]byte("]")) // end of JSON array + } + + if gzw != nil { + gzw.Flush() // flush any remaining compressed data to the client } + + rc.gm.Val(metrics.ReadMatch, int64(count)) rc.inst.Stop("encode-response") } @@ -804,18 +875,17 @@ 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.ReadEntities(ctx, rc.entityType, q, f) + entity, err := api.es.ReadEntity(ctx, rc.entityType, rc.entityId, f) if err != nil { api.readError(rc, w, err) return } - if len(entities) == 0 { + if entity == nil { w.WriteHeader(http.StatusNotFound) return } - json.NewEncoder(w).Encode(entities[0]) + json.NewEncoder(w).Encode(entity) } // getLabelsHandler godoc @@ -837,18 +907,17 @@ 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.ReadEntities(ctx, rc.entityType, q, etre.QueryFilter{}) + entity, err := api.es.ReadEntity(ctx, rc.entityType, rc.entityId, etre.QueryFilter{}) if err != nil { api.readError(rc, w, err) return } - if len(entities) == 0 { + if entity == nil { w.WriteHeader(http.StatusNotFound) return } - json.NewEncoder(w).Encode(entities[0].Labels()) + json.NewEncoder(w).Encode(entity.Labels()) } // -------------------------------------------------------------------------- diff --git a/api/api_test.go b/api/api_test.go index 0b1dbd1..0cef138 100644 --- a/api/api_test.go +++ b/api/api_test.go @@ -198,9 +198,9 @@ func TestClientQueryTimeout(t *testing.T) { // and plumbed all the way down to the entity.Store context var gotCtx context.Context store := mock.EntityStore{} - store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + store.ReadEntityFunc = func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotCtx = ctx - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil } server := setup(t, defaultConfig, store) defer server.ts.Close() @@ -242,9 +242,13 @@ func TestContextPropagation(t *testing.T) { var gotCtx context.Context store := mock.EntityStore{} // We're going to test all operations, so we need to set all of these funcs - store.ReadEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + store.ReadEntityFunc = func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotCtx = ctx - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil + } + store.StreamEntitiesFunc = func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + gotCtx = ctx + return mock.DoStreamEntities(testEntitiesWithObjectIDs[0:1], nil) } store.CreateEntitiesFunc = func(ctx context.Context, op entity.WriteOp, entities []etre.Entity) ([]string, error) { gotCtx = ctx diff --git a/api/query_test.go b/api/query_test.go index c52f50e..fa9a87b 100644 --- a/api/query_test.go +++ b/api/query_test.go @@ -28,10 +28,10 @@ func TestQueryBasic(t *testing.T) { var gotQuery query.Query var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { gotQuery = q gotFilter = f - return testEntitiesWithObjectIDs, nil + return mock.DoStreamEntities(testEntitiesWithObjectIDs, nil) }, } server := setup(t, defaultConfig, store) @@ -171,9 +171,9 @@ 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(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { gotQuery = q - return []etre.Entity{}, nil // no matching queries + return mock.DoStreamEntities(nil, nil) }, } server := setup(t, defaultConfig, store) @@ -220,8 +220,8 @@ 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(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"} + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -269,8 +269,8 @@ 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(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"} + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(nil, entity.DbError{Err: fmt.Errorf("fake error"), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -310,11 +310,11 @@ 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(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() <-ctx.Done() - return nil, entity.DbError{Err: ctx.Err(), Type: "db-read"} + return mock.DoStreamEntities(nil, entity.DbError{Err: ctx.Err(), Type: "db-read"}) }, } server := setup(t, defaultConfig, store) @@ -357,8 +357,8 @@ func TestQueryErrorsTimeout(t *testing.T) { func TestResponseCompression(t *testing.T) { // Stand up the server store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return testEntitiesWithObjectIDs, nil + StreamEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + return mock.DoStreamEntities(testEntitiesWithObjectIDs, nil) }, } server := setup(t, defaultConfig, store) diff --git a/api/single_entity_read_test.go b/api/single_entity_read_test.go index 987c192..a233d9e 100644 --- a/api/single_entity_read_test.go +++ b/api/single_entity_read_test.go @@ -17,7 +17,6 @@ import ( "github.com/square/etre/auth" "github.com/square/etre/entity" "github.com/square/etre/metrics" - "github.com/square/etre/query" "github.com/square/etre/test" "github.com/square/etre/test/mock" ) @@ -27,15 +26,14 @@ import ( // ////////////////////////////////////////////////////////////////////////// func TestGetEntityBasic(t *testing.T) { - // Test the most basic GET /entity/:type/:id gets the entity. This is - // really a wrapper to call ReadEntitiesFunc() with _id=:id. - var gotQuery query.Query + // Test the most basic GET /entity/:type/:id gets the entity. + var gotEntityId string var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - gotQuery = q + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + gotEntityId = entityId gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -49,8 +47,7 @@ func TestGetEntityBasic(t *testing.T) { assert.Equal(t, http.StatusOK, statusCode, "response status = %d, expected %d", statusCode, http.StatusOK) // GET /entity/:type/:id = "_id=:id" - expectQuery, _ := query.Translate("_id=" + testEntityIds[0]) - assert.Equal(t, expectQuery, gotQuery) + assert.Equal(t, testEntityIds[0], gotEntityId) // No filter options provided in URL expectFilter := etre.QueryFilter{} @@ -84,9 +81,9 @@ 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(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -130,8 +127,8 @@ 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(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - return []etre.Entity{}, nil + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + return nil, nil }, } server := setup(t, defaultConfig, store) @@ -164,7 +161,7 @@ func TestGetEntityErrors(t *testing.T) { read := false var dbError error store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { read = true return nil, dbError }, @@ -255,13 +252,13 @@ func TestGetEntityErrors(t *testing.T) { func TestGetEntityLabels(t *testing.T) { // Test that GET /entity/:type/:id/labels works - var gotQuery query.Query + var gotEntityId string var gotFilter etre.QueryFilter store := mock.EntityStore{ - ReadEntitiesFunc: func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - gotQuery = q + ReadEntityFunc: func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + gotEntityId = entityId gotFilter = f - return testEntitiesWithObjectIDs[0:1], nil + return testEntitiesWithObjectIDs[0], nil }, } server := setup(t, defaultConfig, store) @@ -275,8 +272,7 @@ func TestGetEntityLabels(t *testing.T) { assert.Equal(t, http.StatusOK, statusCode) // GET /entity/:type/:id = "_id=:id" - expectQuery, _ := query.Translate("_id=" + testEntityIds[0]) - assert.Equal(t, expectQuery, gotQuery) + assert.Equal(t, testEntityIds[0], gotEntityId) // No filter options provided in URL expectFilter := etre.QueryFilter{} diff --git a/entity/store.go b/entity/store.go index b4bace7..dbc23a6 100644 --- a/entity/store.go +++ b/entity/store.go @@ -19,7 +19,7 @@ import ( // Store interface has methods needed to do CRUD operations on entities. type Store interface { - ReadEntities(context.Context, string, query.Query, etre.QueryFilter) ([]etre.Entity, error) + ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) CreateEntities(context.Context, WriteOp, []etre.Entity) ([]string, error) @@ -28,6 +28,8 @@ type Store interface { DeleteEntities(context.Context, WriteOp, query.Query) ([]etre.Entity, error) DeleteLabel(context.Context, WriteOp, string) (etre.Entity, error) + + StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan EntityResult } type store struct { @@ -45,65 +47,156 @@ func NewStore(entities map[string]*mongo.Collection, cdcStore cdc.Store, cfg con } } -// ReadEntities queries the db and returns a slice of Entity objects if -// something is found, a nil slice if nothing is found, and an error if one -// occurs. -func (s store) ReadEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { +// ReadEntity queries the db for a single entity. Returns nil if not found. +func (s store) ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { c, ok := s.coll[entityType] if !ok { - panic("invalid entity type passed to ReadEntities: " + entityType) - } - - // Distinct optimizaiton: unique values for the one return label. For example, - // "es -u node.metacluster zone=pd" returns a list of unique metacluster names. - // This is 10x faster than "es node.metacluster zone=pd | sort -u". - if len(f.ReturnLabels) == 1 && f.Distinct { - dr := c.Distinct(ctx, f.ReturnLabels[0], Filter(q)) - if err := dr.Err(); err != nil { - nfe := mongo.ErrNoDocuments - if errors.Is(err, nfe) { - // No documents found, return empty slice - return []etre.Entity{}, nil - } - return nil, s.dbError(ctx, err, "db-read-distinct") - } - var values []string - err := dr.Decode(&values) - if err != nil { - return nil, s.dbError(ctx, err, "db-read-distinct") - } - entities := make([]etre.Entity, len(values)) - for i, v := range values { - entities[i] = etre.Entity{f.ReturnLabels[0]: v} - } - return entities, nil + panic("invalid entity type passed to ReadEntity: " + entityType) } - // Find and return all matching entities + // Find and return the matching entity p := bson.M{} if len(f.ReturnLabels) > 0 { for _, label := range f.ReturnLabels { p[label] = 1 } // Only include _id if explicitly set in f.ReturnLabels. If not, - // ok is false and we must explicitly exlude it because MongoDB + // ok is false and we must explicitly exclude it because MongoDB // returns it by default. if _, ok := p["_id"]; !ok { p["_id"] = 0 } } - // Set batch size and projection - opts := options.Find().SetProjection(p).SetBatchSize(int32(s.config.BatchSize)) - cursor, err := c.Find(ctx, Filter(q), opts) + // Query, we should only get one row + q, err := query.Translate("_id=" + entityId) if err != nil { - return nil, s.dbError(ctx, err, "db-query") + return nil, s.dbError(ctx, err, "invalid-query") + } + + result := c.FindOne(ctx, Filter(q), options.FindOne().SetProjection(p)) + if err := result.Err(); err != nil { + nfe := mongo.ErrNoDocuments + if errors.Is(err, nfe) { + return nil, nil + } + return nil, s.dbError(ctx, result.Err(), "db-query") } - entities := []etre.Entity{} - if err := cursor.All(ctx, &entities); err != nil { + entity := etre.Entity{} + if err := result.Decode(&entity); err != nil { return nil, s.dbError(ctx, err, "db-read-cursor") } - return entities, nil + return entity, nil +} + +type EntityResult struct { + Entity etre.Entity + Err error +} + +func (s store) StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan EntityResult { + c, ok := s.coll[entityType] + if !ok { + panic("invalid entity type passed to StreamEntities: " + entityType) + } + + ch := make(chan EntityResult, 2*int32(s.config.BatchSize)) + go func() { + defer close(ch) + + // Distinct optimization: unique values for the one return label. For example, + // "es -u node.metacluster zone=pd" returns a list of unique metacluster names. + // This is 10x faster than "es node.metacluster zone=pd | sort -u". + if len(f.ReturnLabels) == 1 && f.Distinct { + dr := c.Distinct(ctx, f.ReturnLabels[0], Filter(q)) + if err := dr.Err(); err != nil { + nfe := mongo.ErrNoDocuments + if errors.Is(err, nfe) { + // No documents found, return to close the channel + return + } + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query-distinct")) + return + } + + var values []string + err := dr.Decode(&values) + if err != nil { + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query-distinct")) + return + } + // Distinct doesn't return a cursor, so we just have to loop and send the results + for _, v := range values { + s.writeEntityToChannel(ctx, ch, etre.Entity{f.ReturnLabels[0]: v}) + } + return + } + + // Find and return all matching entities + p := bson.M{} + if len(f.ReturnLabels) > 0 { + for _, label := range f.ReturnLabels { + p[label] = 1 + } + // Only include _id if explicitly set in f.ReturnLabels. If not, + // ok is false and we must explicitly exclude it because MongoDB + // returns it by default. + if _, ok := p["_id"]; !ok { + p["_id"] = 0 + } + } + + // Run the query + opts := options.Find().SetProjection(p).SetBatchSize(int32(s.config.BatchSize)) + cursor, err := c.Find(ctx, Filter(q), opts) + if err != nil { + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-query")) + return + } + defer cursor.Close(ctx) + + // Stream results + for cursor.Next(ctx) { + var entity etre.Entity + if err := cursor.Decode(&entity); err != nil { + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-read-cursor")) + return + } + s.writeEntityToChannel(ctx, ch, entity) + } + // Check for errors from iterating over cursor + if err := cursor.Err(); err != nil { + s.writeErrToChannel(ctx, ch, s.dbError(ctx, err, "db-read-cursor")) + return + } + }() + return ch +} + +func (s store) writeEntityToChannel(ctx context.Context, ch chan EntityResult, entity etre.Entity) { + select { + case <-ctx.Done(): + // The context was canceled or timed out. Bail out. + // We can't write the error to the channel because the receiver may have stopped listening. + // We depend on the receiver to see the ctx timeout as well. + return + case ch <- EntityResult{Entity: entity}: + // Successfully wrote to the channel + } + return +} + +func (s store) writeErrToChannel(ctx context.Context, ch chan EntityResult, err error) { + select { + case <-ctx.Done(): + // The context was canceled or timed out. Bail out. + // We can't write the error to the channel because the receiver may have stopped listening. + // We depend on the receiver to see the ctx timeout as well. + return + case ch <- EntityResult{Err: err}: + // Successfully wrote to the channel + } + return } // CreateEntities inserts many entities into DB. This method allows for partial diff --git a/entity/store_test.go b/entity/store_test.go index 84389a8..9fa7be5 100644 --- a/entity/store_test.go +++ b/entity/store_test.go @@ -94,7 +94,7 @@ func docs(entities []etre.Entity) []interface{} { // Read // -------------------------------------------------------------------------- -func TestReadEntitiesWithAllOperators(t *testing.T) { +func TestStreamEntitiesWithAllOperators(t *testing.T) { // Test all operators: in, notin, =, ==, !=, (has) y, (does not have) !foo, <, > // All values are set such that it only matches the first test node to make // testing easier and ensure we don't match the other entities. @@ -115,7 +115,7 @@ func TestReadEntitiesWithAllOperators(t *testing.T) { q, err := query.Translate(qs) require.NoError(t, err) - actual, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + actual, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) assert.Equal(t, expect, actual) } @@ -126,7 +126,7 @@ type readTest struct { expect []etre.Entity } -func TestReadEntitiesMatching(t *testing.T) { +func TestStreamEntitiesMatching(t *testing.T) { // Test various combinations of queries to ensure that we match and return // the correct entities. This is the fundamental job of Etre, so it should // be very thoroughly tested. @@ -157,13 +157,40 @@ func TestReadEntitiesMatching(t *testing.T) { q, err := query.Translate(rt.query) require.NoError(t, err) - got, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) assert.Equal(t, rt.expect, got) } } -func TestReadEntitiesFilterDistinct(t *testing.T) { +func TestReadEntityMatching(t *testing.T) { + // Test various combinations of queries to ensure that we match and return + // the correct entities. This is the fundamental job of Etre, so it should + // be very thoroughly tested. + store := setup(t, &mock.CDCStore{}) + readTests := []struct { + id string + expect etre.Entity + }{ + { + // found + id: testNodes[1]["_id"].(bson.ObjectID).Hex(), + expect: testNodes[1], + }, + { + // not found + id: "bogus", + expect: nil, + }, + } + for _, rt := range readTests { + got, err := store.ReadEntity(context.Background(), entityType, rt.id, etre.QueryFilter{}) + require.NoError(t, err) + assert.Equal(t, rt.expect, got) + } +} + +func TestStreamEntitiesFilterDistinct(t *testing.T) { // Test that etre.QueryFilter{Distinct: true} returns a list of unique values // for one label. The 1st test node has y=a and the 2nd and 3rd both have y=b, // so the unique values are [a,b]. @@ -175,7 +202,7 @@ func TestReadEntitiesFilterDistinct(t *testing.T) { ReturnLabels: []string{"y"}, // only works with 1 return label Distinct: true, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) require.NoError(t, err) expect := []etre.Entity{ @@ -185,7 +212,7 @@ func TestReadEntitiesFilterDistinct(t *testing.T) { assert.Equal(t, expect, got) } -func TestReadEntitiesFilterDistinctNoResult(t *testing.T) { +func TestStreamEntitiesFilterDistinctNoResult(t *testing.T) { // Test that etre.QueryFilter{Distinct: true} returns an empty list // if no rows match the query. store := setup(t, &mock.CDCStore{}) @@ -196,12 +223,12 @@ func TestReadEntitiesFilterDistinctNoResult(t *testing.T) { ReturnLabels: []string{"BOGUS"}, // only works with 1 return label Distinct: true, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) require.NoError(t, err) assert.Empty(t, got) } -func TestReadEntitiesFilterReturnLabels(t *testing.T) { +func TestStreamEntitiesFilterReturnLabels(t *testing.T) { // Test that etre.QueryFilter{ReturnLabels: []string{x}} returns only that // label and not the others (y, z, bar, foo). We'll select/match by label y // but return only label x. @@ -217,17 +244,33 @@ func TestReadEntitiesFilterReturnLabels(t *testing.T) { {"x": int64(4)}, {"x": int64(6)}, } - got, err := store.ReadEntities(context.Background(), entityType, q, f) + got, err := readStream(store.StreamEntities(context.Background(), entityType, q, f)) + require.NoError(t, err) + assert.Equal(t, expect, got) +} + +func TestReadEntityFilterReturnLabels(t *testing.T) { + // Test that etre.QueryFilter{ReturnLabels: []string{x}} returns only that + // label and not the others (y, z, bar, foo). We'll select/match by label y + // but return only label x. + store := setup(t, &mock.CDCStore{}) + + f := etre.QueryFilter{ + ReturnLabels: []string{"x"}, // testing this + } + expect := etre.Entity{"x": int64(2)} + + got, err := store.ReadEntity(context.Background(), entityType, testNodes[0]["_id"].(bson.ObjectID).Hex(), f) require.NoError(t, err) assert.Equal(t, expect, got) } -func TestReadEntitiesFilterReturnMetalabels(t *testing.T) { +func TestStreamEntitiesFilterReturnMetalabels(t *testing.T) { store := setup(t, &mock.CDCStore{}) q, err := query.Translate("y=a") require.NoError(t, err) - actual, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}}) + actual, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}})) require.NoError(t, err) expect := []etre.Entity{ @@ -237,6 +280,16 @@ func TestReadEntitiesFilterReturnMetalabels(t *testing.T) { assert.Equal(t, expect, actual) } +func TestReadEntityFilterReturnMetalabels(t *testing.T) { + store := setup(t, &mock.CDCStore{}) + + actual, err := store.ReadEntity(context.Background(), entityType, testNodes[0]["_id"].(bson.ObjectID).Hex(), etre.QueryFilter{ReturnLabels: []string{"_id", "_type", "_rev", "y", "_created", "_updated"}}) + require.NoError(t, err) + + expect := etre.Entity{"_id": testNodes[0]["_id"], "_type": entityType, "_rev": int64(0), "y": "a", "_created": testNodes[0]["_created"], "_updated": testNodes[0]["_updated"]} + assert.Equal(t, expect, actual) +} + // -------------------------------------------------------------------------- // Create // -------------------------------------------------------------------------- @@ -745,7 +798,7 @@ func TestDeleteLabel(t *testing.T) { // The foo label should no longer be set on the entity q, _ := query.Translate("y=a") - gotNew, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + gotNew, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) e := etre.Entity{} @@ -780,3 +833,14 @@ func TestDeleteLabel(t *testing.T) { } assert.Equal(t, expectEvent, gotEvents) } + +func readStream(ch <-chan entity.EntityResult) ([]etre.Entity, error) { + entities := []etre.Entity{} + for r := range ch { + if r.Err != nil { + return nil, r.Err + } + entities = append(entities, r.Entity) + } + return entities, nil +} diff --git a/entity/v09_test.go b/entity/v09_test.go index 1f8a9e6..8a11ed9 100644 --- a/entity/v09_test.go +++ b/entity/v09_test.go @@ -293,7 +293,7 @@ func TestV09DeleteLabel(t *testing.T) { // The foo label should no longer be set on the entity q, _ := query.Translate("x=a") - gotNew, err := store.ReadEntities(context.Background(), entityType, q, etre.QueryFilter{}) + gotNew, err := readStream(store.StreamEntities(context.Background(), entityType, q, etre.QueryFilter{})) require.NoError(t, err) e := etre.Entity{} diff --git a/test/mock/entity.go b/test/mock/entity.go index edf39bd..43fa3a1 100644 --- a/test/mock/entity.go +++ b/test/mock/entity.go @@ -11,12 +11,13 @@ import ( ) type EntityStore struct { - ReadEntitiesFunc func(context.Context, string, query.Query, etre.QueryFilter) ([]etre.Entity, error) + ReadEntityFunc func(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) DeleteEntityLabelFunc func(context.Context, entity.WriteOp, string) (etre.Entity, error) CreateEntitiesFunc func(context.Context, entity.WriteOp, []etre.Entity) ([]string, error) UpdateEntitiesFunc func(context.Context, entity.WriteOp, query.Query, etre.Entity) ([]etre.Entity, error) DeleteEntitiesFunc func(context.Context, entity.WriteOp, query.Query) ([]etre.Entity, error) DeleteLabelFunc func(context.Context, entity.WriteOp, string) (etre.Entity, error) + StreamEntitiesFunc func(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult } func (s EntityStore) DeleteEntityLabel(ctx context.Context, wo entity.WriteOp, label string) (etre.Entity, error) { @@ -33,9 +34,9 @@ func (s EntityStore) CreateEntities(ctx context.Context, wo entity.WriteOp, enti return nil, nil } -func (s EntityStore) ReadEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) ([]etre.Entity, error) { - if s.ReadEntitiesFunc != nil { - return s.ReadEntitiesFunc(ctx, entityType, q, f) +func (s EntityStore) ReadEntity(ctx context.Context, entityType string, entityId string, f etre.QueryFilter) (etre.Entity, error) { + if s.ReadEntityFunc != nil { + return s.ReadEntityFunc(ctx, entityType, entityId, f) } return nil, nil } @@ -60,3 +61,25 @@ func (s EntityStore) DeleteLabel(ctx context.Context, wo entity.WriteOp, label s } return etre.Entity{}, nil } + +func (s EntityStore) StreamEntities(ctx context.Context, entityType string, q query.Query, f etre.QueryFilter) <-chan entity.EntityResult { + if s.StreamEntitiesFunc != nil { + return s.StreamEntitiesFunc(ctx, entityType, q, f) + } + return DoStreamEntities(nil, nil) +} + +func DoStreamEntities(entities []etre.Entity, err error) <-chan entity.EntityResult { + ch := make(chan entity.EntityResult) + go func() { + defer close(ch) + if err != nil { + ch <- entity.EntityResult{Err: err} + return + } + for _, e := range entities { + ch <- entity.EntityResult{Entity: e} + } + }() + return ch +}