diff --git a/query/querytest/compiler.go b/query/querytest/compiler.go index 6d1de8a71d..8820f36389 100644 --- a/query/querytest/compiler.go +++ b/query/querytest/compiler.go @@ -1,8 +1,10 @@ package querytest import ( + "context" + "github.com/influxdata/flux/plan" - "github.com/influxdata/flux/stdlib/influxdata/influxdb/v1" + v1 "github.com/influxdata/flux/stdlib/influxdata/influxdb/v1" "github.com/influxdata/influxdb/v2/query/influxql" "github.com/influxdata/influxdb/v2/query/stdlib/influxdata/influxdb" ) @@ -24,7 +26,7 @@ func (ReplaceFromRule) Pattern() plan.Pattern { return plan.Pat(influxdb.FromKind) } -func (r ReplaceFromRule) Rewrite(n plan.Node) (plan.Node, bool, error) { +func (r ReplaceFromRule) Rewrite(ctx context.Context, n plan.Node) (plan.Node, bool, error) { if err := n.ReplaceSpec(&v1.FromInfluxJSONProcedureSpec{ File: r.Filename, }); err != nil { diff --git a/query/stdlib/influxdata/influxdb/from_test.go b/query/stdlib/influxdata/influxdb/from_test.go index a77ddd6a88..7f5374bccd 100644 --- a/query/stdlib/influxdata/influxdb/from_test.go +++ b/query/stdlib/influxdata/influxdb/from_test.go @@ -1,6 +1,7 @@ package influxdb_test import ( + "context" "fmt" "testing" "time" @@ -180,7 +181,7 @@ func TestFromValidation(t *testing.T) { influxdb.PushDownFilterRule{}, influxdb.PushDownGroupRule{}, )) - _, err := pp.Plan(ps) + _, err := pp.Plan(context.Background(), ps) if err == nil { t.Error("Expected query with no call to range to fail physical planning") } diff --git a/query/stdlib/influxdata/influxdb/rules.go b/query/stdlib/influxdata/influxdb/rules.go index 4102a2a73e..becb5f6948 100644 --- a/query/stdlib/influxdata/influxdb/rules.go +++ b/query/stdlib/influxdata/influxdb/rules.go @@ -1,6 +1,8 @@ package influxdb import ( + "context" + "github.com/influxdata/flux" "github.com/influxdata/flux/ast" "github.com/influxdata/flux/execute" @@ -31,7 +33,7 @@ func (rule PushDownGroupRule) Pattern() plan.Pattern { return plan.Pat(universe.GroupKind, plan.Pat(ReadRangePhysKind)) } -func (rule PushDownGroupRule) Rewrite(node plan.Node) (plan.Node, bool, error) { +func (rule PushDownGroupRule) Rewrite(ctx context.Context, node plan.Node) (plan.Node, bool, error) { src := node.Predecessors()[0].ProcedureSpec().(*ReadRangePhysSpec) grp := node.ProcedureSpec().(*universe.GroupProcedureSpec) @@ -71,7 +73,7 @@ func (rule PushDownRangeRule) Pattern() plan.Pattern { } // Rewrite converts 'from |> range' into 'ReadRange' -func (rule PushDownRangeRule) Rewrite(node plan.Node) (plan.Node, bool, error) { +func (rule PushDownRangeRule) Rewrite(ctx context.Context, node plan.Node) (plan.Node, bool, error) { fromNode := node.Predecessors()[0] fromSpec := fromNode.ProcedureSpec().(*FromProcedureSpec) @@ -96,7 +98,7 @@ func (PushDownFilterRule) Pattern() plan.Pattern { return plan.Pat(universe.FilterKind, plan.Pat(ReadRangePhysKind)) } -func (PushDownFilterRule) Rewrite(pn plan.Node) (plan.Node, bool, error) { +func (PushDownFilterRule) Rewrite(ctx context.Context, pn plan.Node) (plan.Node, bool, error) { filterSpec := pn.ProcedureSpec().(*universe.FilterProcedureSpec) fromNode := pn.Predecessors()[0] fromSpec := fromNode.ProcedureSpec().(*ReadRangePhysSpec) @@ -183,7 +185,7 @@ func (rule PushDownReadTagKeysRule) Pattern() plan.Pattern { plan.Pat(ReadRangePhysKind)))) } -func (rule PushDownReadTagKeysRule) Rewrite(pn plan.Node) (plan.Node, bool, error) { +func (rule PushDownReadTagKeysRule) Rewrite(ctx context.Context, pn plan.Node) (plan.Node, bool, error) { // Retrieve the nodes and specs for all of the predecessors. distinctSpec := pn.ProcedureSpec().(*universe.DistinctProcedureSpec) keepNode := pn.Predecessors()[0] @@ -245,7 +247,7 @@ func (rule PushDownReadTagValuesRule) Pattern() plan.Pattern { plan.Pat(ReadRangePhysKind)))) } -func (rule PushDownReadTagValuesRule) Rewrite(pn plan.Node) (plan.Node, bool, error) { +func (rule PushDownReadTagValuesRule) Rewrite(ctx context.Context, pn plan.Node) (plan.Node, bool, error) { // Retrieve the nodes and specs for all of the predecessors. distinctNode := pn distinctSpec := distinctNode.ProcedureSpec().(*universe.DistinctProcedureSpec) @@ -556,7 +558,7 @@ func (SortedPivotRule) Pattern() plan.Pattern { return plan.Pat(universe.PivotKind, plan.Pat(ReadRangePhysKind)) } -func (SortedPivotRule) Rewrite(pn plan.Node) (plan.Node, bool, error) { +func (SortedPivotRule) Rewrite(ctx context.Context, pn plan.Node) (plan.Node, bool, error) { pivotSpec := pn.ProcedureSpec().Copy().(*universe.PivotProcedureSpec) pivotSpec.IsSortedByFunc = func(cols []string, desc bool) bool { if desc {