The implementation of FindTraceIDs function for ElasticSearch reader.

plugin/storage/es/spanstore/reader.go

@@ -234,9 +234,35 @@ func (s *SpanReader) FindTraces(ctx context.Context, traceQuery *spanstore.Trace
 	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
+	if err := validateQuery(traceQuery); err != nil {
+		return nil, err
+	}
+	if traceQuery.NumTraces == 0 {
+		traceQuery.NumTraces = defaultNumTraces
+	}
+
+	esTraceIDs, err := s.findTraceIDs(ctx, traceQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	traceIDs := []model.TraceID{}
+	for _, ID := range esTraceIDs {
+		if len(traceIDs) >= traceQuery.NumTraces {
+			break
+		}
+
+		traceID, err := model.TraceIDFromString(ID)
+		if err != nil {
+			return nil, err
+		}
+
+		traceIDs = append(traceIDs, traceID)
+	}
+
+	return traceIDs, nil
 }
 
 func (s *SpanReader) multiRead(ctx context.Context, traceIDs []string, startTime, endTime time.Time) ([]*model.Trace, error) {

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:    1,
+		}
+
+		traceIDs, err := r.reader.FindTraceIDs(context.Background(), traceIDsQuery)
+		require.NoError(t, err)
+		assert.Len(t, traceIDs, 1)
+		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)
+	})
+}
+
+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)
-		assert.EqualError(t, err, "not implemented")
+	})
+}
+
+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, `strconv.ParseUint: parsing "sdfsdfdss": invalid syntax`)
+		assert.Len(t, traceIDs, 0)
 	})
 }