CASSGO-39 Add exec attempt interceptor (round 3)

cluster.go

@@ -258,6 +258,10 @@ type ClusterConfig struct {
 	// See https://issues.apache.org/jira/browse/CASSANDRA-10786
 	DisableSkipMetadata bool
 
+	// ExecAttemptInterceptor will set the provided interceptor on all queries/batches created from this session.
+	// Use it to intercept queries by providing an implementation of ExecAttemptInterceptor.
+	ExecAttemptInterceptor ExecAttemptInterceptor
+
 	// QueryObserver will set the provided query observer on all queries created from this session.
 	// Use it to collect metrics / stats from queries by providing an implementation of QueryObserver.
 	QueryObserver QueryObserver

control.go

@@ -591,7 +591,7 @@ func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter
 				NewLogFieldString("statement", statement), NewLogFieldError("err", iter.err))
 		}
 
-		qry.metrics.attempt(0)
+		qry.metrics.recordAttempt(0, 0, nil)
 		qry.hostMetricsManager.attempt(0, c.getConn().host)
 		if iter.err == nil || !c.retry.Attempt(qry) {
 			break

doc.go

@@ -813,6 +813,19 @@
 //
 // See Example_userDefinedTypesMap, Example_userDefinedTypesStruct, ExampleUDTMarshaler, ExampleUDTUnmarshaler.
 //
+// # Interceptors
+//
+// A ExecAttemptInterceptor wraps query/batch execution and can be used to inject logic that should apply to all query
+// and batch execution attempts. For example, interceptors can be used for rate limiting, logging, attaching
+// distributed tracing metadata to the context, and inspecting query/batch results. However, the query/batch itself
+// cannot be modified.
+//
+// A ExecAttemptInterceptor will be invoked once prior to each query execution attempt, including retry attempts
+// and speculative execution attempts. Interceptors are responsible for calling the provided handler and returning
+// a non-nil Iter or an error.
+//
+// See Example_interceptor for full example.
+//
 // # Metrics and tracing
 //
 // It is possible to provide observer implementations that could be used to gather metrics:

example_interceptor_test.go

@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/*
+ * Content before git sha 34fdeebefcbf183ed7f916f931aa0586fdaa1b40
+ * Copyright (c) 2016, The Gocql authors,
+ * provided under the BSD-3-Clause License.
+ * See the NOTICE file distributed with this work for additional information.
+ */
+
+package gocql_test
+
+import (
+	"context"
+	"log"
+	"time"
+
+	gocql "github.com/apache/cassandra-gocql-driver/v2"
+)
+
+type MyExecAttemptInterceptor struct {
+	injectFault bool
+}
+
+func (q MyExecAttemptInterceptor) Intercept(
+	ctx context.Context,
+	attempt gocql.QueryAttempt,
+	handler gocql.QueryAttemptHandler,
+) (*gocql.Iter, error) {
+	switch attempt.Type {
+	case gocql.OpQuery:
+		// Inspect query
+		log.Println(attempt.Query.Statement())
+	case gocql.OpBatch:
+		// Inspect batch
+		log.Println(attempt.Batch.Entries()[0].Stmt)
+	}
+
+	// Inspect or modify context
+	ctx = context.WithValue(ctx, "trace-id", "123")
+
+	// Optionally bypass the handler and return an error to prevent query execution.
+	// For example, to simulate query timeouts.
+	if q.injectFault && attempt.Attempts == 0 {
+		<-time.After(1 * time.Second)
+		return nil, gocql.RequestErrWriteTimeout{}
+	}
+
+	// The interceptor *must* invoke the handler to execute the query.
+	return handler(ctx)
+}
+
+// Example_interceptor demonstrates how to implement a ExecAttemptInterceptor.
+func Example_interceptor() {
+	cluster := gocql.NewCluster("localhost:9042")
+	cluster.ExecAttemptInterceptor = MyExecAttemptInterceptor{injectFault: true}
+
+	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").
+		RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+		ScanContext(ctx, &stringValue)
+	if err != nil {
+		log.Fatalf("query failed %T", err)
+	}
+}
+
+// Example_interceptor_chain demonstrates how to chain ExecAttemptInterceptors.
+func Example_interceptor_chain() {
+	cluster := gocql.NewCluster("localhost:9042")
+	cluster.ExecAttemptInterceptor = gocql.ExecAttemptInterceptorChain{
+		[]gocql.ExecAttemptInterceptor{
+			MyExecAttemptInterceptor{},
+			MyExecAttemptInterceptor{},
+			MyExecAttemptInterceptor{},
+		},
+	}
+
+	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").
+		RetryPolicy(&gocql.SimpleRetryPolicy{NumRetries: 2}).
+		ScanContext(ctx, &stringValue)
+	if err != nil {
+		log.Fatalf("query failed %T", err)
+	}
+}