From 665daa5a5dcfb9e48172f6979af7d1be83a3e888 Mon Sep 17 00:00:00 2001 From: Andrew Hackmann <5140848+bossinc@users.noreply.github.com> Date: Thu, 4 Dec 2025 11:28:38 -0600 Subject: [PATCH] Elasticsearch: Client refactor (#114745) * split up client.go * split up search_request.go * remove double spaces --- .../client/aggregation_builder.go | 224 ++++++++++++ pkg/tsdb/elasticsearch/client/client.go | 309 +---------------- .../elasticsearch/client/http_transport.go | 66 ++++ .../elasticsearch/client/query_builder.go | 112 ++++++ .../elasticsearch/client/request_encoder.go | 53 +++ .../elasticsearch/client/response_parser.go | 262 ++++++++++++++ .../elasticsearch/client/search_request.go | 327 ------------------ 7 files changed, 734 insertions(+), 619 deletions(-) create mode 100644 pkg/tsdb/elasticsearch/client/aggregation_builder.go create mode 100644 pkg/tsdb/elasticsearch/client/http_transport.go create mode 100644 pkg/tsdb/elasticsearch/client/query_builder.go create mode 100644 pkg/tsdb/elasticsearch/client/request_encoder.go create mode 100644 pkg/tsdb/elasticsearch/client/response_parser.go diff --git a/pkg/tsdb/elasticsearch/client/aggregation_builder.go b/pkg/tsdb/elasticsearch/client/aggregation_builder.go new file mode 100644 index 00000000000..043a0d82f4a --- /dev/null +++ b/pkg/tsdb/elasticsearch/client/aggregation_builder.go @@ -0,0 +1,224 @@ +package es + +const ( + // DefaultGeoHashPrecision is the default precision for geohash grid aggregations + DefaultGeoHashPrecision = 3 + // termsOrderTerm is used internally for ordering terms + termsOrderTerm = "_term" +) + +// AggBuilder represents an aggregation builder +type AggBuilder interface { + Histogram(key, field string, fn func(a *HistogramAgg, b AggBuilder)) AggBuilder + DateHistogram(key, field string, fn func(a *DateHistogramAgg, b AggBuilder)) AggBuilder + Terms(key, field string, fn func(a *TermsAggregation, b AggBuilder)) AggBuilder + Nested(key, path string, fn func(a *NestedAggregation, b AggBuilder)) AggBuilder + Filters(key string, fn func(a *FiltersAggregation, b AggBuilder)) AggBuilder + GeoHashGrid(key, field string, fn func(a *GeoHashGridAggregation, b AggBuilder)) AggBuilder + Metric(key, metricType, field string, fn func(a *MetricAggregation)) AggBuilder + Pipeline(key, pipelineType string, bucketPath any, fn func(a *PipelineAggregation)) AggBuilder + Build() (AggArray, error) +} + +type aggBuilderImpl struct { + AggBuilder + aggDefs []*aggDef +} + +func newAggBuilder() *aggBuilderImpl { + return &aggBuilderImpl{ + aggDefs: make([]*aggDef, 0), + } +} + +func (b *aggBuilderImpl) Build() (AggArray, error) { + aggs := make(AggArray, 0) + + for _, aggDef := range b.aggDefs { + agg := &Agg{ + Key: aggDef.key, + Aggregation: aggDef.aggregation, + } + + for _, cb := range aggDef.builders { + childAggs, err := cb.Build() + if err != nil { + return nil, err + } + + agg.Aggregation.Aggs = append(agg.Aggregation.Aggs, childAggs...) + } + + aggs = append(aggs, agg) + } + + return aggs, nil +} + +func (b *aggBuilderImpl) Histogram(key, field string, fn func(a *HistogramAgg, b AggBuilder)) AggBuilder { + innerAgg := &HistogramAgg{ + Field: field, + } + aggDef := newAggDef(key, &aggContainer{ + Type: "histogram", + Aggregation: innerAgg, + }) + + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) DateHistogram(key, field string, fn func(a *DateHistogramAgg, b AggBuilder)) AggBuilder { + innerAgg := &DateHistogramAgg{ + Field: field, + } + aggDef := newAggDef(key, &aggContainer{ + Type: "date_histogram", + Aggregation: innerAgg, + }) + + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) Terms(key, field string, fn func(a *TermsAggregation, b AggBuilder)) AggBuilder { + innerAgg := &TermsAggregation{ + Field: field, + Order: make(map[string]any), + } + aggDef := newAggDef(key, &aggContainer{ + Type: "terms", + Aggregation: innerAgg, + }) + + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + if len(innerAgg.Order) > 0 { + if orderBy, exists := innerAgg.Order[termsOrderTerm]; exists { + innerAgg.Order["_key"] = orderBy + delete(innerAgg.Order, termsOrderTerm) + } + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) Nested(key, field string, fn func(a *NestedAggregation, b AggBuilder)) AggBuilder { + innerAgg := &NestedAggregation{ + Path: field, + } + aggDef := newAggDef(key, &aggContainer{ + Type: "nested", + Aggregation: innerAgg, + }) + + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) Filters(key string, fn func(a *FiltersAggregation, b AggBuilder)) AggBuilder { + innerAgg := &FiltersAggregation{ + Filters: make(map[string]any), + } + aggDef := newAggDef(key, &aggContainer{ + Type: "filters", + Aggregation: innerAgg, + }) + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) GeoHashGrid(key, field string, fn func(a *GeoHashGridAggregation, b AggBuilder)) AggBuilder { + innerAgg := &GeoHashGridAggregation{ + Field: field, + Precision: DefaultGeoHashPrecision, + } + aggDef := newAggDef(key, &aggContainer{ + Type: "geohash_grid", + Aggregation: innerAgg, + }) + + if fn != nil { + builder := newAggBuilder() + aggDef.builders = append(aggDef.builders, builder) + fn(innerAgg, builder) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) Metric(key, metricType, field string, fn func(a *MetricAggregation)) AggBuilder { + innerAgg := &MetricAggregation{ + Type: metricType, + Field: field, + Settings: make(map[string]any), + } + + aggDef := newAggDef(key, &aggContainer{ + Type: metricType, + Aggregation: innerAgg, + }) + + if fn != nil { + fn(innerAgg) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} + +func (b *aggBuilderImpl) Pipeline(key, pipelineType string, bucketPath any, fn func(a *PipelineAggregation)) AggBuilder { + innerAgg := &PipelineAggregation{ + BucketPath: bucketPath, + Settings: make(map[string]any), + } + aggDef := newAggDef(key, &aggContainer{ + Type: pipelineType, + Aggregation: innerAgg, + }) + + if fn != nil { + fn(innerAgg) + } + + b.aggDefs = append(b.aggDefs, aggDef) + + return b +} diff --git a/pkg/tsdb/elasticsearch/client/client.go b/pkg/tsdb/elasticsearch/client/client.go index 9b5cf6dc54b..fbb3e09f092 100644 --- a/pkg/tsdb/elasticsearch/client/client.go +++ b/pkg/tsdb/elasticsearch/client/client.go @@ -1,16 +1,11 @@ package es import ( - "bytes" "context" - "encoding/json" "errors" "fmt" "io" "net/http" - "net/url" - "path" - "strconv" "strings" "time" @@ -73,6 +68,9 @@ var NewClient = func(ctx context.Context, ds *DatasourceInfo, logger log.Logger) ds: ds, configuredFields: ds.ConfiguredFields, indexPattern: ip, + transport: newHTTPTransport(ctx, ds.HTTPClient, ds.URL, logger), + encoder: newRequestEncoder(logger), + parser: newResponseParser(logger), }, nil } @@ -82,6 +80,9 @@ type baseClientImpl struct { configuredFields ConfiguredFields indexPattern IndexPattern logger log.Logger + transport *httpTransport + encoder *requestEncoder + parser *responseParser } func (c *baseClientImpl) GetConfiguredFields() ConfiguredFields { @@ -95,69 +96,11 @@ type multiRequest struct { } func (c *baseClientImpl) executeBatchRequest(uriPath, uriQuery string, requests []*multiRequest) (*http.Response, error) { - bytes, err := c.encodeBatchRequests(requests) + payload, err := c.encoder.encodeBatchRequests(requests) if err != nil { return nil, err } - return c.executeRequest(http.MethodPost, uriPath, uriQuery, bytes) -} - -func (c *baseClientImpl) encodeBatchRequests(requests []*multiRequest) ([]byte, error) { - start := time.Now() - - payload := bytes.Buffer{} - for _, r := range requests { - reqHeader, err := json.Marshal(r.header) - if err != nil { - return nil, err - } - payload.WriteString(string(reqHeader) + "\n") - - reqBody, err := json.Marshal(r.body) - if err != nil { - return nil, err - } - - body := string(reqBody) - body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.interval.Milliseconds(), 10)) - body = strings.ReplaceAll(body, "$__interval", r.interval.String()) - - payload.WriteString(body + "\n") - } - - elapsed := time.Since(start) - c.logger.Debug("Completed encoding of batch requests to json", "duration", elapsed) - - return payload.Bytes(), nil -} - -func (c *baseClientImpl) executeRequest(method, uriPath, uriQuery string, body []byte) (*http.Response, error) { - c.logger.Debug("Sending request to Elasticsearch", "url", c.ds.URL) - u, err := url.Parse(c.ds.URL) - if err != nil { - return nil, backend.DownstreamError(fmt.Errorf("URL could not be parsed: %w", err)) - } - u.Path = path.Join(u.Path, uriPath) - u.RawQuery = uriQuery - - var req *http.Request - if method == http.MethodPost { - req, err = http.NewRequestWithContext(c.ctx, http.MethodPost, u.String(), bytes.NewBuffer(body)) - } else { - req, err = http.NewRequestWithContext(c.ctx, http.MethodGet, u.String(), nil) - } - if err != nil { - return nil, err - } - - req.Header.Set("Content-Type", "application/x-ndjson") - - //nolint:bodyclose - resp, err := c.ds.HTTPClient.Do(req) - if err != nil { - return nil, err - } - return resp, nil + return c.transport.executeBatchRequest(uriPath, uriQuery, payload) } func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error) { @@ -207,7 +150,6 @@ func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearch c.logger.Info("Response received from Elasticsearch", "status", "ok", "statusCode", res.StatusCode, "contentLength", res.ContentLength, "duration", time.Since(start), "stage", StageDatabaseRequest) - start = time.Now() _, resSpan := tracing.DefaultTracer().Start(c.ctx, "datasource.elasticsearch.queryData.executeMultisearch.decodeResponse") defer func() { if err != nil { @@ -217,239 +159,15 @@ func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearch resSpan.End() }() - var msr MultiSearchResponse improvedParsingEnabled := isFeatureEnabled(c.ctx, featuremgmt.FlagElasticsearchImprovedParsing) - if improvedParsingEnabled { - err = StreamMultiSearchResponse(res.Body, &msr) - } else { - dec := json.NewDecoder(res.Body) - err = dec.Decode(&msr) - if err != nil { - // Invalid JSON response from Elasticsearch - err = backend.DownstreamError(err) - } - } + msr, err := c.parser.parseMultiSearchResponse(res.Body, improvedParsingEnabled) if err != nil { - c.logger.Error("Failed to decode response from Elasticsearch", "error", err, "duration", time.Since(start), "improvedParsingEnabled", improvedParsingEnabled) return nil, err } - c.logger.Debug("Completed decoding of response from Elasticsearch", "duration", time.Since(start), "improvedParsingEnabled", improvedParsingEnabled) - msr.Status = res.StatusCode - return &msr, nil -} - -// StreamMultiSearchResponse processes the JSON response in a streaming fashion -func StreamMultiSearchResponse(body io.Reader, msr *MultiSearchResponse) error { - dec := json.NewDecoder(body) - - _, err := dec.Token() // reads the `{` opening brace - if err != nil { - // Invalid JSON response from Elasticsearch - return backend.DownstreamError(err) - } - - for dec.More() { - tok, err := dec.Token() - if err != nil { - return err - } - - if tok == "responses" { - _, err := dec.Token() // reads the `[` opening bracket for responses array - if err != nil { - return err - } - - for dec.More() { - var sr SearchResponse - - _, err := dec.Token() // reads `{` for each SearchResponse - if err != nil { - return err - } - - for dec.More() { - field, err := dec.Token() - if err != nil { - return err - } - - switch field { - case "hits": - sr.Hits = &SearchResponseHits{} - err := processHits(dec, &sr) - if err != nil { - return err - } - case "aggregations": - err := dec.Decode(&sr.Aggregations) - if err != nil { - return err - } - case "error": - err := dec.Decode(&sr.Error) - if err != nil { - return err - } - default: - // skip over unknown fields - err := skipUnknownField(dec) - if err != nil { - return err - } - } - } - - msr.Responses = append(msr.Responses, &sr) - - _, err = dec.Token() // reads `}` closing for each SearchResponse - if err != nil { - return err - } - } - - _, err = dec.Token() // reads the `]` closing bracket for responses array - if err != nil { - return err - } - } else { - err := skipUnknownField(dec) - if err != nil { - return err - } - } - } - - _, err = dec.Token() // reads the `}` closing brace for the entire JSON - return err -} - -// processHits processes the hits in the JSON response incrementally. -func processHits(dec *json.Decoder, sr *SearchResponse) error { - tok, err := dec.Token() // reads the `{` opening brace for the hits object - if err != nil { - return err - } - - if tok != json.Delim('{') { - return fmt.Errorf("expected '{' for hits object, got %v", tok) - } - - for dec.More() { - tok, err := dec.Token() - if err != nil { - return err - } - - switch tok { - case "hits": - if err := streamHitsArray(dec, sr); err != nil { - return err - } - case "total": - var total *SearchResponseHitsTotal - err := dec.Decode(&total) - if err != nil { - // It's possible that the user is using an older version of Elasticsearch (or one that doesn't return what is expected) - // Attempt to parse the total value as an integer in this case - totalInt := 0 - err = dec.Decode(&totalInt) - if err == nil { - total = &SearchResponseHitsTotal{ - Value: totalInt, - } - } else { - // Log the error but do not fail the query - backend.Logger.Debug("failed to decode total hits", "error", err) - } - } - sr.Hits.Total = total - default: - // ignore these fields as they are not used in the current implementation - err := skipUnknownField(dec) - if err != nil { - return err - } - } - } - - // read the closing `}` for the hits object - _, err = dec.Token() - if err != nil { - return err - } - - return nil -} - -// streamHitsArray processes the hits array field incrementally. -func streamHitsArray(dec *json.Decoder, sr *SearchResponse) error { - tok, err := dec.Token() - if err != nil { - return err - } - - // read the opening `[` for the hits array - if tok != json.Delim('[') { - return fmt.Errorf("expected '[' for hits array, got %v", tok) - } - - for dec.More() { - var hit map[string]interface{} - err = dec.Decode(&hit) - if err != nil { - return err - } - - sr.Hits.Hits = append(sr.Hits.Hits, hit) - } - - // read the closing bracket `]` for the hits array - tok, err = dec.Token() - if err != nil { - return err - } - - if tok != json.Delim(']') { - return fmt.Errorf("expected ']' for closing hits array, got %v", tok) - } - - return nil -} - -// skipUnknownField skips over an unknown JSON field's value in the stream. -func skipUnknownField(dec *json.Decoder) error { - tok, err := dec.Token() - if err != nil { - return err - } - - switch tok { - case json.Delim('{'): - // skip everything inside the object until we reach the closing `}` - for dec.More() { - if err := skipUnknownField(dec); err != nil { - return err - } - } - _, err = dec.Token() // read the closing `}` - return err - case json.Delim('['): - // skip everything inside the array until we reach the closing `]` - for dec.More() { - if err := skipUnknownField(dec); err != nil { - return err - } - } - _, err = dec.Token() // read the closing `]` - return err - default: - // no further action needed for primitives - return nil - } + return msr, nil } func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) ([]*multiRequest, error) { @@ -495,3 +213,10 @@ func (c *baseClientImpl) MultiSearch() *MultiSearchRequestBuilder { func isFeatureEnabled(ctx context.Context, feature string) bool { return backend.GrafanaConfigFromContext(ctx).FeatureToggles().IsEnabled(feature) } + +// StreamMultiSearchResponse processes the JSON response in a streaming fashion +// This is a public wrapper for backward compatibility +func StreamMultiSearchResponse(body io.Reader, msr *MultiSearchResponse) error { + parser := newResponseParser(log.NewNullLogger()) + return parser.streamMultiSearchResponse(body, msr) +} diff --git a/pkg/tsdb/elasticsearch/client/http_transport.go b/pkg/tsdb/elasticsearch/client/http_transport.go new file mode 100644 index 00000000000..9ce33fed6c7 --- /dev/null +++ b/pkg/tsdb/elasticsearch/client/http_transport.go @@ -0,0 +1,66 @@ +package es + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/url" + "path" + + "github.com/grafana/grafana-plugin-sdk-go/backend" + "github.com/grafana/grafana-plugin-sdk-go/backend/log" +) + +// httpTransport handles HTTP communication with Elasticsearch +type httpTransport struct { + ctx context.Context + httpClient *http.Client + baseURL string + logger log.Logger +} + +// newHTTPTransport creates a new HTTP transport +func newHTTPTransport(ctx context.Context, httpClient *http.Client, baseURL string, logger log.Logger) *httpTransport { + return &httpTransport{ + ctx: ctx, + httpClient: httpClient, + baseURL: baseURL, + logger: logger, + } +} + +// executeBatchRequest executes a batch request to Elasticsearch +func (t *httpTransport) executeBatchRequest(uriPath, uriQuery string, body []byte) (*http.Response, error) { + return t.executeRequest(http.MethodPost, uriPath, uriQuery, body) +} + +// executeRequest executes an HTTP request to Elasticsearch +func (t *httpTransport) executeRequest(method, uriPath, uriQuery string, body []byte) (*http.Response, error) { + t.logger.Debug("Sending request to Elasticsearch", "url", t.baseURL) + u, err := url.Parse(t.baseURL) + if err != nil { + return nil, backend.DownstreamError(fmt.Errorf("URL could not be parsed: %w", err)) + } + u.Path = path.Join(u.Path, uriPath) + u.RawQuery = uriQuery + + var req *http.Request + if method == http.MethodPost { + req, err = http.NewRequestWithContext(t.ctx, http.MethodPost, u.String(), bytes.NewBuffer(body)) + } else { + req, err = http.NewRequestWithContext(t.ctx, http.MethodGet, u.String(), nil) + } + if err != nil { + return nil, err + } + + req.Header.Set("Content-Type", "application/x-ndjson") + + //nolint:bodyclose + resp, err := t.httpClient.Do(req) + if err != nil { + return nil, err + } + return resp, nil +} diff --git a/pkg/tsdb/elasticsearch/client/query_builder.go b/pkg/tsdb/elasticsearch/client/query_builder.go new file mode 100644 index 00000000000..ed1d16a306d --- /dev/null +++ b/pkg/tsdb/elasticsearch/client/query_builder.go @@ -0,0 +1,112 @@ +package es + +import ( + "strings" +) + +// QueryBuilder represents a query builder +type QueryBuilder struct { + boolQueryBuilder *BoolQueryBuilder +} + +// NewQueryBuilder create a new query builder +func NewQueryBuilder() *QueryBuilder { + return &QueryBuilder{} +} + +// Build builds and return a query builder +func (b *QueryBuilder) Build() (*Query, error) { + q := Query{} + + if b.boolQueryBuilder != nil { + b, err := b.boolQueryBuilder.Build() + if err != nil { + return nil, err + } + q.Bool = b + } + + return &q, nil +} + +// Bool creates and return a query builder +func (b *QueryBuilder) Bool() *BoolQueryBuilder { + if b.boolQueryBuilder == nil { + b.boolQueryBuilder = NewBoolQueryBuilder() + } + return b.boolQueryBuilder +} + +// BoolQueryBuilder represents a bool query builder +type BoolQueryBuilder struct { + filterQueryBuilder *FilterQueryBuilder +} + +// NewBoolQueryBuilder create a new bool query builder +func NewBoolQueryBuilder() *BoolQueryBuilder { + return &BoolQueryBuilder{} +} + +// Filter creates and return a filter query builder +func (b *BoolQueryBuilder) Filter() *FilterQueryBuilder { + if b.filterQueryBuilder == nil { + b.filterQueryBuilder = NewFilterQueryBuilder() + } + return b.filterQueryBuilder +} + +// Build builds and return a bool query builder +func (b *BoolQueryBuilder) Build() (*BoolQuery, error) { + boolQuery := BoolQuery{} + + if b.filterQueryBuilder != nil { + filters, err := b.filterQueryBuilder.Build() + if err != nil { + return nil, err + } + boolQuery.Filters = filters + } + + return &boolQuery, nil +} + +// FilterQueryBuilder represents a filter query builder +type FilterQueryBuilder struct { + filters []Filter +} + +// NewFilterQueryBuilder creates a new filter query builder +func NewFilterQueryBuilder() *FilterQueryBuilder { + return &FilterQueryBuilder{ + filters: make([]Filter, 0), + } +} + +// Build builds and return a filter query builder +func (b *FilterQueryBuilder) Build() ([]Filter, error) { + return b.filters, nil +} + +// AddDateRangeFilter adds a new time range filter +func (b *FilterQueryBuilder) AddDateRangeFilter(timeField string, lte, gte int64, format string) *FilterQueryBuilder { + b.filters = append(b.filters, &RangeFilter{ + Key: timeField, + Lte: lte, + Gte: gte, + Format: format, + }) + return b +} + +// AddQueryStringFilter adds a new query string filter +func (b *FilterQueryBuilder) AddQueryStringFilter(querystring string, analyseWildcard bool) *FilterQueryBuilder { + if len(strings.TrimSpace(querystring)) == 0 { + return b + } + + b.filters = append(b.filters, &QueryStringFilter{ + Query: querystring, + AnalyzeWildcard: analyseWildcard, + }) + return b +} diff --git a/pkg/tsdb/elasticsearch/client/request_encoder.go b/pkg/tsdb/elasticsearch/client/request_encoder.go new file mode 100644 index 00000000000..ae22c8e2694 --- /dev/null +++ b/pkg/tsdb/elasticsearch/client/request_encoder.go @@ -0,0 +1,53 @@ +package es + +import ( + "bytes" + "encoding/json" + "strconv" + "strings" + "time" + + "github.com/grafana/grafana-plugin-sdk-go/backend/log" +) + +// requestEncoder handles encoding of search requests to Elasticsearch format +type requestEncoder struct { + logger log.Logger +} + +// newRequestEncoder creates a new request encoder +func newRequestEncoder(logger log.Logger) *requestEncoder { + return &requestEncoder{ + logger: logger, + } +} + +// encodeBatchRequests encodes multiple requests into NDJSON format +func (e *requestEncoder) encodeBatchRequests(requests []*multiRequest) ([]byte, error) { + start := time.Now() + + payload := bytes.Buffer{} + for _, r := range requests { + reqHeader, err := json.Marshal(r.header) + if err != nil { + return nil, err + } + payload.WriteString(string(reqHeader) + "\n") + + reqBody, err := json.Marshal(r.body) + if err != nil { + return nil, err + } + + body := string(reqBody) + body = strings.ReplaceAll(body, "$__interval_ms", strconv.FormatInt(r.interval.Milliseconds(), 10)) + body = strings.ReplaceAll(body, "$__interval", r.interval.String()) + + payload.WriteString(body + "\n") + } + + elapsed := time.Since(start) + e.logger.Debug("Completed encoding of batch requests to json", "duration", elapsed) + + return payload.Bytes(), nil +} diff --git a/pkg/tsdb/elasticsearch/client/response_parser.go b/pkg/tsdb/elasticsearch/client/response_parser.go new file mode 100644 index 00000000000..a54b238c87a --- /dev/null +++ b/pkg/tsdb/elasticsearch/client/response_parser.go @@ -0,0 +1,262 @@ +package es + +import ( + "encoding/json" + "fmt" + "io" + "time" + + "github.com/grafana/grafana-plugin-sdk-go/backend" + "github.com/grafana/grafana-plugin-sdk-go/backend/log" +) + +// responseParser handles parsing of Elasticsearch responses +type responseParser struct { + logger log.Logger +} + +// newResponseParser creates a new response parser +func newResponseParser(logger log.Logger) *responseParser { + return &responseParser{ + logger: logger, + } +} + +// parseMultiSearchResponse parses a multi-search response using streaming +func (p *responseParser) parseMultiSearchResponse(body io.Reader, improvedParsingEnabled bool) (*MultiSearchResponse, error) { + start := time.Now() + + var msr MultiSearchResponse + var err error + + if improvedParsingEnabled { + err = p.streamMultiSearchResponse(body, &msr) + } else { + dec := json.NewDecoder(body) + err = dec.Decode(&msr) + if err != nil { + // Invalid JSON response from Elasticsearch + err = backend.DownstreamError(err) + } + } + + if err != nil { + p.logger.Error("Failed to decode response from Elasticsearch", "error", err, "duration", time.Since(start), "improvedParsingEnabled", improvedParsingEnabled) + return nil, err + } + + p.logger.Debug("Completed decoding of response from Elasticsearch", "duration", time.Since(start), "improvedParsingEnabled", improvedParsingEnabled) + + return &msr, nil +} + +// streamMultiSearchResponse processes the JSON response in a streaming fashion +func (p *responseParser) streamMultiSearchResponse(body io.Reader, msr *MultiSearchResponse) error { + dec := json.NewDecoder(body) + + _, err := dec.Token() // reads the `{` opening brace + if err != nil { + // Invalid JSON response from Elasticsearch + return backend.DownstreamError(err) + } + + for dec.More() { + tok, err := dec.Token() + if err != nil { + return err + } + + if tok == "responses" { + _, err := dec.Token() // reads the `[` opening bracket for responses array + if err != nil { + return err + } + + for dec.More() { + var sr SearchResponse + + _, err := dec.Token() // reads `{` for each SearchResponse + if err != nil { + return err + } + + for dec.More() { + field, err := dec.Token() + if err != nil { + return err + } + + switch field { + case "hits": + sr.Hits = &SearchResponseHits{} + err := p.processHits(dec, &sr) + if err != nil { + return err + } + case "aggregations": + err := dec.Decode(&sr.Aggregations) + if err != nil { + return err + } + case "error": + err := dec.Decode(&sr.Error) + if err != nil { + return err + } + default: + // skip over unknown fields + err := skipUnknownField(dec) + if err != nil { + return err + } + } + } + + msr.Responses = append(msr.Responses, &sr) + + _, err = dec.Token() // reads `}` closing for each SearchResponse + if err != nil { + return err + } + } + + _, err = dec.Token() // reads the `]` closing bracket for responses array + if err != nil { + return err + } + } else { + err := skipUnknownField(dec) + if err != nil { + return err + } + } + } + + _, err = dec.Token() // reads the `}` closing brace for the entire JSON + return err +} + +// processHits processes the hits in the JSON response incrementally. +func (p *responseParser) processHits(dec *json.Decoder, sr *SearchResponse) error { + tok, err := dec.Token() // reads the `{` opening brace for the hits object + if err != nil { + return err + } + + if tok != json.Delim('{') { + return fmt.Errorf("expected '{' for hits object, got %v", tok) + } + + for dec.More() { + tok, err := dec.Token() + if err != nil { + return err + } + + switch tok { + case "hits": + if err := streamHitsArray(dec, sr); err != nil { + return err + } + case "total": + var total *SearchResponseHitsTotal + err := dec.Decode(&total) + if err != nil { + // It's possible that the user is using an older version of Elasticsearch (or one that doesn't return what is expected) + // Attempt to parse the total value as an integer in this case + totalInt := 0 + err = dec.Decode(&totalInt) + if err == nil { + total = &SearchResponseHitsTotal{ + Value: totalInt, + } + } else { + // Log the error but do not fail the query + backend.Logger.Debug("failed to decode total hits", "error", err) + } + } + sr.Hits.Total = total + default: + // ignore these fields as they are not used in the current implementation + err := skipUnknownField(dec) + if err != nil { + return err + } + } + } + + // read the closing `}` for the hits object + _, err = dec.Token() + if err != nil { + return err + } + + return nil +} + +// streamHitsArray processes the hits array field incrementally. +func streamHitsArray(dec *json.Decoder, sr *SearchResponse) error { + tok, err := dec.Token() + if err != nil { + return err + } + + // read the opening `[` for the hits array + if tok != json.Delim('[') { + return fmt.Errorf("expected '[' for hits array, got %v", tok) + } + + for dec.More() { + var hit map[string]interface{} + err = dec.Decode(&hit) + if err != nil { + return err + } + + sr.Hits.Hits = append(sr.Hits.Hits, hit) + } + + // read the closing bracket `]` for the hits array + tok, err = dec.Token() + if err != nil { + return err + } + + if tok != json.Delim(']') { + return fmt.Errorf("expected ']' for closing hits array, got %v", tok) + } + + return nil +} + +// skipUnknownField skips over an unknown JSON field's value in the stream. +func skipUnknownField(dec *json.Decoder) error { + tok, err := dec.Token() + if err != nil { + return err + } + + switch tok { + case json.Delim('{'): + // skip everything inside the object until we reach the closing `}` + for dec.More() { + if err := skipUnknownField(dec); err != nil { + return err + } + } + _, err = dec.Token() // read the closing `}` + return err + case json.Delim('['): + // skip everything inside the array until we reach the closing `]` + for dec.More() { + if err := skipUnknownField(dec); err != nil { + return err + } + } + _, err = dec.Token() // read the closing `]` + return err + default: + // no further action needed for primitives + return nil + } +} diff --git a/pkg/tsdb/elasticsearch/client/search_request.go b/pkg/tsdb/elasticsearch/client/search_request.go index cb0cf004c11..f898517ab07 100644 --- a/pkg/tsdb/elasticsearch/client/search_request.go +++ b/pkg/tsdb/elasticsearch/client/search_request.go @@ -1,7 +1,6 @@ package es import ( - "strings" "time" "github.com/grafana/grafana-plugin-sdk-go/backend" @@ -11,9 +10,6 @@ const ( HighlightPreTagsString = "@HIGHLIGHT@" HighlightPostTagsString = "@/HIGHLIGHT@" HighlightFragmentSize = 2147483647 - DefaultGeoHashPrecision = 3 - - termsOrderTerm = "_term" ) type SortOrder string @@ -192,326 +188,3 @@ func (m *MultiSearchRequestBuilder) Build() (*MultiSearchRequest, error) { Requests: requests, }, nil } - -// QueryBuilder represents a query builder -type QueryBuilder struct { - boolQueryBuilder *BoolQueryBuilder -} - -// NewQueryBuilder create a new query builder -func NewQueryBuilder() *QueryBuilder { - return &QueryBuilder{} -} - -// Build builds and return a query builder -func (b *QueryBuilder) Build() (*Query, error) { - q := Query{} - - if b.boolQueryBuilder != nil { - b, err := b.boolQueryBuilder.Build() - if err != nil { - return nil, err - } - q.Bool = b - } - - return &q, nil -} - -// Bool creates and return a query builder -func (b *QueryBuilder) Bool() *BoolQueryBuilder { - if b.boolQueryBuilder == nil { - b.boolQueryBuilder = NewBoolQueryBuilder() - } - return b.boolQueryBuilder -} - -// BoolQueryBuilder represents a bool query builder -type BoolQueryBuilder struct { - filterQueryBuilder *FilterQueryBuilder -} - -// NewBoolQueryBuilder create a new bool query builder -func NewBoolQueryBuilder() *BoolQueryBuilder { - return &BoolQueryBuilder{} -} - -// Filter creates and return a filter query builder -func (b *BoolQueryBuilder) Filter() *FilterQueryBuilder { - if b.filterQueryBuilder == nil { - b.filterQueryBuilder = NewFilterQueryBuilder() - } - return b.filterQueryBuilder -} - -// Build builds and return a bool query builder -func (b *BoolQueryBuilder) Build() (*BoolQuery, error) { - boolQuery := BoolQuery{} - - if b.filterQueryBuilder != nil { - filters, err := b.filterQueryBuilder.Build() - if err != nil { - return nil, err - } - boolQuery.Filters = filters - } - - return &boolQuery, nil -} - -// FilterQueryBuilder represents a filter query builder -type FilterQueryBuilder struct { - filters []Filter -} - -// NewFilterQueryBuilder creates a new filter query builder -func NewFilterQueryBuilder() *FilterQueryBuilder { - return &FilterQueryBuilder{ - filters: make([]Filter, 0), - } -} - -// Build builds and return a filter query builder -func (b *FilterQueryBuilder) Build() ([]Filter, error) { - return b.filters, nil -} - -// AddDateRangeFilter adds a new time range filter -func (b *FilterQueryBuilder) AddDateRangeFilter(timeField string, lte, gte int64, format string) *FilterQueryBuilder { - b.filters = append(b.filters, &RangeFilter{ - Key: timeField, - Lte: lte, - Gte: gte, - Format: format, - }) - return b -} - -// AddQueryStringFilter adds a new query string filter -func (b *FilterQueryBuilder) AddQueryStringFilter(querystring string, analyseWildcard bool) *FilterQueryBuilder { - if len(strings.TrimSpace(querystring)) == 0 { - return b - } - - b.filters = append(b.filters, &QueryStringFilter{ - Query: querystring, - AnalyzeWildcard: analyseWildcard, - }) - return b -} - -// AggBuilder represents an aggregation builder -type AggBuilder interface { - Histogram(key, field string, fn func(a *HistogramAgg, b AggBuilder)) AggBuilder - DateHistogram(key, field string, fn func(a *DateHistogramAgg, b AggBuilder)) AggBuilder - Terms(key, field string, fn func(a *TermsAggregation, b AggBuilder)) AggBuilder - Nested(key, path string, fn func(a *NestedAggregation, b AggBuilder)) AggBuilder - Filters(key string, fn func(a *FiltersAggregation, b AggBuilder)) AggBuilder - GeoHashGrid(key, field string, fn func(a *GeoHashGridAggregation, b AggBuilder)) AggBuilder - Metric(key, metricType, field string, fn func(a *MetricAggregation)) AggBuilder - Pipeline(key, pipelineType string, bucketPath any, fn func(a *PipelineAggregation)) AggBuilder - Build() (AggArray, error) -} - -type aggBuilderImpl struct { - AggBuilder - aggDefs []*aggDef -} - -func newAggBuilder() *aggBuilderImpl { - return &aggBuilderImpl{ - aggDefs: make([]*aggDef, 0), - } -} - -func (b *aggBuilderImpl) Build() (AggArray, error) { - aggs := make(AggArray, 0) - - for _, aggDef := range b.aggDefs { - agg := &Agg{ - Key: aggDef.key, - Aggregation: aggDef.aggregation, - } - - for _, cb := range aggDef.builders { - childAggs, err := cb.Build() - if err != nil { - return nil, err - } - - agg.Aggregation.Aggs = append(agg.Aggregation.Aggs, childAggs...) - } - - aggs = append(aggs, agg) - } - - return aggs, nil -} - -func (b *aggBuilderImpl) Histogram(key, field string, fn func(a *HistogramAgg, b AggBuilder)) AggBuilder { - innerAgg := &HistogramAgg{ - Field: field, - } - aggDef := newAggDef(key, &aggContainer{ - Type: "histogram", - Aggregation: innerAgg, - }) - - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) DateHistogram(key, field string, fn func(a *DateHistogramAgg, b AggBuilder)) AggBuilder { - innerAgg := &DateHistogramAgg{ - Field: field, - } - aggDef := newAggDef(key, &aggContainer{ - Type: "date_histogram", - Aggregation: innerAgg, - }) - - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) Terms(key, field string, fn func(a *TermsAggregation, b AggBuilder)) AggBuilder { - innerAgg := &TermsAggregation{ - Field: field, - Order: make(map[string]any), - } - aggDef := newAggDef(key, &aggContainer{ - Type: "terms", - Aggregation: innerAgg, - }) - - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - if len(innerAgg.Order) > 0 { - if orderBy, exists := innerAgg.Order[termsOrderTerm]; exists { - innerAgg.Order["_key"] = orderBy - delete(innerAgg.Order, termsOrderTerm) - } - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) Nested(key, field string, fn func(a *NestedAggregation, b AggBuilder)) AggBuilder { - innerAgg := &NestedAggregation{ - Path: field, - } - aggDef := newAggDef(key, &aggContainer{ - Type: "nested", - Aggregation: innerAgg, - }) - - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) Filters(key string, fn func(a *FiltersAggregation, b AggBuilder)) AggBuilder { - innerAgg := &FiltersAggregation{ - Filters: make(map[string]any), - } - aggDef := newAggDef(key, &aggContainer{ - Type: "filters", - Aggregation: innerAgg, - }) - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) GeoHashGrid(key, field string, fn func(a *GeoHashGridAggregation, b AggBuilder)) AggBuilder { - innerAgg := &GeoHashGridAggregation{ - Field: field, - Precision: DefaultGeoHashPrecision, - } - aggDef := newAggDef(key, &aggContainer{ - Type: "geohash_grid", - Aggregation: innerAgg, - }) - - if fn != nil { - builder := newAggBuilder() - aggDef.builders = append(aggDef.builders, builder) - fn(innerAgg, builder) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) Metric(key, metricType, field string, fn func(a *MetricAggregation)) AggBuilder { - innerAgg := &MetricAggregation{ - Type: metricType, - Field: field, - Settings: make(map[string]any), - } - - aggDef := newAggDef(key, &aggContainer{ - Type: metricType, - Aggregation: innerAgg, - }) - - if fn != nil { - fn(innerAgg) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -} - -func (b *aggBuilderImpl) Pipeline(key, pipelineType string, bucketPath any, fn func(a *PipelineAggregation)) AggBuilder { - innerAgg := &PipelineAggregation{ - BucketPath: bucketPath, - Settings: make(map[string]any), - } - aggDef := newAggDef(key, &aggContainer{ - Type: pipelineType, - Aggregation: innerAgg, - }) - - if fn != nil { - fn(innerAgg) - } - - b.aggDefs = append(b.aggDefs, aggDef) - - return b -}