Use graphql.Register func to register cacheControl plugin (#5)

* fix: register extension using graphql.RegisterExtension func

* refactor: creates an specific type for context key and reorganize declarations

---------

Co-authored-by: Raphael Brasil <[email protected]>

Author: raphaabrasil

cache/cache.go

@@ -9,29 +9,35 @@ import (
 	"github.com/vektah/gqlparser/v2/ast"
 )
 
-type Scope string
+type (
+	Scope string
+
+	ctxKey string
+
+	OverallCachePolicy struct {
+		MaxAge float64
+		Scope  Scope
+	}
+
+	CacheControlExtension struct {
+		Version int    `json:"version"`
+		Hints   []Hint `json:"hints"`
+		mu      sync.Mutex
+	}
+
+	Hint struct {
+		Path   ast.Path `json:"path"`
+		MaxAge float64  `json:"maxAge"`
+		Scope  Scope    `json:"scope"`
+	}
+)
 
 const (
 	ScopePublic  = Scope("PUBLIC")
 	ScopePrivate = Scope("PRIVATE")
-)
 
-type Hint struct {
-	Path   ast.Path `json:"path"`
-	MaxAge float64  `json:"maxAge"`
-	Scope  Scope    `json:"scope"`
-}
-
-type OverallCachePolicy struct {
-	MaxAge float64
-	Scope  Scope
-}
-
-type CacheControlExtension struct {
-	Version int    `json:"version"`
-	Hints   []Hint `json:"hints"`
-	mu      sync.Mutex
-}
+	cacheCtxKey ctxKey = "key"
+)
 
 func (cache *CacheControlExtension) AddHint(h Hint) {
 	cache.mu.Lock()
@@ -65,15 +71,14 @@ func (cache *CacheControlExtension) OverallPolicy() OverallCachePolicy {
 	}
 }
 
-const key = "key"
-
 func WithCacheControlExtension(ctx context.Context) context.Context {
 	cache := &CacheControlExtension{Version: 1}
-	return context.WithValue(ctx, key, cache)
+
+	return context.WithValue(ctx, cacheCtxKey, cache)
 }
 
 func CacheControl(ctx context.Context) *CacheControlExtension {
-	c := ctx.Value(key)
+	c := ctx.Value(cacheCtxKey)
 	if c, ok := c.(*CacheControlExtension); ok {
 		return c
 	}
@@ -82,7 +87,7 @@ func CacheControl(ctx context.Context) *CacheControlExtension {
 }
 
 func SetHint(ctx context.Context, scope Scope, maxAge time.Duration) {
-	c := ctx.Value(key)
+	c := ctx.Value(cacheCtxKey)
 	if c, ok := c.(*CacheControlExtension); ok {
 		c.AddHint(Hint{
 			Path:   graphql.GetFieldContext(ctx).Path(),

cache/extension.go

@@ -8,8 +8,10 @@ import (
 
 type Extension struct{}
 
-var _ graphql.HandlerExtension = Extension{}
-var _ graphql.ResponseInterceptor = Extension{}
+var _ interface {
+	graphql.HandlerExtension
+	graphql.ResponseInterceptor
+} = Extension{}
 
 func (c Extension) ExtensionName() string {
 	return "cache"
@@ -20,23 +22,18 @@ func (c Extension) Validate(_ graphql.ExecutableSchema) error {
 }
 
 func (c Extension) InterceptResponse(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
+	if !graphql.HasOperationContext(ctx) {
+		return next(ctx)
+	}
+
 	cache := CacheControl(ctx)
 	if cache == nil {
 		ctx = WithCacheControlExtension(ctx)
-	}
 
-	result := next(ctx)
-
-	if result != nil {
-		cache := CacheControl(ctx)
-
-		if len(cache.Hints) > 0 {
-			if result.Extensions == nil {
-				result.Extensions = make(map[string]interface{})
-			}
-			result.Extensions["cacheControl"] = cache
-		}
+		cache = CacheControl(ctx)
 	}
 
-	return result
+	graphql.RegisterExtension(ctx, "cacheControl", cache)
+
+	return next(ctx)
 }

cache/extension_test.go

@@ -18,59 +18,62 @@ func TestCache(t *testing.T) {
 	})
 
 	t.Run("InterceptResponse", func(t *testing.T) {
-		t.Run("should inject CacheControl in context", func(t *testing.T) {
-			ext := Extension{}
+		query := "query{latestPost {id}}"
 
-			ctx := context.Background()
-			_ = ext.InterceptResponse(ctx, func(ctx context.Context) *graphql.Response {
-				cc := CacheControl(ctx)
-				require.NotNil(t, cc)
-				return &graphql.Response{}
-			})
-		})
+		setupCtx := func(t *testing.T) context.Context {
+			t.Helper()
+
+			ctxWithOperation := graphql.WithOperationContext(
+				context.Background(),
+				&graphql.OperationContext{
+					RawQuery: query,
+				})
 
-		t.Run("should not inject cacheControl extension", func(t *testing.T) {
+			return graphql.WithResponseContext(
+				ctxWithOperation,
+				graphql.DefaultErrorPresenter,
+				graphql.DefaultRecover,
+			)
+		}
+
+		t.Run("Does not try to register cacheControl extension if context is not part of an ongoing operation", func(t *testing.T) {
 			ext := Extension{}
 
-			ctx := context.Background()
+			ctx := context.TODO()
+
 			resp := ext.InterceptResponse(ctx, func(ctx context.Context) *graphql.Response {
 				return &graphql.Response{}
 			})
 
 			require.Nil(t, resp.Extensions["cacheControl"])
 		})
 
-		t.Run("should inject cacheControl extension", func(t *testing.T) {
+		t.Run("Injects CacheControl in context", func(t *testing.T) {
 			ext := Extension{}
 
-			ctx := context.Background()
-			resp := ext.InterceptResponse(ctx, func(ctx context.Context) *graphql.Response {
+			ctx := setupCtx(t)
+
+			_ = ext.InterceptResponse(ctx, func(ctx context.Context) *graphql.Response {
 				cc := CacheControl(ctx)
-				cc.AddHint(Hint{
-					MaxAge: 10,
-					Scope:  ScopePrivate,
-				})
+				require.NotNil(t, cc)
 				return &graphql.Response{}
 			})
-
-			require.NotNil(t, resp.Extensions["cacheControl"])
 		})
 
-		t.Run("should not override extensions", func(t *testing.T) {
+		t.Run("Registers cacheControl extension around the graphql response operation", func(t *testing.T) {
 			ext := Extension{}
 
-			ctx := context.Background()
+			ctx := setupCtx(t)
+
+			var registeredExtensions map[string]interface{}
+
 			resp := ext.InterceptResponse(ctx, func(ctx context.Context) *graphql.Response {
-				return &graphql.Response{
-					Extensions: map[string]interface{}{
-						"foo": "bar",
-					},
-				}
+				registeredExtensions = graphql.GetExtensions(ctx)
+
+				return &graphql.Response{Extensions: registeredExtensions}
 			})
 
-			require.NotNil(t, resp.Extensions["foo"])
-			require.Nil(t, resp.Extensions["cacheControl"])
+			require.NotNil(t, resp.Extensions["cacheControl"])
 		})
 	})
-
 }