From 5130edde3e55c9f2042f01b3cee4ff10c98e1b27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Thu, 14 Nov 2024 12:42:08 +0100
Subject: [PATCH 1/2] feat: add method for getting batched statements

Adds a GetBatchedStatements function that returns the currently
buffered statements that will be executed if RunBatch is called.
This gives applications the possibility of inspecting the list of
buffered statements before actually executing them.
---
 driver.go      | 13 +++++++++++++
 driver_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/driver.go b/driver.go
index 647dfc20..c42869be 100644
--- a/driver.go
+++ b/driver.go
@@ -535,6 +535,10 @@ type SpannerConn interface {
 	InDDLBatch() bool
 	// InDMLBatch returns true if the connection is currently in a DML batch.
 	InDMLBatch() bool
+	// GetBatchedStatements returns a copy of the statements that are currently
+	// buffered to be executed as a DML or DDL batch. It returns an empty slice
+	// if no batch is active, or if there are no statements buffered.
+	GetBatchedStatements() []spanner.Statement
 
 	// RetryAbortsInternally returns true if the connection automatically
 	// retries all aborted transactions.
@@ -755,6 +759,15 @@ func (c *conn) InDMLBatch() bool {
 	return (c.batch != nil && c.batch.tp == dml) || (c.inReadWriteTransaction() && c.tx.(*readWriteTransaction).batch != nil)
 }
 
+func (c *conn) GetBatchedStatements() []spanner.Statement {
+	if c.batch == nil {
+		return []spanner.Statement{}
+	}
+	res := make([]spanner.Statement, len(c.batch.statements))
+	copy(res, c.batch.statements)
+	return res
+}
+
 func (c *conn) inBatch() bool {
 	return c.InDDLBatch() || c.InDMLBatch()
 }
diff --git a/driver_test.go b/driver_test.go
index 9b7d328b..869ff0f5 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -17,6 +17,7 @@ package spannerdriver
 import (
 	"context"
 	"database/sql/driver"
+	"reflect"
 	"testing"
 	"time"
 
@@ -424,6 +425,52 @@ func TestConn_NonDmlStatementsInDmlBatch(t *testing.T) {
 	}
 }
 
+func TestConn_GetBatchedStatements(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	c := &conn{}
+	if !reflect.DeepEqual(c.GetBatchedStatements(), []spanner.Statement{}) {
+		t.Fatal("conn should return an empty slice when no batch is active")
+	}
+	if err := c.StartBatchDDL(); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(c.GetBatchedStatements(), []spanner.Statement{}) {
+		t.Fatal("conn should return an empty slice when a batch contains no statements")
+	}
+	if _, err := c.ExecContext(ctx, "create table table1", []driver.NamedValue{}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := c.ExecContext(ctx, "create table table2", []driver.NamedValue{}); err != nil {
+		t.Fatal(err)
+	}
+	batchedStatements := c.GetBatchedStatements()
+	if !reflect.DeepEqual([]spanner.Statement{
+		{SQL: "create table table1", Params: map[string]interface{}{}},
+		{SQL: "create table table2", Params: map[string]interface{}{}},
+	}, batchedStatements) {
+		t.Errorf("unexpected batched statements: %v", batchedStatements)
+	}
+
+	// Changing the returned slice does not change the batched statements.
+	batchedStatements[0] = spanner.Statement{SQL: "drop table table1"}
+	batchedStatements2 := c.GetBatchedStatements()
+	if !reflect.DeepEqual([]spanner.Statement{
+		{SQL: "create table table1", Params: map[string]interface{}{}},
+		{SQL: "create table table2", Params: map[string]interface{}{}},
+	}, batchedStatements2) {
+		t.Errorf("unexpected batched statements: %v", batchedStatements2)
+	}
+
+	if err := c.AbortBatch(); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(c.GetBatchedStatements(), []spanner.Statement{}) {
+		t.Fatal("conn should return an empty slice when no batch is active")
+	}
+}
+
 func TestConn_GetCommitTimestampAfterAutocommitDml(t *testing.T) {
 	want := time.Now()
 	c := &conn{

From ea2db3987d7ea109ddd1b29721669d4866a0a900 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Thu, 14 Nov 2024 13:36:10 +0100
Subject: [PATCH 2/2] chore: use slices.Clone

---
 driver.go | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/driver.go b/driver.go
index c42869be..9de4cc92 100644
--- a/driver.go
+++ b/driver.go
@@ -22,6 +22,7 @@ import (
 	"math/big"
 	"reflect"
 	"regexp"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -760,12 +761,10 @@ func (c *conn) InDMLBatch() bool {
 }
 
 func (c *conn) GetBatchedStatements() []spanner.Statement {
-	if c.batch == nil {
+	if c.batch == nil || c.batch.statements == nil {
 		return []spanner.Statement{}
 	}
-	res := make([]spanner.Statement, len(c.batch.statements))
-	copy(res, c.batch.statements)
-	return res
+	return slices.Clone(c.batch.statements)
 }
 
 func (c *conn) inBatch() bool {