Make Traces span limit settable

grafana · Jan 21, 2025 · 732c00d · 732c00d
1 parent ea3de6b
commit 732c00d
Show file tree

Hide file tree

Showing 9 changed files with 117 additions and 30 deletions.
diff --git a/pkg/opensearch/client/search_request.go b/pkg/opensearch/client/search_request.go
@@ -429,7 +429,7 @@ type AggBuilder interface {
 	DateHistogram(key, field string, fn func(a *DateHistogramAgg, b AggBuilder)) AggBuilder
 	Terms(key, field string, fn func(a *TermsAggregation, b AggBuilder)) AggBuilder
 	Filters(key string, fn func(a *FiltersAggregation, b AggBuilder)) AggBuilder
-	TraceList() AggBuilder
+	TraceList(spanLimit int) AggBuilder
 	ServiceMap() AggBuilder
 	Stats() AggBuilder
 	GeoHashGrid(key, field string, fn func(a *GeoHashGridAggregation, b AggBuilder)) AggBuilder
@@ -618,7 +618,10 @@ func (b *aggBuilderImpl) Stats() AggBuilder {
 }
 
 // TraceList sets the "aggs" object of the query to OpenSearch for the trace list
-func (b *aggBuilderImpl) TraceList() AggBuilder {
+func (b *aggBuilderImpl) TraceList(spanLimit int) AggBuilder {
+	if spanLimit > 10000 {
+		spanLimit = 10000
+	}
 	aggDef := &aggDefinition{
 		key: "traces",
 		aggregation: &AggContainer{
@@ -629,7 +632,7 @@ func (b *aggBuilderImpl) TraceList() AggBuilder {
 				Order map[string]string `json:"order"`
 			}{
 				Field: "traceId",
-				Size:  100,
+				Size:  spanLimit,
 				Order: map[string]string{"_key": "asc"},
 			},
 			Aggs: AggArray{

diff --git a/pkg/opensearch/lucene_handler.go b/pkg/opensearch/lucene_handler.go
@@ -17,6 +17,11 @@ import (
 	"github.com/grafana/opensearch-datasource/pkg/utils"
 )
 
+const (
+	defaultLogsSize  = 500
+	defaultSpanLimit = 1000
+)
+
 type luceneHandler struct {
 	client     client.Client
 	reqQueries []backend.DataQuery
@@ -74,14 +79,15 @@ func (h *luceneHandler) processQuery(q *Query) error {
 			aggBuilder := b.Agg()
 			aggBuilder.Stats()
 		default:
+			limit := utils.StringToIntWithDefaultValue(q.SpanLimit, defaultSpanLimit)
 			if traceId != "" {
-				b.Size(1000)
+				b.Size(limit)
 				b.SetTraceSpansFilters(toMs, fromMs, traceId)
 			} else {
-				b.Size(1000)
+				b.Size(limit)
 				b.SetTraceListFilters(toMs, fromMs, q.RawQuery)
 				aggBuilder := b.Agg()
-				aggBuilder.TraceList()
+				aggBuilder.TraceList(limit)
 			}
 		}
 		return nil
@@ -123,7 +129,7 @@ func processLogsQuery(q *Query, b *client.SearchRequestBuilder, from, to int64,
 	sizeString := metric.Settings.Get("size").MustString()
 	size, err := strconv.Atoi(sizeString)
 	if err != nil {
-		size = defaultSize
+		size = defaultLogsSize
 	}
 	b.Size(size)
 
@@ -156,8 +162,6 @@ func setIntPath(settings *simplejson.Json, path ...string) {
 	}
 }
 
-const defaultSize = 500
-
 func processDocumentQuery(q *Query, b *client.SearchRequestBuilder, defaultTimeField string) {
 	metric := q.Metrics[0]
 	order := metric.Settings.Get("order").MustString()
@@ -167,7 +171,7 @@ func processDocumentQuery(q *Query, b *client.SearchRequestBuilder, defaultTimeF
 	sizeString := metric.Settings.Get("size").MustString()
 	size, err := strconv.Atoi(sizeString)
 	if err != nil {
-		size = defaultSize
+		size = defaultLogsSize
 	}
 	b.Size(size)
 }

diff --git a/pkg/opensearch/models.go b/pkg/opensearch/models.go
@@ -21,6 +21,7 @@ type Query struct {
 	RefID           string
 	Format          string
 	TimeRange       backend.TimeRange
+	SpanLimit       string `json:"spanLimit"`
 
 	// serviceMapInfo is used on the backend to pass information for service map queries
 	serviceMapInfo serviceMapInfo

diff --git a/pkg/opensearch/query_request.go b/pkg/opensearch/query_request.go
@@ -90,6 +90,8 @@ func parse(reqQueries []backend.DataQuery) ([]*Query, error) {
 		alias := model.Get("alias").MustString("")
 		format := model.Get("format").MustString("")
 
+		tracesSpanLimit := model.Get("spanLimit").MustString()
+
 		// For queries requesting the service map, we inject extra queries to handle retrieving
 		// the required information
 		hasServiceMap := model.Get("serviceMap").MustBool(false)
@@ -136,6 +138,7 @@ func parse(reqQueries []backend.DataQuery) ([]*Query, error) {
 					serviceMapInfo:  serviceMapInfo{Type: ServiceMap},
 				},
 			)
+			tracesSpanLimit = ""
 		}
 
 		queries = append(queries, &Query{
@@ -148,6 +151,7 @@ func parse(reqQueries []backend.DataQuery) ([]*Query, error) {
 			Interval:        q.Interval,
 			RefID:           q.RefID,
 			Format:          format,
+			SpanLimit:       tracesSpanLimit,
 		})
 	}
 

diff --git a/pkg/opensearch/query_request_test.go b/pkg/opensearch/query_request_test.go
@@ -1366,7 +1366,8 @@ func Test_trace_list(t *testing.T) {
 		  "timeField": "@timestamp",
 		  "datasourceId": 2020,
 		  "intervalMs": 10000,
-		  "maxDataPoints": 1124
+		  "maxDataPoints": 1124,
+      "spanLimit": "500"
 		}`, from, to, 15*time.Second)
 	require.NoError(t, err)
 
@@ -1400,7 +1401,7 @@ func Test_trace_list(t *testing.T) {
 		Order map[string]string `json:"order"`
 	}{
 		Field: "traceId",
-		Size:  100,
+		Size:  500,
 		Order: map[string]string{"_key": "asc"},
 	}, actualRequest.Aggs[0].Aggregation.Aggregation)
 

diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
@@ -3,6 +3,7 @@ package utils
 import (
 	"encoding/json"
 	"errors"
+	"strconv"
 	"strings"
 	"time"
 
@@ -96,3 +97,15 @@ func SpanHasError(spanEvents []interface{}) bool {
 }
 
 func Pointer[T any](v T) *T { return &v }
+
+func StringToIntWithDefaultValue(valueStr string, defaultValue int) int {
+	value, err := strconv.Atoi(valueStr)
+	if err != nil {
+		value = defaultValue
+	}
+	// In our case, 0 is not a valid value and in this case we default to defaultValue
+	if value == 0 {
+		value = defaultValue
+	}
+	return value
+}
diff --git a/src/components/QueryEditor/LuceneQueryEditor/LuceneQueryEditor.test.tsx b/src/components/QueryEditor/LuceneQueryEditor/LuceneQueryEditor.test.tsx
@@ -1,5 +1,5 @@
 import { DataSourcePluginMeta } from '@grafana/data';
-import { render, screen } from '@testing-library/react';
+import { fireEvent, render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { OpenSearchDatasource } from 'opensearchDatasource';
 import React from 'react';
@@ -80,19 +80,55 @@ describe('LuceneQueryEditor', () => {
     expect(screen.queryByText('Traces')).not.toBeInTheDocument();
 
     await userEvent.click(screen.getByText('Metric'));
-
     expect(screen.queryByText('Traces')).toBeInTheDocument();
-
     await userEvent.click(screen.getByText('Traces'));
 
     expect(mockOnChange).toHaveBeenCalledTimes(1);
-
     expect(mockOnChange.mock.calls[0][0].luceneQueryType).toBe('Traces');
   });
 
-  it('renders the service map switch when traces is selected', async () => {
+  it('renders the spanLimit field when traces is selected', async () => {
+    const mockQuery = createMockQuery();
+    mockQuery.luceneQueryType = LuceneQueryType.Traces;
+    const mockOnChange = createMockOnChange();
+    const mockDatasource = createMockDatasource();
+
+    render(
+      <OpenSearchProvider query={mockQuery} onChange={mockOnChange} datasource={mockDatasource}>
+        <LuceneQueryEditor query={mockQuery} onChange={mockOnChange} />
+      </OpenSearchProvider>
+    );
+
+    const limitElement = screen.getByTestId('span-limit-input');
+    expect(limitElement).toBeInTheDocument();
+    await userEvent.type(limitElement!, '200');
+    fireEvent.blur(limitElement!);
+
+    expect(mockOnChange).toHaveBeenCalledTimes(1);
+    expect(mockOnChange.mock.calls[0][0].spanLimit).toBe('200');
+  });
+
+  it('renders the service map switch and span limit field when traces is selected', async () => {
+    const mockQuery = createMockQuery({
+      luceneQueryType: LuceneQueryType.Traces,
+    });
+    const mockOnChange = createMockOnChange();
+    const mockDatasource = createMockDatasource();
+
+    render(
+      <OpenSearchProvider query={mockQuery} onChange={mockOnChange} datasource={mockDatasource}>
+        <LuceneQueryEditor query={mockQuery} onChange={mockOnChange} />
+      </OpenSearchProvider>
+    );
+
+    expect(screen.queryByText('Service Map')).toBeInTheDocument();
+    expect(screen.queryByText('Span Limit')).toBeInTheDocument();
+    jest.clearAllMocks();
+  });
+  it('does not render the span limit input if service map is selected', async () => {
     const mockQuery = createMockQuery({
       luceneQueryType: LuceneQueryType.Traces,
+      serviceMap: true,
     });
     const mockOnChange = createMockOnChange();
     const mockDatasource = createMockDatasource();
@@ -104,6 +140,8 @@ describe('LuceneQueryEditor', () => {
     );
 
     expect(screen.queryByText('Service Map')).toBeInTheDocument();
+    // Span Limit should not be rendered if service map is selected
+    expect(screen.queryByText('Span Limit')).not.toBeInTheDocument();
     jest.clearAllMocks();
   });
 });
diff --git a/src/components/QueryEditor/LuceneQueryEditor/LuceneQueryEditor.tsx b/src/components/QueryEditor/LuceneQueryEditor/LuceneQueryEditor.tsx
@@ -1,4 +1,4 @@
-import { Segment, InlineSegmentGroup, InlineField, InlineSwitch } from '@grafana/ui';
+import { Segment, InlineSegmentGroup, InlineField, InlineSwitch, Input } from '@grafana/ui';
 import { useNextId } from 'hooks/useNextId';
 import React from 'react';
 import { LuceneQueryType, OpenSearchQuery } from 'types';
@@ -20,6 +20,7 @@ const toOption = (queryType: LuceneQueryType) => {
 
 export const LuceneQueryEditor = (props: LuceneQueryEditorProps) => {
   const luceneQueryType = props.query.luceneQueryType || LuceneQueryType.Metric;
+  const serviceMapSet = props.query.serviceMap || false;
   const nextId = useNextId();
 
   const setLuceneQueryType = (newQueryType: LuceneQueryType) => {
@@ -43,18 +44,39 @@ export const LuceneQueryEditor = (props: LuceneQueryEditorProps) => {
             value={toOption(luceneQueryType)}
           />
           {luceneQueryType === LuceneQueryType.Traces && (
-            <InlineField label="Service Map" tooltip={"Request and display service map data for trace(s)"}>
-              <InlineSwitch
-                value={props.query.serviceMap || false}
-                onChange={(event) => {
-                  const newVal = event.currentTarget.checked;
-                  props.onChange({
-                    ...props.query,
-                    serviceMap: newVal,
-                  });
-                }}
-              />
-            </InlineField>
+            <>
+              <InlineField label="Service Map" tooltip={'Request and display service map data for trace(s)'}>
+                <InlineSwitch
+                  value={props.query.serviceMap || false}
+                  onChange={(event) => {
+                    const newVal = event.currentTarget.checked;
+                    props.onChange({
+                      ...props.query,
+                      serviceMap: newVal,
+                    });
+                  }}
+                />
+              </InlineField>
+              {!serviceMapSet && (
+                <InlineField
+                  label="Span Limit"
+                  tooltip={'Maximum returned spans. Defaults to 1000, maximum value of 10000'}
+                >
+                  <Input
+                    data-testid="span-limit-input"
+                    placeholder="1000"
+                    defaultValue={props.query.spanLimit}
+                    onBlur={(event) => {
+                      const newVal = event.target.value;
+                      props.onChange({
+                        ...props.query,
+                        spanLimit: newVal,
+                      });
+                    }}
+                  />
+                </InlineField>
+              )}
+            </>
           )}
         </InlineSegmentGroup>
       </QueryEditorRow>

diff --git a/src/types.ts b/src/types.ts
@@ -82,6 +82,7 @@ export interface OpenSearchQuery extends DataQuery {
   format?: PPLFormatType;
   luceneQueryType?: LuceneQueryType;
   serviceMap?: boolean;
+  spanLimit?: string;
 }
 
 export interface OpenSearchAnnotationQuery {