Add multiple engine support to GRPCAPI

Fix build fail for thanos, adds support for multiple engine in GRPCAPI. Signed-off-by: Pradyumna Krishna <git@onpy.in>
thanos-io · Mar 22, 2023 · 63df9a9 · 63df9a9
1 parent ba50f4d
commit 63df9a9
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 48 deletions.
diff --git a/cmd/thanos/query.go b/cmd/thanos/query.go
@@ -802,7 +802,7 @@ func runQuery(
 			info.WithQueryAPIInfoFunc(),
 		)
 
-		grpcAPI := apiv1.NewGRPCAPI(time.Now, queryReplicaLabels, queryableCreator, prometheusEngine, lookbackDeltaCreator, instantDefaultMaxSourceResolution)
+		grpcAPI := apiv1.NewGRPCAPI(time.Now, queryReplicaLabels, queryableCreator, prometheusEngine, thanosEngine, lookbackDeltaCreator, instantDefaultMaxSourceResolution)
 		storeServer := store.NewLimitedStoreServer(store.NewInstrumentedStoreServer(reg, proxy), reg, storeRateLimits)
 		s := grpcserver.New(logger, reg, tracer, grpcLogOpts, tagOpts, comp, grpcProbe,
 			grpcserver.WithServer(apiv1.RegisterQueryServer(grpcAPI)),

diff --git a/pkg/api/query/grpc.go b/pkg/api/query/grpc.go
@@ -24,6 +24,7 @@ type GRPCAPI struct {
 	replicaLabels               []string
 	queryableCreate             query.QueryableCreator
 	queryEngine                 v1.QueryEngine
+	thanosEngine                v1.QueryEngine
 	lookbackDeltaCreate         func(int64) time.Duration
 	defaultMaxResolutionSeconds time.Duration
 }
@@ -33,6 +34,7 @@ func NewGRPCAPI(
 	replicaLabels []string,
 	creator query.QueryableCreator,
 	queryEngine v1.QueryEngine,
+	thanosEngine v1.QueryEngine,
 	lookbackDeltaCreate func(int64) time.Duration,
 	defaultMaxResolutionSeconds time.Duration,
 ) *GRPCAPI {
@@ -41,6 +43,7 @@ func NewGRPCAPI(
 		replicaLabels:               replicaLabels,
 		queryableCreate:             creator,
 		queryEngine:                 queryEngine,
+		thanosEngine:                thanosEngine,
 		lookbackDeltaCreate:         lookbackDeltaCreate,
 		defaultMaxResolutionSeconds: defaultMaxResolutionSeconds,
 	}
@@ -99,7 +102,17 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
 		request.ShardInfo,
 		query.NoopSeriesStatsReporter,
 	)
-	qry, err := g.queryEngine.NewInstantQuery(queryable, &promql.QueryOpts{LookbackDelta: lookbackDelta}, request.Query, ts)
+
+	var engine v1.QueryEngine
+	switch promqlEngineType(request.Engine) {
+	case promqlEnginePrometheus:
+		engine = g.queryEngine
+	case promqlEngineThanos:
+		engine = g.thanosEngine
+	default:
+		engine = g.queryEngine
+	}
+	qry, err := engine.NewInstantQuery(queryable, &promql.QueryOpts{LookbackDelta: lookbackDelta}, request.Query, ts)
 	if err != nil {
 		return err
 	}
@@ -187,7 +200,16 @@ func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Que
 	endTime := time.Unix(request.EndTimeSeconds, 0)
 	interval := time.Duration(request.IntervalSeconds) * time.Second
 
-	qry, err := g.queryEngine.NewRangeQuery(queryable, &promql.QueryOpts{LookbackDelta: lookbackDelta}, request.Query, startTime, endTime, interval)
+	var engine v1.QueryEngine
+	switch promqlEngineType(request.Engine) {
+	case promqlEnginePrometheus:
+		engine = g.queryEngine
+	case promqlEngineThanos:
+		engine = g.thanosEngine
+	default:
+		engine = g.queryEngine
+	}
+	qry, err := engine.NewRangeQuery(queryable, &promql.QueryOpts{LookbackDelta: lookbackDelta}, request.Query, startTime, endTime, interval)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/api/query/querypb/query.pb.go b/pkg/api/query/querypb/query.pb.go
diff --git a/pkg/api/query/querypb/query.proto b/pkg/api/query/querypb/query.proto
@@ -41,6 +41,7 @@ message QueryRequest {
   ShardInfo shard_info = 11;
 
   int64 lookback_delta_seconds = 12;
+  string engine = 13;
 }
 
 message StoreMatchers {
@@ -78,6 +79,7 @@ message QueryRangeRequest {
 
   ShardInfo shard_info = 13;
   int64 lookback_delta_seconds = 14;
+  string engine = 15;
 }
 
 message QueryRangeResponse {