Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
laynax committed Jul 24, 2024
1 parent 1b55b44 commit 492b2e7
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 10 deletions.
8 changes: 4 additions & 4 deletions pkg/context/statsig/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ func ExtractExperiment(ctx context.Context) (statsig.GetExperimentFunc, error) {

// ExperimentToContext adds the sdkstatsig.GetExperimentFunc to the context for extraction later.
// Returning the new context that has been created.
func ExperimentToContext(ctx context.Context, experiment string) context.Context {
r := &ctxStatsigExperimentMarker{fn: statsig.GetExperiment(experiment)}
func ExperimentToContext(ctx context.Context, x statsig.GetExperimentFunc) context.Context {
r := &ctxStatsigExperimentMarker{fn: x}
return context.WithValue(ctx, ctxStatsigExperminetKey, r)
}

Expand All @@ -48,7 +48,7 @@ func ExtractFeatureFlag(ctx context.Context) (statsig.GetFeatureFlagFunc, error)

// FeatureFlagToContext adds the sdkstatsig.GetFeatureFlagFunc to the context for extraction later.
// Returning the new context that has been created.
func FeatureFlagToContext(ctx context.Context, featureFlag string) context.Context {
r := &ctxStatsigFeatureFlagMarker{fn: statsig.GetFeatureFlag(featureFlag)}
func FeatureFlagToContext(ctx context.Context, f statsig.GetFeatureFlagFunc) context.Context {
r := &ctxStatsigFeatureFlagMarker{fn: f}
return context.WithValue(ctx, ctxStatsigFeatureFlagKey, r)
}
13 changes: 9 additions & 4 deletions pkg/interceptors/statsig.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"google.golang.org/grpc"

sdkstatsigcontext "github.com/scribd/go-sdk/pkg/context/statsig"
"github.com/scribd/go-sdk/pkg/statsig"
)

// ExperimentUnaryServerInterceptor returns a unary server interceptor
Expand All @@ -18,7 +19,8 @@ func ExperimentUnaryServerInterceptor(experiment string) grpc.UnaryServerInterce
info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler,
) (interface{}, error) {
newCtx := sdkstatsigcontext.ExperimentToContext(ctx, experiment)
x := statsig.GetExperiment(experiment)
newCtx := sdkstatsigcontext.ExperimentToContext(ctx, x)
return handler(newCtx, req)
}
}
Expand All @@ -32,7 +34,8 @@ func ExperimentStreamServerInterceptor(experiment string) grpc.StreamServerInter
info *grpc.StreamServerInfo,
handler grpc.StreamHandler,
) error {
newCtx := sdkstatsigcontext.ExperimentToContext(stream.Context(), experiment)
x := statsig.GetExperiment(experiment)
newCtx := sdkstatsigcontext.ExperimentToContext(stream.Context(), x)

wrapped := grpcmiddleware.WrapServerStream(stream)
wrapped.WrappedContext = newCtx
Expand All @@ -50,7 +53,8 @@ func FeatureFlagUnaryServerInterceptor(featureFlag string) grpc.UnaryServerInter
info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler,
) (interface{}, error) {
newCtx := sdkstatsigcontext.FeatureFlagToContext(ctx, featureFlag)
ff := statsig.GetFeatureFlag(featureFlag)
newCtx := sdkstatsigcontext.FeatureFlagToContext(ctx, ff)

return handler(newCtx, req)
}
Expand All @@ -65,7 +69,8 @@ func FeatureFlagStreamServerInterceptor(featureFlag string) grpc.StreamServerInt
info *grpc.StreamServerInfo,
handler grpc.StreamHandler,
) error {
newCtx := sdkstatsigcontext.FeatureFlagToContext(stream.Context(), featureFlag)
ff := statsig.GetFeatureFlag(featureFlag)
newCtx := sdkstatsigcontext.FeatureFlagToContext(stream.Context(), ff)

wrapped := grpcmiddleware.WrapServerStream(stream)
wrapped.WrappedContext = newCtx
Expand Down
7 changes: 5 additions & 2 deletions pkg/middleware/statsig.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"net/http"

sdkstatsigcontext "github.com/scribd/go-sdk/pkg/context/statsig"
"github.com/scribd/go-sdk/pkg/statsig"
)

// StatsigExperimentMiddleware holds the experiment name to be injected in the request context.
Expand All @@ -20,7 +21,8 @@ func NewStatsigExperimentMiddleware(experiment string) StatsigExperimentMiddlewa
// http.Handler to be mounted as middleware. The handler injects the experiment key to the request context.
func (sm StatsigExperimentMiddleware) Handler(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := sdkstatsigcontext.ExperimentToContext(r.Context(), sm.experiment)
x := statsig.GetExperiment(sm.experiment)
ctx := sdkstatsigcontext.ExperimentToContext(r.Context(), x)
next.ServeHTTP(w, r.WithContext(ctx))
})
}
Expand All @@ -37,7 +39,8 @@ func NewStatsigFeatureFlagMiddleware(featureFlag string) StatsigFeatureFlagMiddl
// http.Handler to be mounted as middleware. The handler injects the feature flag key to the request context.
func (sm StatsigFeatureFlagMiddleware) Handler(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := sdkstatsigcontext.FeatureFlagToContext(r.Context(), sm.featureFlag)
ff := statsig.GetFeatureFlag(sm.featureFlag)
ctx := sdkstatsigcontext.FeatureFlagToContext(r.Context(), ff)
next.ServeHTTP(w, r.WithContext(ctx))
})
}

0 comments on commit 492b2e7

Please sign in to comment.