Update QueryAttemptHandler signature

To facilitate chaining interceptors

Author: BenEddy

example_interceptor_test.go

@@ -61,7 +61,7 @@ func (q MyQueryAttemptInterceptor) Intercept(
 	}
 
 	// The interceptor *must* invoke the handler to execute the query.
-	return handler(ctx, attempt), nil
+	return handler(ctx, attempt)
 }
 
 // Example_interceptor demonstrates how to implement a QueryAttemptInterceptor.
@@ -86,3 +86,54 @@ func Example_interceptor() {
 		log.Fatalf("query failed %T", err)
 	}
 }
+
+type QueryAttemptInterceptorChain struct {
+	interceptors []gocql.QueryAttemptInterceptor
+}
+
+func (c QueryAttemptInterceptorChain) Intercept(
+	ctx context.Context,
+	attempt gocql.QueryAttempt,
+	handler gocql.QueryAttemptHandler,
+) (*gocql.Iter, error) {
+	return c.interceptors[0].Intercept(ctx, attempt, c.getNextHandler(0, handler))
+}
+
+func (c QueryAttemptInterceptorChain) getNextHandler(curr int, final gocql.QueryAttemptHandler) gocql.QueryAttemptHandler {
+	if curr == len(c.interceptors)-1 {
+		return final
+	}
+
+	return func(ctx context.Context, attempt gocql.QueryAttempt) (*gocql.Iter, error) {
+		return c.interceptors[curr+1].Intercept(ctx, attempt, c.getNextHandler(curr+1, final))
+	}
+}
+
+// Example_interceptor_chain demonstrates how to chain QueryAttemptInterceptors.
+func Example_interceptor_chain() {
+	cluster := gocql.NewCluster("localhost:9042")
+	cluster.QueryAttemptInterceptor = QueryAttemptInterceptorChain{
+		[]gocql.QueryAttemptInterceptor{
+			MyQueryAttemptInterceptor{},
+			MyQueryAttemptInterceptor{},
+			MyQueryAttemptInterceptor{},
+		},
+	}
+
+	session, err := cluster.CreateSession()
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer session.Close()
+
+	ctx := context.Background()
+
+	var stringValue string
+	err = session.Query("select now() from system.local").
+		WithContext(ctx).
+		RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+		Scan(&stringValue)
+	if err != nil {
+		log.Fatalf("query failed %T", err)
+	}
+}

query_executor.go

@@ -65,7 +65,7 @@ type QueryAttempt struct {
 }
 
 // QueryAttemptHandler is a function that attempts query execution.
-type QueryAttemptHandler = func(context.Context, QueryAttempt) *Iter
+type QueryAttemptHandler = func(context.Context, QueryAttempt) (*Iter, error)
 
 // QueryAttemptInterceptor is the interface implemented by query interceptors / middleware.
 //
@@ -93,9 +93,10 @@ func (q *queryExecutor) attemptQuery(ctx context.Context, qry ExecutableQuery, c
 			Host:     conn.host,
 			Attempts: qry.Attempts(),
 		}
-		iter, err = q.interceptor.Intercept(_ctx, attempt, func(_ctx context.Context, attempt QueryAttempt) *Iter {
+		iter, err = q.interceptor.Intercept(_ctx, attempt, func(_ctx context.Context, attempt QueryAttempt) (*Iter, error) {
 			ctx = _ctx
-			return attempt.Query.execute(ctx, attempt.Conn)
+			iter := attempt.Query.execute(ctx, attempt.Conn)
+			return iter, iter.err
 		})
 		if err != nil {
 			iter = &Iter{err: err}