Add Tracer support to batches (apache#1569)

Implement missing Batch tracing similar to the current Query tracing.

AUTHORS

@@ -122,3 +122,4 @@ Piotr Dulikowski <piodul@scylladb.com>
 Árni Dagur <arni@dagur.eu>
 Tushar Das <tushar.das5@gmail.com>
 Maxim Vladimirskiy <horkhe@gmail.com>
+Bogdan-Ciprian Rusu <bogdanciprian.rusu@crowdstrike.com>

conn.go

@@ -1327,7 +1327,7 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 		b := &req.statements[i]
 
 		if len(entry.Args) > 0 || entry.binding != nil {
-			info, err := c.prepareStatement(batch.Context(), entry.Stmt, nil)
+			info, err := c.prepareStatement(batch.Context(), entry.Stmt, batch.trace)
 			if err != nil {
 				return &Iter{err: err}
 			}
@@ -1369,8 +1369,7 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 		}
 	}
 
-	// TODO: should batch support tracing?
-	framer, err := c.exec(batch.Context(), req, nil)
+	framer, err := c.exec(batch.Context(), req, batch.trace)
 	if err != nil {
 		return &Iter{err: err}
 	}
@@ -1380,6 +1379,10 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 		return &Iter{err: err, framer: framer}
 	}
 
+	if len(framer.traceID) > 0 && batch.trace != nil {
+		batch.trace.Trace(framer.traceID)
+	}
+
 	switch x := resp.(type) {
 	case *resultVoidFrame:
 		return &Iter{}

session.go

@@ -1624,6 +1624,7 @@ type Batch struct {
 	CustomPayload         map[string][]byte
 	rt                    RetryPolicy
 	spec                  SpeculativeExecutionPolicy
+	trace                 Tracer
 	observer              BatchObserver
 	session               *Session
 	serialCons            SerialConsistency
@@ -1653,6 +1654,7 @@ func (s *Session) NewBatch(typ BatchType) *Batch {
 		Type:             typ,
 		rt:               s.cfg.RetryPolicy,
 		serialCons:       s.cfg.SerialConsistency,
+		trace:            s.trace,
 		observer:         s.batchObserver,
 		session:          s,
 		Cons:             s.cons,
@@ -1666,6 +1668,13 @@ func (s *Session) NewBatch(typ BatchType) *Batch {
 	return batch
 }
 
+// Trace enables tracing of this batch. Look at the documentation of the
+// Tracer interface to learn more about tracing.
+func (b *Batch) Trace(trace Tracer) *Batch {
+	b.trace = trace
+	return b
+}
+
 // Observer enables batch-level observer on this batch.
 // The provided observer will be called every time this batched query is executed.
 func (b *Batch) Observer(observer BatchObserver) *Batch {

session_test.go

@@ -236,6 +236,12 @@ func TestBatchBasicAPI(t *testing.T) {
 		t.Fatalf("expected batch.GetConsistency() to return 'One', got '%s'", b.GetConsistency())
 	}
 
+	trace := &traceWriter{}
+	b.Trace(trace)
+	if b.trace != trace {
+		t.Fatalf("expected batch.Trace to be '%v', got '%v'", trace, b.trace)
+	}
+
 	// Test batch.Query()
 	b.Query("test", 1)
 	if b.Entries[0].Stmt != "test" {