feat: expose the static cost to the middleware and configs #2470

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.

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

Draft

ysmolski wants to merge 12 commits into main from yury/eng-8636-router-expose-the-static-cost-to-the-middleware

router-tests/go.mod

-Original file line number
+Diff line change
@@ Expand Up / @@ -26,7 +26,7 @@ require ( @@
     	github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects v0.0.0-20250715110703-10f2e5f9c79e
     	github.com/wundergraph/cosmo/router v0.0.0-20251125205644-175f80c4e6d9
     	github.com/wundergraph/cosmo/router-plugin v0.0.0-20250808194725-de123ba1c65e
-    	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.245
+    	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.246
     	go.opentelemetry.io/otel v1.36.0
     	go.opentelemetry.io/otel/sdk v1.36.0
     	go.opentelemetry.io/otel/sdk/metric v1.36.0
@@ Expand Down @@

router-tests/go.sum

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -352,8 +352,8 @@ github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/
  
    github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=

    github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 h1:8/D7f8gKxTBjW+SZK4mhxTTBVpxcqeBgWF1Rfmltbfk=

    github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083/go.mod h1:eOTL6acwctsN4F3b7YE+eE2t8zcJ/doLm9sZzsxxxrE=

    github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.245 h1:MYewlXgIhI9jusocPUeyo346J3M5cqzc6ddru1qp+S8=

    github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.245/go.mod h1:mX25ASEQiKamxaFSK6NZihh0oDCigIuzro30up4mFH4=

    github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.246 h1:B4/E6zJ5PbHNqcNw4Bvki7o+alcgyz8L5+Nlq1FZbLM=

    github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.246/go.mod h1:mX25ASEQiKamxaFSK6NZihh0oDCigIuzro30up4mFH4=

    github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 h1:FnBeRrxr7OU4VvAzt5X7s6266i6cSVkkFPS0TuXWbIg=

    github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=

    github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=

router-tests/modules/custom_trace_propagator_test.go

-Original file line number
+Diff line change
@@ -1,4 +1,4 @@
-    package module
+    package module_test
     import (
     	"encoding/json"
@@ Expand Down @@

router-tests/modules/streams_hooks_combined_test.go

-Original file line number
+Diff line change
@@ -1,4 +1,4 @@
-    package module
+    package module_test
     import (
     	"encoding/json"
@@ Expand All / @@ -8,6 +8,8 @@ import ( @@
     	"github.com/hasura/go-graphql-client"
     	"github.com/stretchr/testify/assert"
     	"github.com/stretchr/testify/require"
+    	"go.uber.org/zap/zapcore"
     	"github.com/wundergraph/cosmo/router-tests/events"
     	stream_publish "github.com/wundergraph/cosmo/router-tests/modules/stream-publish"
     	stream_receive "github.com/wundergraph/cosmo/router-tests/modules/stream-receive"
@@ Expand All / @@ -16,7 +18,6 @@ import ( @@
     	"github.com/wundergraph/cosmo/router/pkg/config"
     	"github.com/wundergraph/cosmo/router/pkg/pubsub/datasource"
     	"github.com/wundergraph/cosmo/router/pkg/pubsub/kafka"
-    	"go.uber.org/zap/zapcore"
     )
     func TestStreamsHooksCombined(t *testing.T) {
@@ Expand Down @@

router-tests/modules/verify-cost-analysis/module.go

-Original file line number
+Diff line change
@@ -0,0 +1,63 @@
+    package verify_cost_analysis
+    import (
+    	"net/http"
+    	"go.uber.org/zap"
+    	"github.com/wundergraph/cosmo/router/core"
+    )
+    const myModuleID = "verifyCost"
+    // CapturedCost holds the captured cost from operation context
+    type CapturedCost struct {
+    	Cost  core.OperationCost
+    	Error error
+    }
+    // VerifyCostModule captures cost for verification in tests
+    type VerifyCostModule struct {
+    	ResultsChan chan CapturedCost
+    	Logger      *zap.Logger
+    }
+    func (m *VerifyCostModule) Provision(ctx *core.ModuleContext) error {
+    	m.Logger = ctx.Logger
+    	if m.ResultsChan == nil {
+    		m.ResultsChan = make(chan CapturedCost, 1)
+    	}
+    	return nil
+    }
+    func (m *VerifyCostModule) Middleware(ctx core.RequestContext, next http.Handler) {
+    	operation := ctx.Operation()
+    	cost, err := operation.Cost()
+    	captured := CapturedCost{Cost: cost, Error: err}
+    	// Send the captured values to the test
+    	select {
+    	case m.ResultsChan <- captured:
+    	default:
+    		// Channel is full, skip
+    	}
+    	next.ServeHTTP(ctx.ResponseWriter(), ctx.Request())
+    }
+    func (m *VerifyCostModule) Module() core.ModuleInfo {
+    	return core.ModuleInfo{
+    		ID:       myModuleID,
+    		Priority: 1,
+    		New: func() core.Module {
+    			return &VerifyCostModule{
+    				ResultsChan: make(chan CapturedCost, 1),
+    			}
+    		},
+    	}
+    }
+    var (
+    	_ core.RouterMiddlewareHandler = (*VerifyCostModule)(nil)
+    )

router-tests/modules/verify_cost_analysis_test.go

-Original file line number
+Diff line change
@@ -0,0 +1,103 @@
+    package module_test
+    import (
+    	"testing"
+    	"time"
+    	"github.com/stretchr/testify/assert"
+    	"github.com/stretchr/testify/require"
+    	verifyModule "github.com/wundergraph/cosmo/router-tests/modules/verify-cost-analysis"
+    	"github.com/wundergraph/cosmo/router-tests/testenv"
+    	"github.com/wundergraph/cosmo/router/core"
+    	"github.com/wundergraph/cosmo/router/pkg/config"
+    )
+    func TestCostModuleExposition(t *testing.T) {
+    	t.Parallel()
+    	t.Run("module can access cost when cost analysis is enabled", func(t *testing.T) {
+    		t.Parallel()
+    		resultsChan := make(chan verifyModule.CapturedCost, 1)
+    		cfg := config.Config{
+    			Graph: config.Graph{},
+    			Modules: map[string]any{
+    				"verifyCost": verifyModule.VerifyCostModule{
+    					ResultsChan: resultsChan,
+    				},
+    			},
+    		}
+    		testenv.Run(t, &testenv.Config{
+    			RouterOptions: []core.Option{
+    				core.WithModulesConfig(cfg.Modules),
+    				core.WithCustomModules(&verifyModule.VerifyCostModule{}),
+    			},
+    			ModifySecurityConfiguration: func(securityConfiguration *config.SecurityConfiguration) {
+    				securityConfiguration.CostAnalysis = &config.CostAnalysis{
+    					Enabled:           true,
+    					Mode:              config.CostAnalysisModeMeasure,
+    					EstimatedListSize: 10,
+    				}
+    			},
+    		}, func(t *testing.T, xEnv *testenv.Environment) {
+    			res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+    				Query: `query GetEmployee { employee(id: 1) { id } }`,
+    			})
+    			require.NoError(t, err)
+    			assert.Equal(t, 200, res.Response.StatusCode)
+    			testenv.AwaitChannelWithT(t, 10*time.Second, resultsChan, func(t *testing.T, captured verifyModule.CapturedCost) {
+    				assert.NoError(t, captured.Error, "Cost() should not return an error when cost analysis is enabled")
+    				assert.Greater(t, captured.Cost.Estimated, 0, "Estimated cost should be greater than 0 for a query with object fields")
+    				// Log the cost for demonstration purposes
+    				t.Logf("Query estimated cost: %d (could be used for rate limiting)", captured.Cost.Estimated)
+    				// In a real module, you could use this cost for:
+    				// - Rate limiting (deduct from user's cost budget)
+    				// - Logging/metrics
+    				// - Custom rejection logic
+    				// - Billing/monetization
+    			})
+    		})
+    	})
+    	t.Run("module receives error when cost analysis is disabled", func(t *testing.T) {
+    		t.Parallel()
+    		resultsChan := make(chan verifyModule.CapturedCost, 1)
+    		cfg := config.Config{
+    			Graph: config.Graph{},
+    			Modules: map[string]any{
+    				"verifyCost": verifyModule.VerifyCostModule{
+    					ResultsChan: resultsChan,
+    				},
+    			},
+    		}
+    		testenv.Run(t, &testenv.Config{
+    			RouterOptions: []core.Option{
+    				core.WithModulesConfig(cfg.Modules),
+    				core.WithCustomModules(&verifyModule.VerifyCostModule{}),
+    			},
+    			ModifySecurityConfiguration: func(securityConfiguration *config.SecurityConfiguration) {
+    				securityConfiguration.CostAnalysis = nil
+    			},
+    		}, func(t *testing.T, xEnv *testenv.Environment) {
+    			res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+    				Query: `query GetEmployee { employee(id: 1) { id } }`,
+    			})
+    			require.NoError(t, err)
+    			assert.Equal(t, 200, res.Response.StatusCode)
+    			testenv.AwaitChannelWithT(t, 10*time.Second, resultsChan, func(t *testing.T, captured verifyModule.CapturedCost) {
+    				assert.Error(t, captured.Error, "Cost() should return an error when cost analysis is disabled")
+    				assert.Equal(t, 0, captured.Cost.Estimated, "Estimated cost should be 0 when cost analysis is disabled")
+    			})
+    		})
+    	})
+    }

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: expose the static cost to the middleware and configs #2470

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

feat: expose the static cost to the middleware and configs #2470

Are you sure you want to change the base?

Uh oh!

feat: expose the static cost to the middleware and configs #2470

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!