Skip to content

fix: use remapped variables in cost calculation #1505

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ysmolski merged 7 commits into from
May 21, 2026
Merged

fix: use remapped variables in cost calculation #1505

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
35c2ca9
fix: use remapped variables in cost calculation
ysmolski May 20, 2026
ca40d86
go lint
ysmolski May 20, 2026
4846ba8
move variables.Set to resolve.VariablesView
ysmolski May 21, 2026
4425b46
Merge branch 'master' of github.com:wundergraph/graphql-go-tools into…
ysmolski May 21, 2026
93102fe
comment a null behavior
ysmolski May 21, 2026
6bdf5cc
remove stale comments
ysmolski May 21, 2026
a59d00e
remap vars in execution tests
ysmolski May 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions execution/engine/execution_engine.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,13 +215,14 @@ func (e *ExecutionEngine) Execute(ctx context.Context, operation *graphql.Reques
if report.HasErrors() {
return report
}
varSet := execContext.resolveContext.VariableSet()
if costCalculator != nil {
costCalculator.ValidateSliceArguments(execContext.resolveContext.Variables, &report)
costCalculator.ValidateSliceArguments(varSet, &report)
if report.HasErrors() {
return report
}
}
operation.ComputeEstimatedCost(costCalculator, execContext.resolveContext.Variables)
operation.ComputeEstimatedCost(costCalculator, varSet)

if execContext.resolveContext.TracingOptions.Enable && !execContext.resolveContext.TracingOptions.ExcludePlannerStats {
planningTime := resolve.GetDurationNanoSinceTraceStart(execContext.resolveContext.Context()) - tracePlanStart
Expand All @@ -240,7 +241,7 @@ func (e *ExecutionEngine) Execute(ctx context.Context, operation *graphql.Reques
return err
}
if resp != nil {
operation.ComputeActualCost(costCalculator, execContext.resolveContext.Variables, execContext.resolveContext.ActualListSizes)
operation.ComputeActualCost(costCalculator, varSet, execContext.resolveContext.ActualListSizes)
}
return nil
case *plan.SubscriptionResponsePlan:
Expand Down
23 changes: 11 additions & 12 deletions execution/graphql/request.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,11 @@ import (
"io"
"net/http"

"github.com/wundergraph/astjson"

"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
"github.com/wundergraph/graphql-go-tools/v2/pkg/astparser"
"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan"
"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/variables"
"github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport"
)

Expand All @@ -23,10 +22,10 @@ const (
type OperationType ast.OperationType

const (
OperationTypeUnknown OperationType = OperationType(ast.OperationTypeUnknown)
OperationTypeQuery OperationType = OperationType(ast.OperationTypeQuery)
OperationTypeMutation OperationType = OperationType(ast.OperationTypeMutation)
OperationTypeSubscription OperationType = OperationType(ast.OperationTypeSubscription)
OperationTypeUnknown = OperationType(ast.OperationTypeUnknown)
OperationTypeQuery = OperationType(ast.OperationTypeQuery)
OperationTypeMutation = OperationType(ast.OperationTypeMutation)
OperationTypeSubscription = OperationType(ast.OperationTypeSubscription)
)

var (
Expand Down Expand Up @@ -196,11 +195,11 @@ func (r *Request) OperationType() (OperationType, error) {
return OperationTypeUnknown, nil
}

func (r *Request) ComputeEstimatedCost(calc *plan.CostCalculator, variables *astjson.Value) {
func (r *Request) ComputeEstimatedCost(calc *plan.CostCalculator, vars variables.Set) {
if calc != nil {
r.estimatedCost = calc.EstimateCost(variables)
r.estimatedCost = calc.EstimateCost(vars)
// Debugging of cost trees. Uncomment to debug.
// fmt.Println(calc.DebugPrint(variables, nil))
// fmt.Println(calc.DebugPrint(vars, nil))
} else {
r.estimatedCost = 0
}
Expand All @@ -210,11 +209,11 @@ func (r *Request) EstimatedCost() int {
return r.estimatedCost
}

func (r *Request) ComputeActualCost(calc *plan.CostCalculator, variables *astjson.Value, actualListSizes map[string]int) {
func (r *Request) ComputeActualCost(calc *plan.CostCalculator, vars variables.Set, actualListSizes map[string]int) {
if calc != nil {
r.actualCost = calc.ActualCost(variables, actualListSizes)
r.actualCost = calc.ActualCost(vars, actualListSizes)
// Debugging of cost trees. Uncomment to debug.
// fmt.Println(calc.DebugPrint(variables, actualListSizes))
// fmt.Println(calc.DebugPrint(vars, actualListSizes))
} else {
r.actualCost = 0
}
Expand Down
70 changes: 34 additions & 36 deletions v2/pkg/engine/plan/cost.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/wundergraph/astjson"

"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/variables"
"github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport"
)

Expand Down Expand Up @@ -89,7 +90,7 @@ type FieldListSize struct {
// multiplier returns the multiplier based on arguments and variables.
// It picks the maximum value among slicing arguments, otherwise it tries to use AssumedSize.
// If neither is available, it falls back to defaultListSize.
func (ls *FieldListSize) multiplier(args map[string]ArgumentInfo, vars *astjson.Value, defaultListSize int) int {
func (ls *FieldListSize) multiplier(args map[string]ArgumentInfo, vars variables.Set, defaultListSize int) int {
multiplier := -1
for _, slicingArg := range ls.SlicingArguments {
value, found := ls.resolveSlicingArg(slicingArg, args, vars)
Expand All @@ -112,7 +113,7 @@ func (ls *FieldListSize) multiplier(args map[string]ArgumentInfo, vars *astjson.
// It falls back to SlicingArgumentDefaults when no value is provided.
// The slicingArg may be a simple argument name or a dot-path into an input object argument.
// An explicitly provided [null] value in variables overrides the default value in schema.
func (ls *FieldListSize) resolveSlicingArg(slicingArg string, args map[string]ArgumentInfo, vars *astjson.Value) (int, bool) {
func (ls *FieldListSize) resolveSlicingArg(slicingArg string, args map[string]ArgumentInfo, vars variables.Set) (int, bool) {
defaultValue, hasDefault := ls.SlicingArgumentDefaults[slicingArg]
if strings.Contains(slicingArg, ".") {
value := extractSlicingArgValue(slicingArg, args, vars)
Expand Down Expand Up @@ -144,10 +145,7 @@ func (ls *FieldListSize) resolveSlicingArg(slicingArg string, args map[string]Ar

// extractSlicingArgValue extracts a value from variables using slicingArg that contains
// a string in the format: "<argumentName>.<inputField1>.<inputField2>..."
func extractSlicingArgValue(slicingArg string, args map[string]ArgumentInfo, vars *astjson.Value) *astjson.Value {
if vars == nil {
return nil
}
func extractSlicingArgValue(slicingArg string, args map[string]ArgumentInfo, vars variables.Set) *astjson.Value {
path := strings.Split(slicingArg, ".")
inputArg := path[0]
arg, found := args[inputArg]
Expand Down Expand Up @@ -282,11 +280,11 @@ type inputObjectField struct {

// inputFieldsCost computes the cost of input object fields from the variable value.
// It handles both single objects and arrays of objects.
func (arg *ArgumentInfo) inputFieldsCost(variables *astjson.Value, weights map[FieldCoordinate]*FieldCost) int {
func (arg *ArgumentInfo) inputFieldsCost(vars variables.Set, weights map[FieldCoordinate]*FieldCost) int {
if !arg.hasVariable {
return 0
}
varValue := variables.Get(arg.varName)
varValue := vars.Get(arg.varName)
if varValue == nil {
return 0
}
Expand Down Expand Up @@ -322,7 +320,7 @@ func (node *CostTreeNode) maxWeightImplementingField(config *DataSourceCostConfi
return maxWeight
}

func (node *CostTreeNode) maxMultiplierImplementingField(config *DataSourceCostConfig, fieldName string, arguments map[string]ArgumentInfo, vars *astjson.Value, defaultListSize int) *FieldListSize {
func (node *CostTreeNode) maxMultiplierImplementingField(config *DataSourceCostConfig, fieldName string, arguments map[string]ArgumentInfo, vars variables.Set, defaultListSize int) *FieldListSize {
var maxMultiplier int
var maxListSize *FieldListSize
for _, implTypeName := range node.implementingTypeNames {
Expand Down Expand Up @@ -403,17 +401,17 @@ func (node *CostTreeNode) maxDirectiveArgumentWeightsImplementingFields(config *
// When it is positive, then its value is used as a fallback value of list sizes for the estimated cost.
// When it is negative, then it computes the actual cost. And it uses the actualListSizes map.
// For actual cost, multipliers are computed as averages (totalCount/parentCount).
func (node *CostTreeNode) cost(configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, defaultListSize int, actualListSizes map[string]int) int {
func (node *CostTreeNode) cost(configs map[DSHash]*DataSourceCostConfig, vars variables.Set, defaultListSize int, actualListSizes map[string]int) int {
if node == nil {
return 0
}

fieldCost, argsCost, directivesCost, multiplier := node.costsAndMultiplier(configs, variables, defaultListSize, actualListSizes)
fieldCost, argsCost, directivesCost, multiplier := node.costsAndMultiplier(configs, vars, defaultListSize, actualListSizes)

// Sum children costs
var childrenCost int
for _, child := range node.children {
childrenCost += child.cost(configs, variables, defaultListSize, actualListSizes)
childrenCost += child.cost(configs, vars, defaultListSize, actualListSizes)
}

// We enforce multiplier=1 for non-list fields.
Expand Down Expand Up @@ -458,7 +456,7 @@ func (node *CostTreeNode) cost(configs map[DSHash]*DataSourceCostConfig, variabl
// When estimating cost, it picks the highest multiplier among different data sources.
// Also, it picks the maximum field weight of implementing types and then
// the maximum among slicing arguments.
func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, defaultListSize int, actualListSizes map[string]int) (fieldCost, argsCost, directivesCost int, multiplier float64) {
func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostConfig, vars variables.Set, defaultListSize int, actualListSizes map[string]int) (fieldCost, argsCost, directivesCost int, multiplier float64) {
if len(node.dataSourceHashes) <= 0 {
// no data source is responsible for this field
return
Expand Down Expand Up @@ -497,7 +495,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
fieldWeight = parent.maxWeightImplementingField(dsCostConfig, node.fieldCoords.FieldName)
// If this field has listSize defined, then do not look into implementing types.
if isEstimation && listSize == nil && node.returnsListType {
listSize = parent.maxMultiplierImplementingField(dsCostConfig, node.fieldCoords.FieldName, node.arguments, variables, defaultListSize)
listSize = parent.maxMultiplierImplementingField(dsCostConfig, node.fieldCoords.FieldName, node.arguments, vars, defaultListSize)
}
}

Expand Down Expand Up @@ -536,7 +534,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
// Input objects always add field-level costs, as the spec says.
// For other types, the explicit argument weight replaces the default type weight.
if arg.isInputObject {
argsCost += arg.inputFieldsCost(variables, dsCostConfig.Weights)
argsCost += arg.inputFieldsCost(vars, dsCostConfig.Weights)
} else if !argumentWeightFound {
if arg.isSimple {
argsCost += dsCostConfig.EnumScalarTypeWeight(arg.typeName)
Expand Down Expand Up @@ -566,7 +564,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
// Pick the maximum multiplier of all data sources.

if listSize != nil {
m := float64(listSize.multiplier(node.arguments, variables, defaultListSize))
m := float64(listSize.multiplier(node.arguments, vars, defaultListSize))
// If this node returns a list of abstract types, then it could have listSize defined.
// Spec allows defining listSize on the fields of interfaces.
if m > multiplier {
Expand All @@ -586,7 +584,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
if sf != node.fieldCoords.FieldName {
continue
}
m := float64(parentLS.multiplier(parent.arguments, variables, defaultListSize))
m := float64(parentLS.multiplier(parent.arguments, vars, defaultListSize))
if m > multiplier {
multiplier = m
}
Expand All @@ -603,7 +601,7 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
dsCostConfig, parent.fieldCoords.FieldName, node.fieldCoords.FieldName,
)
for _, implLS := range implementing {
m := float64(implLS.multiplier(parent.arguments, variables, defaultListSize))
m := float64(implLS.multiplier(parent.arguments, vars, defaultListSize))
if m > multiplier {
multiplier = m
}
Expand Down Expand Up @@ -724,28 +722,28 @@ func NewCostCalculator(config Configuration) *CostCalculator {
}

// EstimateCost returns the calculated total static cost.
// config should be static per process or instance. variables could change between requests.
func (c *CostCalculator) EstimateCost(variables *astjson.Value) int {
return c.tree.cost(c.costConfigs, variables, c.defaultListSize, nil)
// config should be static per process or instance. vars could change between requests.
func (c *CostCalculator) EstimateCost(vars variables.Set) int {
return c.tree.cost(c.costConfigs, vars, c.defaultListSize, nil)
}

const (
actualCostMode = -1 // -1 signals actual mode
)

// ActualCost returns the actual cost of the operation that is based on the actual sizes of lists.
func (c *CostCalculator) ActualCost(variables *astjson.Value, actualListSizes map[string]int) int {
return c.tree.cost(c.costConfigs, variables, actualCostMode, actualListSizes)
func (c *CostCalculator) ActualCost(vars variables.Set, actualListSizes map[string]int) int {
return c.tree.cost(c.costConfigs, vars, actualCostMode, actualListSizes)
}

// ValidateSliceArguments checks that all fields with slicingArguments and
// requireOneSlicingArgument are valid against the arguments passed to those fields.
// Violations are collected as external errors into the report.
func (c *CostCalculator) ValidateSliceArguments(variables *astjson.Value, report *operationreport.Report) {
c.tree.validateSliceArguments(c.costConfigs, variables, report)
func (c *CostCalculator) ValidateSliceArguments(vars variables.Set, report *operationreport.Report) {
c.tree.validateSliceArguments(c.costConfigs, vars, report)
}

func (node *CostTreeNode) validateSliceArguments(configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, report *operationreport.Report) {
func (node *CostTreeNode) validateSliceArguments(configs map[DSHash]*DataSourceCostConfig, vars variables.Set, report *operationreport.Report) {
if node == nil {
return
}
Expand All @@ -771,7 +769,7 @@ func (node *CostTreeNode) validateSliceArguments(configs map[DSHash]*DataSourceC
// The engine has all inlined literals converted to variables at this stage.
// No need to check for literals.
for _, slicingArg := range listSize.SlicingArguments {
if _, found := listSize.resolveSlicingArg(slicingArg, node.arguments, variables); found {
if _, found := listSize.resolveSlicingArg(slicingArg, node.arguments, vars); found {
count++
}
}
Expand All @@ -796,7 +794,7 @@ func (node *CostTreeNode) validateSliceArguments(configs map[DSHash]*DataSourceC
}

for _, child := range node.children {
child.validateSliceArguments(configs, variables, report)
child.validateSliceArguments(configs, vars, report)
}
}

Expand All @@ -818,7 +816,7 @@ func (node *CostTreeNode) buildASTPath() ast.Path {

// DebugPrint prints the cost tree structure for debugging purposes.
// It shows each node's field coordinate, costs, multipliers, and computed totals.
func (c *CostCalculator) DebugPrint(variables *astjson.Value, actualListSizes map[string]int) string {
func (c *CostCalculator) DebugPrint(vars variables.Set, actualListSizes map[string]int) string {
if c.tree == nil || len(c.tree.children) == 0 {
return "<empty cost tree>"
}
Expand All @@ -833,12 +831,12 @@ func (c *CostCalculator) DebugPrint(variables *astjson.Value, actualListSizes ma
sb.WriteString("Estimated Cost Tree Debug\n")
sb.WriteString("=========================\n")
}
c.tree.children[0].debugPrint(&sb, costConfigs, variables, defaultListSize, actualListSizes, 0)
c.tree.children[0].debugPrint(&sb, costConfigs, vars, defaultListSize, actualListSizes, 0)
return sb.String()
}

// debugPrint recursively prints a node and its children with indentation.
func (node *CostTreeNode) debugPrint(sb *strings.Builder, configs map[DSHash]*DataSourceCostConfig, variables *astjson.Value, defaultListSize int, actualListSizes map[string]int, depth int) {
func (node *CostTreeNode) debugPrint(sb *strings.Builder, configs map[DSHash]*DataSourceCostConfig, vars variables.Set, defaultListSize int, actualListSizes map[string]int, depth int) {
// implementation is a bit crude and redundant, we could skip calculating nodes all over again.
// but it should suffice for debugging tests.
if node == nil {
Expand Down Expand Up @@ -872,7 +870,7 @@ func (node *CostTreeNode) debugPrint(sb *strings.Builder, configs map[DSHash]*Da
sb.WriteString("\n")

// Compute costs for this node to display in debug output
fieldCost, argsCost, dirsCost, multiplier := node.costsAndMultiplier(configs, variables, defaultListSize, actualListSizes)
fieldCost, argsCost, dirsCost, multiplier := node.costsAndMultiplier(configs, vars, defaultListSize, actualListSizes)
// We enforce multiplier=1 for non-list fields.
if multiplier == 0 && !node.returnsListType {
multiplier = 1
Expand All @@ -899,12 +897,12 @@ func (node *CostTreeNode) debugPrint(sb *strings.Builder, configs map[DSHash]*Da
var argStrs []string
for name, arg := range node.arguments {
if arg.hasVariable {
if variables == nil {
if vars.IsEmpty() {
// actual cost
argStrs = append(argStrs, fmt.Sprintf("%s=$%s", name, arg.varName))
} else {
// estimated cost
v := variables.Get(arg.varName)
v := vars.Get(arg.varName)
argStrs = append(argStrs, fmt.Sprintf("%s=%s($%s)", name, v, arg.varName))
}
} else {
Expand All @@ -921,10 +919,10 @@ func (node *CostTreeNode) debugPrint(sb *strings.Builder, configs map[DSHash]*Da
// This is somewhat redundant, but it should not be used in production.
// If there is a need to present a cost tree to the user,
// printing should be embedded into the tree calculation process.
subtreeCost := node.cost(configs, variables, defaultListSize, actualListSizes)
subtreeCost := node.cost(configs, vars, defaultListSize, actualListSizes)
fmt.Fprintf(sb, "%s subCost=%d\n", indent, subtreeCost)

for _, child := range node.children {
child.debugPrint(sb, configs, variables, defaultListSize, actualListSizes, depth+1)
child.debugPrint(sb, configs, vars, defaultListSize, actualListSizes, depth+1)
}
}
6 changes: 6 additions & 0 deletions v2/pkg/engine/resolve/context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"github.com/wundergraph/astjson"

"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient"
"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/variables"
)

// Context should not ever be initialized directly, and should be initialized via the NewContext function
Expand All @@ -20,6 +21,7 @@ type Context struct {

// Variables contains the variables to be used to render values of variables for the subgraph.
// Resolver takes into account RemapVariables for variable names.
// Recommented read-only use via variables.Set returned by VariableSet().
Variables *astjson.Value

// RemapVariables contains a map from new names to old names. When variables are renamed,
Expand Down Expand Up @@ -323,6 +325,10 @@ func (c *Context) Free() {
c.ActualListSizes = nil
}

func (c *Context) VariableSet() variables.Set {
return variables.NewSet(c.Variables, c.RemapVariables)
}

type traceStartKey struct{}

type TraceInfo struct {
Expand Down
12 changes: 1 addition & 11 deletions v2/pkg/engine/resolve/inputtemplate.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,20 +133,10 @@ func (i *InputTemplate) renderResolvableObjectVariable(ctx context.Context, obje
}

func (i *InputTemplate) renderContextVariable(ctx *Context, segment TemplateSegment, preparedInput InputTemplateWriter) (variableWasUndefined bool, err error) {
variableSourcePath := segment.VariableSourcePath
if len(variableSourcePath) == 1 && ctx.RemapVariables != nil {
nameToUse, hasMapping := ctx.RemapVariables[variableSourcePath[0]]
if hasMapping && nameToUse != variableSourcePath[0] {
variableSourcePath = []string{nameToUse}
}
}

value := ctx.Variables.Get(variableSourcePath...)
value := ctx.VariableSet().Get(segment.VariableSourcePath...)
if value == nil {
_, _ = preparedInput.Write(literal.NULL)
return true, nil
} else if value.Type() == astjson.TypeNull {
return false, segment.Renderer.RenderVariable(ctx.Context(), value, preparedInput)
}
return false, segment.Renderer.RenderVariable(ctx.Context(), value, preparedInput)
}
Expand Down
Loading
Loading