The implementation of FindTraceIDs function for ElasticSearch reader.

plugin/storage/es/spanstore/reader.go

@@ -221,22 +221,34 @@ func (s *SpanReader) FindTraces(ctx context.Context, traceQuery *spanstore.Trace
 	span, ctx := opentracing.StartSpanFromContext(ctx, "FindTraces")
 	defer span.Finish()
 
-	if err := validateQuery(traceQuery); err != nil {
-		return nil, err
-	}
-	if traceQuery.NumTraces == 0 {
-		traceQuery.NumTraces = defaultNumTraces
-	}
-	uniqueTraceIDs, err := s.findTraceIDs(ctx, traceQuery)
+	uniqueTraceIDs, err := s.validateQueryAndFindTraceIDs(ctx, traceQuery)
 	if err != nil {
 		return nil, err
 	}
 	return s.multiRead(ctx, uniqueTraceIDs, traceQuery.StartTimeMin, traceQuery.StartTimeMax)
 }
 
-// FindTraceIDs is not implemented.
+// FindTraceIDs retrieves traces IDs that match the traceQuery
 func (s *SpanReader) FindTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]model.TraceID, error) {
-	return nil, errors.New("not implemented") // TODO: Implement
+	span, ctx := opentracing.StartSpanFromContext(ctx, "FindTraceIDs")
+	defer span.Finish()
+
+	esTraceIDs, err := s.validateQueryAndFindTraceIDs(ctx, traceQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	traceIDs := make([]model.TraceID, len(esTraceIDs))
+	for i, ID := range esTraceIDs {
+		traceID, err := model.TraceIDFromString(ID)
+		if err != nil {
+			return nil, errors.Wrap(err, fmt.Sprintf("Making traceID from string '%s' failed", ID))
+		}
+
+		traceIDs[i] = traceID
+	}
+
+	return traceIDs, nil
 }
 
 func (s *SpanReader) multiRead(ctx context.Context, traceIDs []string, startTime, endTime time.Time) ([]*model.Trace, error) {
@@ -317,6 +329,20 @@ func (s *SpanReader) multiRead(ctx context.Context, traceIDs []string, startTime
 	return traces, nil
 }
 
+func (s *SpanReader) validateQueryAndFindTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]string, error) {
+	childSpan, _ := opentracing.StartSpanFromContext(ctx, "validateQueryAndFindTraceIDs")
+	defer childSpan.Finish()
+
+	if err := validateQuery(traceQuery); err != nil {
+		return nil, err
+	}
+	if traceQuery.NumTraces == 0 {
+		traceQuery.NumTraces = defaultNumTraces
+	}
+
+	return s.findTraceIDs(ctx, traceQuery)
+}
+
 func validateQuery(p *spanstore.TraceQueryParameters) error {
 	if p == nil {
 		return ErrMalformedRequestObject
@@ -339,6 +365,7 @@ func validateQuery(p *spanstore.TraceQueryParameters) error {
 func (s *SpanReader) findTraceIDs(ctx context.Context, traceQuery *spanstore.TraceQueryParameters) ([]string, error) {
 	childSpan, _ := opentracing.StartSpanFromContext(ctx, "findTraceIDs")
 	defer childSpan.Finish()
+
 	//  Below is the JSON body to our HTTP GET request to ElasticSearch. This function creates this.
 	// {
 	//      "size": 0,

plugin/storage/es/spanstore/reader_test.go

@@ -668,11 +668,142 @@ func TestFindTraceIDs(t *testing.T) {
 	testGet(traceIDAggregation, t)
 }
 
-func TestFindTraceIDNotImplemented(t *testing.T) {
+func TestSpanReader_TestFindTraceIDs(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+	rawMessage := []byte(`{"buckets": [{"key": "1","doc_count": 16},{"key": "2","doc_count": 16},{"key": "3","doc_count": 16}]}`)
+	goodAggregations[traceIDAggregation] = (*json.RawMessage)(&rawMessage)
+
+	withSpanReader(func(r *spanReaderTest) {
+		// find trace IDs
+		mockSearchService(r).Return(&elastic.SearchResult{Aggregations: elastic.Aggregations(goodAggregations)}, nil)
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: serviceName,
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+			NumTraces:    3,
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.NoError(t, err)
+		assert.Len(t, traceIDs, 3)
+		assert.EqualValues(t, 1, traceIDs[0].Low)
+	})
+}
+
+func TestSpanReader_FindTraceIDsInvalidQuery(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+	rawMessage := []byte(`{"buckets": [{"key": "1","doc_count": 16},{"key": "2","doc_count": 16},{"key": "3","doc_count": 16}]}`)
+	goodAggregations[traceIDAggregation] = (*json.RawMessage)(&rawMessage)
+
 	withSpanReader(func(r *spanReaderTest) {
-		traceIDs, err := r.reader.FindTraceIDs(context.Background(), nil)
+		mockSearchService(r).Return(&elastic.SearchResult{Aggregations: elastic.Aggregations(goodAggregations)}, nil)
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: "",
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.Error(t, err)
 		assert.Nil(t, traceIDs)
-		assert.EqualError(t, err, "not implemented")
+	})
+}
+
+func TestSpanReader_FindTraceIDsAggregationFailure(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+
+	withSpanReader(func(r *spanReaderTest) {
+		mockSearchService(r).Return(&elastic.SearchResult{Aggregations: elastic.Aggregations(goodAggregations)}, nil)
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: serviceName,
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.Error(t, err)
+		assert.Nil(t, traceIDs)
+	})
+}
+
+func TestSpanReader_FindTraceIDsNoTraceIDs(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+	rawMessage := []byte(`{"buckets": []}`)
+	goodAggregations[traceIDAggregation] = (*json.RawMessage)(&rawMessage)
+
+	withSpanReader(func(r *spanReaderTest) {
+		mockSearchService(r).Return(&elastic.SearchResult{Aggregations: elastic.Aggregations(goodAggregations)}, nil)
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: serviceName,
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.NoError(t, err)
+		assert.Len(t, traceIDs, 0)
+	})
+}
+
+func TestSpanReader_FindTraceIDsReadTraceIDsFailure(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+	rawMessage := []byte(`{"buckets": [{"key": "1","doc_count": 16},{"key": "2","doc_count": 16}]}`)
+	goodAggregations[traceIDAggregation] = (*json.RawMessage)(&rawMessage)
+
+	withSpanReader(func(r *spanReaderTest) {
+		mockSearchService(r).Return(nil, errors.New("read error"))
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: serviceName,
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.EqualError(t, err, "Search service failed: read error")
+		assert.Len(t, traceIDs, 0)
+	})
+}
+
+func TestSpanReader_FindTraceIDsIncorrectTraceIDFailure(t *testing.T) {
+	goodAggregations := make(map[string]*json.RawMessage)
+	rawMessage := []byte(`{"buckets": [{"key": "sdfsdfdssd234nsdvsdfldjsf","doc_count": 16},{"key": "2","doc_count": 16}]}`)
+	goodAggregations[traceIDAggregation] = (*json.RawMessage)(&rawMessage)
+
+	withSpanReader(func(r *spanReaderTest) {
+		mockSearchService(r).Return(&elastic.SearchResult{Aggregations: elastic.Aggregations(goodAggregations)}, nil)
+
+		traceIDsQuery := &spanstore.TraceQueryParameters{
+			ServiceName: serviceName,
+			Tags: map[string]string{
+				"hello": "world",
+			},
+			StartTimeMin: time.Now().Add(-1 * time.Hour),
+			StartTimeMax: time.Now(),
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.EqualError(t, err, `Making traceID from string 'sdfsdfdssd234nsdvsdfldjsf' failed: strconv.ParseUint: parsing "sdfsdfdss": invalid syntax`)
+		assert.Len(t, traceIDs, 0)
 	})
 }