-
Notifications
You must be signed in to change notification settings - Fork 535
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Allow enabling PromQL experimental functions by tenant (#9798)
* Allow enabling PromQL experimental functions by tenant * Return name of experimental function * Add comments and initial check * Rename middleware * Expand doc message for new config * Don't rerun parsing for the query expression in the experimental functions middleware * Only include the new middleware if experimental functions are enabled globally * Add note * Simplify containsExperimentalFunction * Support configuring enabled experimental functions Signed-off-by: Arve Knudsen <[email protected]> * Fix build and update comment * Add cautionary note * Check experimental functions earlier * Revise according to discussion * Rerun parsing for the query expression in the experimental functions middleware and remove `GetQueryExpr` --------- Signed-off-by: Arve Knudsen <[email protected]> Co-authored-by: Arve Knudsen <[email protected]>
- Loading branch information
Showing
11 changed files
with
278 additions
and
21 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| // SPDX-License-Identifier: AGPL-3.0-only | ||
|
|
||
| package querymiddleware | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "github.com/go-kit/log" | ||
| "github.com/grafana/dskit/tenant" | ||
| "github.com/prometheus/prometheus/promql/parser" | ||
| "golang.org/x/exp/slices" | ||
|
|
||
| apierror "github.com/grafana/mimir/pkg/api/error" | ||
| ) | ||
|
|
||
| const ( | ||
| allExperimentalFunctions = "all" | ||
| ) | ||
|
|
||
| type experimentalFunctionsMiddleware struct { | ||
| next MetricsQueryHandler | ||
| limits Limits | ||
| logger log.Logger | ||
| } | ||
|
|
||
| // newExperimentalFunctionsMiddleware creates a middleware that blocks queries that contain PromQL experimental functions | ||
| // that are not enabled for the active tenant(s), allowing us to enable specific functions only for selected tenants. | ||
| func newExperimentalFunctionsMiddleware(limits Limits, logger log.Logger) MetricsQueryMiddleware { | ||
| return MetricsQueryMiddlewareFunc(func(next MetricsQueryHandler) MetricsQueryHandler { | ||
| return &experimentalFunctionsMiddleware{ | ||
| next: next, | ||
| limits: limits, | ||
| logger: logger, | ||
| } | ||
| }) | ||
| } | ||
|
|
||
| func (m *experimentalFunctionsMiddleware) Do(ctx context.Context, req MetricsQueryRequest) (Response, error) { | ||
| tenantIDs, err := tenant.TenantIDs(ctx) | ||
| if err != nil { | ||
| return nil, apierror.New(apierror.TypeBadData, err.Error()) | ||
| } | ||
|
|
||
| enabledExperimentalFunctions := make(map[string][]string, len(tenantIDs)) | ||
| allExperimentalFunctionsEnabled := true | ||
| for _, tenantID := range tenantIDs { | ||
| enabled := m.limits.EnabledPromQLExperimentalFunctions(tenantID) | ||
| enabledExperimentalFunctions[tenantID] = enabled | ||
| if len(enabled) == 0 || enabled[0] != allExperimentalFunctions { | ||
| allExperimentalFunctionsEnabled = false | ||
| } | ||
| } | ||
|
|
||
| if allExperimentalFunctionsEnabled { | ||
| // If all experimental functions are enabled for all tenants here, we don't need to check the query | ||
| // for those functions and can skip this middleware. | ||
| return m.next.Do(ctx, req) | ||
| } | ||
|
|
||
| expr, err := parser.ParseExpr(req.GetQuery()) | ||
| if err != nil { | ||
| return nil, apierror.New(apierror.TypeBadData, DecorateWithParamName(err, "query").Error()) | ||
| } | ||
| funcs := containedExperimentalFunctions(expr) | ||
| if len(funcs) == 0 { | ||
| // This query does not contain any experimental functions, so we can continue to the next middleware. | ||
| return m.next.Do(ctx, req) | ||
| } | ||
|
|
||
| // Make sure that every used experimental function is enabled for all the tenants here. | ||
| for name := range funcs { | ||
| for tenantID, enabled := range enabledExperimentalFunctions { | ||
| if len(enabled) > 0 && enabled[0] == allExperimentalFunctions { | ||
| // If the first item matches the const value of allExperimentalFunctions, then all experimental | ||
| // functions are enabled for this tenant. | ||
| continue | ||
| } | ||
| if !slices.Contains(enabled, name) { | ||
| err := fmt.Errorf("function %q is not enabled for tenant %s", name, tenantID) | ||
| return nil, apierror.New(apierror.TypeBadData, DecorateWithParamName(err, "query").Error()) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Every used experimental function is enabled for the tenant(s). | ||
| return m.next.Do(ctx, req) | ||
| } | ||
|
|
||
| // containedExperimentalFunctions returns any PromQL experimental functions used in the query. | ||
| func containedExperimentalFunctions(expr parser.Expr) map[string]struct{} { | ||
| expFuncNames := map[string]struct{}{} | ||
| parser.Inspect(expr, func(node parser.Node, _ []parser.Node) error { | ||
| call, ok := node.(*parser.Call) | ||
| if ok { | ||
| if parser.Functions[call.Func.Name].Experimental { | ||
| expFuncNames[call.Func.Name] = struct{}{} | ||
| } | ||
| return nil | ||
| } | ||
| agg, ok := node.(*parser.AggregateExpr) | ||
| if ok { | ||
| // Note that unlike most PromQL functions, the experimental nature of the aggregation functions are manually | ||
| // defined and enforced, so they have to be hardcoded here and updated along with changes in Prometheus. | ||
| switch agg.Op { | ||
| case parser.LIMITK, parser.LIMIT_RATIO: | ||
| expFuncNames[agg.Op.String()] = struct{}{} | ||
| } | ||
| } | ||
| return nil | ||
| }) | ||
| return expFuncNames | ||
| } |
64 changes: 64 additions & 0 deletions
64
pkg/frontend/querymiddleware/experimental_functions_test.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,64 @@ | ||
| // SPDX-License-Identifier: AGPL-3.0-only | ||
|
|
||
| package querymiddleware | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| "github.com/prometheus/prometheus/promql/parser" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestContainedExperimentalFunctions(t *testing.T) { | ||
| t.Cleanup(func() { parser.EnableExperimentalFunctions = false }) | ||
| parser.EnableExperimentalFunctions = true | ||
|
|
||
| testCases := map[string]struct { | ||
| query string | ||
| expect []string | ||
| }{ | ||
| "sum by": { | ||
| query: `sum(up) by (namespace)`, | ||
| }, | ||
| "mad_over_time": { | ||
| query: `mad_over_time(up[5m])`, | ||
| expect: []string{"mad_over_time"}, | ||
| }, | ||
| "mad_over_time with sum and by": { | ||
| query: `sum(mad_over_time(up[5m])) by (namespace)`, | ||
| expect: []string{"mad_over_time"}, | ||
| }, | ||
| "sort_by_label": { | ||
| query: `sort_by_label({__name__=~".+"}, "__name__")`, | ||
| expect: []string{"sort_by_label"}, | ||
| }, | ||
| "sort_by_label_desc": { | ||
| query: `sort_by_label_desc({__name__=~".+"}, "__name__")`, | ||
| expect: []string{"sort_by_label_desc"}, | ||
| }, | ||
| "limitk": { | ||
| query: `limitk by (group) (0, up)`, | ||
| expect: []string{"limitk"}, | ||
| }, | ||
| "limit_ratio": { | ||
| query: `limit_ratio(0.5, up)`, | ||
| expect: []string{"limit_ratio"}, | ||
| }, | ||
| "limit_ratio with mad_over_time": { | ||
| query: `limit_ratio(0.5, mad_over_time(up[5m]))`, | ||
| expect: []string{"limit_ratio", "mad_over_time"}, | ||
| }, | ||
| } | ||
|
|
||
| for name, tc := range testCases { | ||
| t.Run(name, func(t *testing.T) { | ||
| expr, err := parser.ParseExpr(tc.query) | ||
| require.NoError(t, err) | ||
| var enabled []string | ||
| for key := range containedExperimentalFunctions(expr) { | ||
| enabled = append(enabled, key) | ||
| } | ||
| require.ElementsMatch(t, tc.expect, enabled) | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.