Skip to content

Commit

Permalink
minor fixes
Browse files Browse the repository at this point in the history 
Signed-off-by: Manik2708 <[email protected]>
  • Loading branch information
@Manik2708
Manik2708 committed Jan 21, 2025
1 parent 6c26cbe commit 084ec66
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 32 deletions.
2 changes: 1 addition & 1 deletion plugin/storage/es/spanstore/reader.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -298,7 +298,7 @@ func (s *SpanReader) GetOperations(
currentTime,
cfg.RolloverFrequencyAsNegativeDuration(s.serviceIndex.RolloverFrequency),
)
operations, err := s.serviceOperationStorage.getOperations(ctx, jaegerIndices, query.ServiceName, query.SpanKind, s.maxDocCount)
operations, err := s.serviceOperationStorage.getOperationsWithKind(ctx, jaegerIndices, query.ServiceName, query.SpanKind, s.maxDocCount)
if err != nil {
return nil, err
}
Expand Down
28 changes: 14 additions & 14 deletions plugin/storage/es/spanstore/reader_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -629,7 +629,7 @@ func testGetWithKind(typ string, t *testing.T, testKind bool) {
rawMessage := []byte(`{"buckets": [{"key": "123","doc_count": 16}]}`)
goodAggregations[typ] = (*json.RawMessage)(&rawMessage)
var filterRawMessage json.RawMessage
if typ == operationsAggregation {
if typ == operationsAggName {
if testKind {
filterRawMessage = rawMessage
}
Expand All @@ -646,15 +646,15 @@ func testGetWithKind(typ string, t *testing.T, testKind bool) {
caption: typ + " search error",
searchError: errors.New("Search failure"),
expectedError: func() string {
if typ == operationsAggregation {
if typ == operationsAggName {
return "search operations failed: Search failure"
}
return "search services failed: Search failure"
},
},
}

if (typ == operationsAggregation && testKind) || (typ != operationsAggregation) {
if (typ == operationsAggName && testKind) || (typ != operationsAggName) {
testCase := testGetStruct{
caption: typ + " search error",
searchResult: &elastic.SearchResult{Aggregations: elastic.Aggregations(badAggregations)},
Expand All @@ -670,16 +670,16 @@ func testGetWithKind(typ string, t *testing.T, testKind bool) {
caption: typ + " full behavior with kind",
searchResult: &elastic.SearchResult{Aggregations: goodAggregations},
expectedOutput: map[string]any{
operationsAggregation: []spanstore.Operation{{Name: "123", SpanKind: "server"}},
"default": []string{"123"},
operationsAggName: []spanstore.Operation{{Name: "123", SpanKind: "server"}},
"default": []string{"123"},
},
expectedError: func() string {
return ""
},
})
}

if typ == operationsAggregation && !testKind {
if typ == operationsAggName && !testKind {
score := 0.6931471
msg := json.RawMessage(`{"operationName": "123"}`)
hitModel := &elastic.SearchHits{
Expand All @@ -699,22 +699,22 @@ func testGetWithKind(typ string, t *testing.T, testKind bool) {
caption: typ + " full behavior",
searchResult: &elastic.SearchResult{Hits: hitModel},
expectedOutput: map[string]any{
operationsAggregation: []spanstore.Operation{{Name: "123"}},
"default": []string{"123"},
operationsAggName: []spanstore.Operation{{Name: "123"}},
"default": []string{"123"},
},
expectedError: func() string {
return ""
},
})
}

if typ != operationsAggregation {
if typ != operationsAggName {
testCases = append(testCases, testGetStruct{
caption: typ + " full behavior",
searchResult: &elastic.SearchResult{Aggregations: goodAggregations},
expectedOutput: map[string]any{
operationsAggregation: []spanstore.Operation{{Name: "123"}},
"default": []string{"123"},
operationsAggName: []spanstore.Operation{{Name: "123"}},
"default": []string{"123"},
},
expectedError: func() string {
return ""
Expand Down Expand Up @@ -749,7 +749,7 @@ func returnSearchFunc(typ string, r *spanReaderTest, testKind bool) (any, error)
switch typ {
case servicesAggregation:
return r.reader.GetServices(context.Background())
case operationsAggregation:
case operationsAggName:
if testKind {
return r.reader.GetOperations(
context.Background(),
Expand Down Expand Up @@ -1025,7 +1025,7 @@ func TestFindTraceIDs(t *testing.T) {
}{
{traceIDAggregation},
{servicesAggregation},
{operationsAggregation},
{operationsAggName},
}
for _, testCase := range testCases {
t.Run(testCase.aggregrationID, func(t *testing.T) {
Expand Down Expand Up @@ -1083,7 +1083,7 @@ func mockSearchServiceWithSpanKind(r *spanReaderTest, inputHasSpanKind bool) *mo
searchService.On("IgnoreUnavailable", mock.AnythingOfType("bool")).Return(searchService)
searchService.On("Aggregation", stringMatcher(servicesAggregation), mock.MatchedBy(matchTermsAggregation)).Return(searchService)
if inputHasSpanKind {
searchService.On("Aggregation", stringMatcher(operationsAggregation), mock.MatchedBy(matchTermsAggregation)).Return(searchService)
searchService.On("Aggregation", stringMatcher(operationsAggName), mock.MatchedBy(matchTermsAggregation)).Return(searchService)
} else {
searchService.On("FetchSourceContext", mock.AnythingOfType("*elastic.FetchSourceContext"), mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(searchService)
searchService.On("Size", mock.AnythingOfType("int")).Return(searchService)
Expand Down
28 changes: 13 additions & 15 deletions plugin/storage/es/spanstore/service_operation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@ import (
)

const (
serviceName = "serviceName"
spanKind = "spanKind"
operationsAggregation = "distinct_operations"
servicesAggregation = "distinct_services"
serviceName = "serviceName"
spanKindField = "spanKindField"
operationsAggName = "distinct_operations"
servicesAggregation = "distinct_services"
)

// ServiceOperationStorage stores service to operation pairs.
Expand Down Expand Up @@ -100,28 +100,26 @@ func getServicesAggregation(maxDocCount int) elastic.Query {
Size(maxDocCount) // ES deprecated size omission for aggregating all. https://github.com/elastic/elasticsearch/issues/18838
}

func (s *ServiceOperationStorage) getOperations(ctx context.Context, indices []string, service, kind string, maxDocCount int) ([]spanstore.Operation, error) {
func (s *ServiceOperationStorage) getOperationsWithKind(ctx context.Context, indices []string, service, kind string, maxDocCount int) ([]spanstore.Operation, error) {
var searchService es.SearchService
if kind != "" {
serviceFilter := getOperationsAggregation(maxDocCount)
serviceNameQuery := elastic.NewTermQuery(serviceName, service)
serviceKindQuery := elastic.NewTermQuery(spanKind, kind)
serviceQuery := elastic.NewBoolQuery().Must(serviceNameQuery, serviceKindQuery)
searchService = s.client().Search(indices...).
Size(0).
Query(serviceQuery).
Size(0). // set to 0 because we don't want actual documents.
Query(elastic.NewBoolQuery().Must(
elastic.NewTermQuery(serviceName, service),
elastic.NewTermQuery(spanKindField, kind))).
IgnoreUnavailable(true).
Aggregation(operationsAggregation, serviceFilter)
Aggregation(operationsAggName, getOperationsAggregation(maxDocCount))
searchResult, err := searchService.Do(ctx)
if err != nil {
return nil, fmt.Errorf("search operations failed: %w", es.DetailedError(err))
}
if searchResult.Aggregations == nil {
return []spanstore.Operation{}, nil
}
bucket, found := searchResult.Aggregations.Terms(operationsAggregation)
bucket, found := searchResult.Aggregations.Terms(operationsAggName)
if !found {
return nil, errors.New("could not find aggregation of " + operationsAggregation)
return nil, errors.New("could not find aggregation of " + operationsAggName)
}
operationNamesBucket := bucket.Buckets
return bucketOfOperationNamesToOperationsArray(operationNamesBucket, kind)
Expand All @@ -130,7 +128,7 @@ func (s *ServiceOperationStorage) getOperations(ctx context.Context, indices []s
searchService = s.client().Search(indices...).
Query(serviceQuery).
IgnoreUnavailable(true).
FetchSourceContext(elastic.NewFetchSourceContext(true).Include(spanKind, operationNameField)).
FetchSourceContext(elastic.NewFetchSourceContext(true).Include(spanKindField, operationNameField)).
Size(maxDocCount)
searchResult, err := searchService.Do(ctx)
if err != nil {
Expand Down
4 changes: 2 additions & 2 deletions plugin/storage/es/spanstore/service_operation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,8 +122,8 @@ func TestSpanReader_GetServices(t *testing.T) {
}

func TestSpanReader_GetOperations(t *testing.T) {
testGet(operationsAggregation, t)
testGetWithKind(operationsAggregation, t, true)
testGet(operationsAggName, t)
testGetWithKind(operationsAggName, t, true)
}

func TestSpanReader_GetServicesEmptyIndex(t *testing.T) {
Expand Down

0 comments on commit 084ec66

Please sign in to comment.