engine,execution: make remote endpoints injectable

Signed-off-by: Michael Hoffmann <mhoffmann@cloudflare.com>

Author: MichaHoffmann

engine/distributed.go

@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/thanos-io/promql-engine/api"
-	"github.com/thanos-io/promql-engine/logicalplan"
 
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
@@ -48,66 +47,3 @@ func (l remoteEngine) LabelSets() []labels.Labels {
 func (l remoteEngine) NewRangeQuery(ctx context.Context, opts promql.QueryOpts, plan api.RemoteQuery, start, end time.Time, interval time.Duration) (promql.Query, error) {
 	return l.engine.NewRangeQuery(ctx, l.q, opts, plan.String(), start, end, interval)
 }
-
-type DistributedEngine struct {
-	endpoints    api.RemoteEndpoints
-	remoteEngine *Engine
-}
-
-func NewDistributedEngine(opts Opts, endpoints api.RemoteEndpoints) *DistributedEngine {
-	opts.LogicalOptimizers = []logicalplan.Optimizer{
-		logicalplan.PassthroughOptimizer{Endpoints: endpoints},
-		logicalplan.DistributedExecutionOptimizer{Endpoints: endpoints},
-	}
-
-	return &DistributedEngine{
-		endpoints:    endpoints,
-		remoteEngine: New(opts),
-	}
-}
-
-func (l DistributedEngine) SetQueryLogger(log promql.QueryLogger) {}
-
-func (l DistributedEngine) NewInstantQuery(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, ts time.Time) (promql.Query, error) {
-	return l.MakeInstantQuery(ctx, q, fromPromQLOpts(opts), qs, ts)
-}
-
-func (l DistributedEngine) NewRangeQuery(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, start, end time.Time, interval time.Duration) (promql.Query, error) {
-	return l.MakeRangeQuery(ctx, q, fromPromQLOpts(opts), qs, start, end, interval)
-}
-
-func (l DistributedEngine) MakeInstantQueryFromPlan(ctx context.Context, q storage.Queryable, opts *QueryOpts, plan logicalplan.Node, ts time.Time) (promql.Query, error) {
-	// Truncate milliseconds to avoid mismatch in timestamps between remote and local engines.
-	// Some clients might only support second precision when executing queries.
-	ts = ts.Truncate(time.Second)
-
-	return l.remoteEngine.MakeInstantQueryFromPlan(ctx, q, opts, plan, ts)
-}
-
-func (l DistributedEngine) MakeRangeQueryFromPlan(ctx context.Context, q storage.Queryable, opts *QueryOpts, plan logicalplan.Node, start, end time.Time, interval time.Duration) (promql.Query, error) {
-	// Truncate milliseconds to avoid mismatch in timestamps between remote and local engines.
-	// Some clients might only support second precision when executing queries.
-	start = start.Truncate(time.Second)
-	end = end.Truncate(time.Second)
-	interval = interval.Truncate(time.Second)
-
-	return l.remoteEngine.MakeRangeQueryFromPlan(ctx, q, opts, plan, start, end, interval)
-}
-
-func (l DistributedEngine) MakeInstantQuery(ctx context.Context, q storage.Queryable, opts *QueryOpts, qs string, ts time.Time) (promql.Query, error) {
-	// Truncate milliseconds to avoid mismatch in timestamps between remote and local engines.
-	// Some clients might only support second precision when executing queries.
-	ts = ts.Truncate(time.Second)
-
-	return l.remoteEngine.MakeInstantQuery(ctx, q, opts, qs, ts)
-}
-
-func (l DistributedEngine) MakeRangeQuery(ctx context.Context, q storage.Queryable, opts *QueryOpts, qs string, start, end time.Time, interval time.Duration) (promql.Query, error) {
-	// Truncate milliseconds to avoid mismatch in timestamps between remote and local engines.
-	// Some clients might only support second precision when executing queries.
-	start = start.Truncate(time.Second)
-	end = end.Truncate(time.Second)
-	interval = interval.Truncate(time.Second)
-
-	return l.remoteEngine.MakeRangeQuery(ctx, q, opts, qs, start, end, interval)
-}

engine/distributed_test.go

@@ -299,14 +299,15 @@ func TestDistributedAggregations(t *testing.T) {
 							))
 						}
 
