using new version of engine

Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com>
thanos-io · May 15, 2024 · 6f0aac1 · 6f0aac1
1 parent 9f5c5b5
commit 6f0aac1
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 17 deletions.
diff --git a/go.mod b/go.mod
@@ -290,4 +290,4 @@ replace (
  k8s.io/klog/v2 => github.com/simonpasquier/klog-gokit/v3 v3.0.0
 )
 
-replace github.com/thanos-io/promql-engine => github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc
+replace github.com/thanos-io/promql-engine => github.com/pedro-stanaka/promql-engine v0.0.0-20240515101825-a418ffe88a7d
diff --git a/go.sum b/go.sum
@@ -1386,8 +1386,8 @@ github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FI
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y=
 github.com/pborman/uuid v1.2.0/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
-github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc h1:24B4ydZHfM9IeJLn9MZtiieV5yW3GKY3v8i664dlZ3E=
-github.com/pedro-stanaka/promql-engine v0.0.0-20240515092514-c90ad37df1bc/go.mod h1:FEPnabuTql1bDA4OUM41mwcZOJ20R436k8vq+xtGEG0=
+github.com/pedro-stanaka/promql-engine v0.0.0-20240515101825-a418ffe88a7d h1:ppcJWGo/RxQRtNAHbpNSiSZcesNWPPraN76UoWWM0vk=
+github.com/pedro-stanaka/promql-engine v0.0.0-20240515101825-a418ffe88a7d/go.mod h1:FEPnabuTql1bDA4OUM41mwcZOJ20R436k8vq+xtGEG0=
 github.com/performancecopilot/speed v3.0.0+incompatible/go.mod h1:/CLtqpZ5gBg1M9iaPbIdPPGyKcA8hKdoy6hAWba7Yac=
 github.com/phpdave11/gofpdf v1.4.2/go.mod h1:zpO6xFn9yxo3YLyMvW8HcKWVdbNqgIfOOp2dXMnm1mY=
 github.com/phpdave11/gofpdi v1.0.12/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI=

diff --git a/pkg/api/query/grpc.go b/pkg/api/query/grpc.go
@@ -110,16 +110,7 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
  }
  }
 
- stats := &querypb.QueryStats{
- SamplesTotal: 0,
- PeakSamples: 0,
- }
- if explQry, ok := qry.(engine.ExplainableQuery); ok {
- analyze := explQry.Analyze()
- stats.SamplesTotal = analyze.TotalSamples()
- stats.PeakSamples = analyze.PeakSamples()
- }
-
+ stats := extractQueryStats(qry)
  switch vector := result.Value.(type) {
  case promql.Scalar:
  series := &prompb.TimeSeries{
@@ -231,6 +222,7 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
  }
  }
 
+ stats := extractQueryStats(qry)
  switch value := result.Value.(type) {
  case promql.Matrix:
  for _, series := range value {
@@ -240,7 +232,7 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
  Samples: floats,
  Histograms: histograms,
  }
- if err := srv.Send(querypb.NewQueryRangeResponse(series)); err != nil {
+ if err := srv.Send(querypb.NewQueryRangeResponse(series, stats)); err != nil {
  return err
  }
  }
@@ -252,7 +244,7 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
  Samples: floats,
  Histograms: histograms,
  }
- if err := srv.Send(querypb.NewQueryRangeResponse(series)); err != nil {
+ if err := srv.Send(querypb.NewQueryRangeResponse(series, stats)); err != nil {
  return err
  }
  }
@@ -261,12 +253,26 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
  series := &prompb.TimeSeries{
  Samples: []prompb.Sample{{Value: value.V, Timestamp: value.T}},
  }
- return srv.Send(querypb.NewQueryRangeResponse(series))
+ return srv.Send(querypb.NewQueryRangeResponse(series, stats))
  }
 
  return nil
 }
 
+func extractQueryStats(qry promql.Query) *querypb.QueryStats {
+ stats := &querypb.QueryStats{
+ SamplesTotal: 0,
+ PeakSamples: 0,
+ }
+ if explQry, ok := qry.(engine.ExplainableQuery); ok {
+ analyze := explQry.Analyze()
+ stats.SamplesTotal = analyze.TotalSamples()
+ stats.PeakSamples = analyze.PeakSamples()
+ }
+
+ return stats
+}
+
 func (g *GRPCAPI) getRangeQueryForEngine(
  ctx context.Context,
  request *querypb.QueryRangeRequest,

diff --git a/pkg/api/query/querypb/responses.go b/pkg/api/query/querypb/responses.go
@@ -30,11 +30,12 @@ func NewQueryWarningsResponse(errs ...error) *QueryResponse {
  }
 }
 
-func NewQueryRangeResponse(series *prompb.TimeSeries) *QueryRangeResponse {
+func NewQueryRangeResponse(series *prompb.TimeSeries, stats *QueryStats) *QueryRangeResponse {
  return &QueryRangeResponse{
  Result: &QueryRangeResponse_Timeseries{
  Timeseries: series,
  },
+ Stats: stats,
  }
 }