Skip to content

Commit

Permalink
feat: add method for getting batched statements (#318)
Browse files Browse the repository at this point in the history
* 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.

* chore: use slices.Clone
  • Loading branch information
olavloite authored Nov 14, 2024
1 parent 5cb6fc6 commit 03d4818
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 0 deletions.
12 changes: 12 additions & 0 deletions driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"math/big"
"reflect"
"regexp"
"slices"
"strconv"
"strings"
"sync"
Expand Down Expand Up @@ -535,6 +536,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.
Expand Down Expand Up @@ -755,6 +760,13 @@ 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 || c.batch.statements == nil {
return []spanner.Statement{}
}
return slices.Clone(c.batch.statements)
}

func (c *conn) inBatch() bool {
return c.InDDLBatch() || c.InDMLBatch()
}
Expand Down
47 changes: 47 additions & 0 deletions driver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"context"
"database/sql/driver"
"net"
"reflect"
"testing"
"time"

Expand Down Expand Up @@ -425,6 +426,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{
Expand Down

0 comments on commit 03d4818

Please sign in to comment.