+						endpoints := api.NewStaticEndpoints(remoteEngines)
 						for _, queryOpts := range allQueryOpts {
 							ctx := context.Background()
 							distOpts := localOpts
 							distOpts.DisableFallback = !query.expectFallback
 							for _, instantTS := range instantTSs {
 								t.Run(fmt.Sprintf("instant/ts=%d", instantTS.Unix()), func(t *testing.T) {
-									distEngine := engine.NewDistributedEngine(distOpts, api.NewStaticEndpoints(remoteEngines))
-									distQry, err := distEngine.NewInstantQuery(ctx, completeSeriesSet, queryOpts, query.query, instantTS)
+									distEngine := engine.New(distOpts)
+									distQry, err := distEngine.MakeDistributedInstantQuery(ctx, completeSeriesSet, endpoints, engine.FromPromQLOpts(queryOpts), query.query, instantTS)
 									testutil.Ok(t, err)
 
 									distResult := distQry.Exec(ctx)
@@ -326,8 +327,8 @@ func TestDistributedAggregations(t *testing.T) {
 								if test.rangeEnd == (time.Time{}) {
 									test.rangeEnd = rangeEnd
 								}
-								distEngine := engine.NewDistributedEngine(distOpts, api.NewStaticEndpoints(remoteEngines))
-								distQry, err := distEngine.NewRangeQuery(ctx, completeSeriesSet, queryOpts, query.query, query.rangeStart, test.rangeEnd, rangeStep)
+								distEngine := engine.New(distOpts)
+								distQry, err := distEngine.MakeDistributedRangeQuery(ctx, completeSeriesSet, endpoints, engine.FromPromQLOpts(queryOpts), query.query, query.rangeStart, test.rangeEnd, rangeStep)
 								testutil.Ok(t, err)
 
 								distResult := distQry.Exec(ctx)
@@ -363,58 +364,16 @@ func TestDistributedEngineWarnings(t *testing.T) {
 		},
 	}
 	remote := engine.NewRemoteEngine(opts, querier, math.MinInt64, math.MaxInt64, nil)
-	ng := engine.NewDistributedEngine(opts, api.NewStaticEndpoints([]api.RemoteEngine{remote}))
+	endpoints := api.NewStaticEndpoints([]api.RemoteEngine{remote})
+	ng := engine.New(opts)
 	var (
 		start = time.UnixMilli(0)
 		end   = time.UnixMilli(600)
 		step  = 30 * time.Second
 	)
-	q, err := ng.NewRangeQuery(context.Background(), querier, nil, "test", start, end, step)
+	q, err := ng.MakeDistributedRangeQuery(context.Background(), querier, endpoints, nil, "test", start, end, step)
 	testutil.Ok(t, err)
 
 	res := q.Exec(context.Background())
 	testutil.Equals(t, 1, len(res.Warnings))
 }
-
-func TestDistributedEnginePartialResponses(t *testing.T) {
-	t.Parallel()
-
-	querierErr := &storage.MockQueryable{
-		MockQuerier: &storage.MockQuerier{
-			SelectMockFunction: func(sortSeries bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
-				return newErrorSeriesSet(errors.New("test error"))
-			},
-		},
-	}
-	querierOk := storageWithMockSeries(newMockSeries([]string{labels.MetricName, "foo", "zone", "west"}, []int64{0, 30, 60, 90}, []float64{0, 3, 4, 5}))
-	querierNoop := &storage.MockQueryable{MockQuerier: storage.NoopQuerier()}
-
-	opts := engine.Opts{
-		EnablePartialResponses: true,
-		EngineOpts: promql.EngineOpts{
-			MaxSamples: math.MaxInt64,
-			Timeout:    1 * time.Minute,
-		},
-	}
-
-	remoteErr := engine.NewRemoteEngine(opts, querierErr, math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("zone", "east")})
-	remoteOk := engine.NewRemoteEngine(opts, querierOk, math.MinInt64, math.MaxInt64, []labels.Labels{labels.FromStrings("zone", "west")})
-	ng := engine.NewDistributedEngine(opts, api.NewStaticEndpoints([]api.RemoteEngine{remoteErr, remoteOk}))
-	var (
-		start = time.UnixMilli(0)
-		end   = time.UnixMilli(600 * 1000)
-		step  = 30 * time.Second
-	)
-	q, err := ng.NewRangeQuery(context.Background(), querierNoop, nil, "sum by (zone) (foo)", start, end, step)
-	testutil.Ok(t, err)
-
-	res := q.Exec(context.Background())
-	testutil.Ok(t, res.Err)
-	testutil.Equals(t, 1, len(res.Warnings))
-	testutil.Equals(t, `remote exec error [[{zone="east"}]]: test error`, res.Warnings.AsErrors()[0].Error())
-
-	m, err := res.Matrix()
-	testutil.Ok(t, err)
-	testutil.Equals(t, 1, m.Len())
-	testutil.Equals(t, labels.FromStrings("zone", "west"), m[0].Metric)
-}