From b0b3ded8da573a57a08433d78ac1ed3e6817798c Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Wed, 11 Jun 2025 06:36:46 -0400
Subject: [PATCH] add ContextBatch api to create batch with context (#326)

---
 batchx.go   | 19 +++++++++++++++++--
 qb/batch.go |  2 +-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/batchx.go b/batchx.go
index f9c8bd8..0627b34 100644
--- a/batchx.go
+++ b/batchx.go
@@ -3,7 +3,6 @@ package gocqlx
 import (
 	"context"
 	"fmt"
-
 	"github.com/gocql/gocql"
 )
 
@@ -13,9 +12,25 @@ type Batch struct {
 }
 
 // NewBatch creates a new batch operation using defaults defined in the cluster.
+//
+// Deprecated: use session.Batch instead
 func (s *Session) NewBatch(bt gocql.BatchType) *Batch {
 	return &Batch{
-		Batch: s.Session.NewBatch(bt),
+		Batch: s.Session.Batch(bt),
+	}
+}
+
+// Batch creates a new batch operation using defaults defined in the cluster.
+func (s *Session) Batch(bt gocql.BatchType) *Batch {
+	return &Batch{
+		Batch: s.Session.Batch(bt),
+	}
+}
+
+// ContextBatch creates a new batch operation using defaults defined in the cluster with attached context.
+func (s *Session) ContextBatch(ctx context.Context, bt gocql.BatchType) *Batch {
+	return &Batch{
+		Batch: s.Session.Batch(bt).WithContext(ctx),
 	}
 }
 
diff --git a/qb/batch.go b/qb/batch.go
index cbd4138..10a6b06 100644
--- a/qb/batch.go
+++ b/qb/batch.go
@@ -34,7 +34,7 @@ type BatchBuilder struct {
 //   - gocql.Batch prepares the included statements separately, which is more efficient.
 //     In contrast, gocqlx.qb.BatchBuilder, which is based on gocql.Query, prepares the whole batch statements as one ordinary query.
 //
-// Deprecated: Please use gocql.Session.NewBatch() instead.
+// Deprecated: please use gocql.Session.Batch() instead.
 func Batch() *BatchBuilder {
 	return &BatchBuilder{}
 }