fix: return batch update counts also when retries are disabled (#566)

The RunDmlBatch function of SpannerConn returns a Spanner-specific result type that
also contains the actual update counts per statement in a batch, and not only the
total update count. This was however not correctly returned if a DML batch was executed
in a read/write transaction with retries disabled.

Author: olavloite

conn_with_mockserver_test.go

@@ -1501,6 +1501,7 @@ func TestGenericConnectionState_GoogleSQL(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
 func TestGenericConnectionState_PostgreSQL(t *testing.T) {
 	t.Parallel()
 
@@ -1751,6 +1752,40 @@ func TestGenericConnectionState_PostgreSQL(t *testing.T) {
 	}
 }
 
+func TestDmlBatchReturnsBatchUpdateCountsOutsideTransaction(t *testing.T) {
+	t.Parallel()
+	db, _, teardown := setupTestDBConnection(t)
+	defer teardown()
+	ctx := context.Background()
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer silentClose(conn)
+
+	if _, err := conn.ExecContext(ctx, "start batch dml"); err != nil {
+		t.Fatal(err)
+	}
+	_, _ = conn.ExecContext(ctx, testutil.UpdateBarSetFoo)
+	_, _ = conn.ExecContext(ctx, testutil.UpdateSingersSetLastName)
+	var res SpannerResult
+	if err := conn.Raw(func(driverConn interface{}) error {
+		spannerConn, _ := driverConn.(SpannerConn)
+		res, err = spannerConn.RunDmlBatch(ctx)
+		return err
+	}); err != nil {
+		t.Fatal(err)
+	}
+	results, err := res.BatchRowsAffected()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if g, w := results, []int64{testutil.UpdateBarSetFooRowCount, testutil.UpdateSingersSetLastNameRowCount}; !reflect.DeepEqual(g, w) {
+		t.Fatalf("batch affected mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func verifyConnectionPropertyValue[T comparable](t *testing.T, c *sql.Conn, name string, value T) {
 	ctx := context.Background()
 	row := c.QueryRowContext(ctx, getShowStatement(c)+name)

transaction.go

@@ -744,7 +744,7 @@ func (tx *readWriteTransaction) runDmlBatch(ctx context.Context) (*result, error
 
 	if !tx.retryAborts() {
 		affected, err := tx.rwTx.BatchUpdateWithOptions(ctx, statements, options.QueryOptions)
-		return &result{rowsAffected: sum(affected)}, err
+		return &result{rowsAffected: sum(affected), batchUpdateCounts: affected}, err
 	}
 
 	var affected []int64

transaction_test.go

@@ -3,6 +3,7 @@ package spannerdriver
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"reflect"
 	"testing"
 
@@ -224,3 +225,49 @@ func TestBeginTransactionDeferrable(t *testing.T) {
 		t.Fatalf("deferrable mismatch\n Got: %v\nWant: %v", g, w)
 	}
 }
+
+func TestDmlBatchReturnsBatchUpdateCounts(t *testing.T) {
+	t.Parallel()
+	db, _, teardown := setupTestDBConnection(t)
+	defer teardown()
+	ctx := context.Background()
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer silentClose(conn)
+
+	for _, retry := range []bool{true, false} {
+		_, err := conn.ExecContext(ctx, "begin transaction")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := conn.ExecContext(ctx, fmt.Sprintf("set local retry_aborts_internally=%v", retry)); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := conn.ExecContext(ctx, "start batch dml"); err != nil {
+			t.Fatal(err)
+		}
+		_, _ = conn.ExecContext(ctx, testutil.UpdateBarSetFoo)
+		_, _ = conn.ExecContext(ctx, testutil.UpdateSingersSetLastName)
+		var res SpannerResult
+		if err := conn.Raw(func(driverConn interface{}) error {
+			spannerConn, _ := driverConn.(SpannerConn)
+			res, err = spannerConn.RunDmlBatch(ctx)
+			return err
+		}); err != nil {
+			t.Fatal(err)
+		}
+		results, err := res.BatchRowsAffected()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if g, w := results, []int64{testutil.UpdateBarSetFooRowCount, testutil.UpdateSingersSetLastNameRowCount}; !reflect.DeepEqual(g, w) {
+			t.Fatalf("batch affected mismatch\n Got: %v\nWant: %v", g, w)
+		}
+		if _, err := conn.ExecContext(ctx, "commit"); err != nil {
+			t.Fatal(err)
+		}
+	}
+}