From 492b2e71845262e4ad71e32e2ab7ce27e32fa61f Mon Sep 17 00:00:00 2001
From: Daniel Bedrood <danielb@scribd.com>
Date: Wed, 24 Jul 2024 13:38:18 +0200
Subject: [PATCH] WIP

---
 pkg/context/statsig/context.go |  8 ++++----
 pkg/interceptors/statsig.go    | 13 +++++++++----
 pkg/middleware/statsig.go      |  7 +++++--
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/pkg/context/statsig/context.go b/pkg/context/statsig/context.go
index 10011b5..2ee992f 100644
--- a/pkg/context/statsig/context.go
+++ b/pkg/context/statsig/context.go
@@ -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)
 }
 
@@ -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)
 }
diff --git a/pkg/interceptors/statsig.go b/pkg/interceptors/statsig.go
index 2cda7c9..43b00c0 100644
--- a/pkg/interceptors/statsig.go
+++ b/pkg/interceptors/statsig.go
@@ -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
@@ -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)
 	}
 }
@@ -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
@@ -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)
 	}
@@ -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
diff --git a/pkg/middleware/statsig.go b/pkg/middleware/statsig.go
index 6e20a61..cb4709d 100644
--- a/pkg/middleware/statsig.go
+++ b/pkg/middleware/statsig.go
@@ -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.
@@ -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))
 	})
 }
@@ -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))
 	})
 }