From 6831bd8496223be2b2e1315417932ee6fedd40c1 Mon Sep 17 00:00:00 2001
From: Olfa Karoui <olfaakaroui@gmail.com>
Date: Wed, 22 Oct 2025 16:41:20 +0200
Subject: [PATCH] Add Pre and Post execution hooks on the step executor level

---
 execute.go | 35 +++++++++++++++++++++++++++++++++++
 gateway.go | 19 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/execute.go b/execute.go
index c041af87..97e8d847 100644
--- a/execute.go
+++ b/execute.go
@@ -35,6 +35,9 @@ type queryExecutionResult struct {
 	Err            error
 }
 
+// StepHook is a function that can be executed before or after a step
+type StepHook func(ctx *ExecutionContext) error
+
 // execution is broken up into two phases:
 // - the first walks down the dependency graph execute the network request
 // - the second strips the id fields from the response and  provides a
@@ -48,6 +51,8 @@ type ExecutionContext struct {
 	Variables          map[string]interface{}
 	RequestContext     context.Context
 	RequestMiddlewares []graphql.NetworkMiddleware
+	PreExecutionHook   StepHook
+	PostExecutionHook  StepHook
 }
 
 // Execute returns the result of the query plan
@@ -202,6 +207,21 @@ func executeOneStep(
 	// log the query
 	ctx.logger.QueryPlanStep(step)
 
+	// Execute pre-execution hook if present
+	if ctx.PreExecutionHook != nil {
+		func() {
+			defer func() {
+				if r := recover(); r != nil {
+					ctx.logger.Warn("Pre-execution hook panicked: ", r)
+				}
+			}()
+			if hookErr := ctx.PreExecutionHook(ctx); hookErr != nil {
+				ctx.logger.Warn("Pre-execution hook failed: ", hookErr)
+				// Continue execution even if pre-hook fails, but log the error
+			}
+		}()
+	}
+
 	// the list of variables and their definitions that pertain to this query
 	variables := map[string]interface{}{}
 
@@ -286,6 +306,21 @@ func executeOneStep(
 		queryResult = resultObj
 	}
 
+	// Execute post-execution hook if present
+	if ctx.PostExecutionHook != nil {
+		func() {
+			defer func() {
+				if r := recover(); r != nil {
+					ctx.logger.Warn("Post-execution hook panicked: ", r)
+				}
+			}()
+			if hookErr := ctx.PostExecutionHook(ctx); hookErr != nil {
+				ctx.logger.Warn("Post-execution hook failed: ", hookErr)
+				// Continue execution even if post-hook fails, but log the error
+			}
+		}()
+	}
+	
 	// if there are next steps
 	var dependentSteps []dependentStepArgs
 	if len(step.Then) > 0 {
diff --git a/gateway.go b/gateway.go
index 12ffe022..11b6ac09 100644
--- a/gateway.go
+++ b/gateway.go
@@ -27,7 +27,8 @@ type Gateway struct {
 	queryerFactory     *QueryerFactory
 	queryPlanCache     QueryPlanCache
 	locationPriorities []string
-
+	preExecutionHook   StepHook
+	postExecutionHook  StepHook
 	// group up the list of middlewares at startup to avoid it during execution
 	requestMiddlewares  []graphql.NetworkMiddleware
 	responseMiddlewares []ResponseMiddleware
@@ -86,6 +87,8 @@ func (g *Gateway) Execute(ctx *RequestContext, plans QueryPlanList) (map[string]
 		RequestMiddlewares: g.requestMiddlewares,
 		Plan:               plan,
 		Variables:          ctx.Variables,
+		PreExecutionHook:   g.preExecutionHook,
+		PostExecutionHook:  g.postExecutionHook,
 	}
 
 	// TODO: handle plans of more than one query
@@ -292,6 +295,20 @@ func WithLocationPriorities(priorities []string) Option {
 	}
 }
 
+// WithPreExecutionHook returns an Option that sets a pre-execution hook for all steps
+func WithPreExecutionHook(hook StepHook) Option {
+	return func(g *Gateway) {
+		g.preExecutionHook = hook
+	}
+}
+
+// WithPostExecutionHook returns an Option that sets a post-execution hook for all steps
+func WithPostExecutionHook(hook StepHook) Option {
+	return func(g *Gateway) {
+		g.postExecutionHook = hook
+	}
+}
+
 // WithLogger returns an Option that sets the logger of the gateway
 func WithLogger(l Logger) Option {
 	return func(g *Gateway) {