feat: add ExecuteBatchDml func with batch result (#454)

Adds an ExecuteBatchDml function that returns the actual batch result. This
function can be used to execute a batch of DML statements and get back the
exact number of rows affected per statement, instead of only the total number
of rows affected by the statements in the batch.

Fixes #377

Author: olavloite

conn.go

@@ -521,7 +521,11 @@ func (c *conn) runDDLBatch(ctx context.Context) (driver.Result, error) {
 	return c.execDDL(ctx, statements...)
 }
 
-func (c *conn) runDMLBatch(ctx context.Context) (driver.Result, error) {
+func (c *conn) runDMLBatch(ctx context.Context) (SpannerResult, error) {
+	if c.inTransaction() {
+		return c.tx.RunDmlBatch(ctx)
+	}
+
 	statements := c.batch.statements
 	options := c.batch.options
 	options.QueryOptions.LastStatement = true
@@ -566,7 +570,7 @@ func (c *conn) execDDL(ctx context.Context, statements ...spanner.Statement) (dr
 	return driver.ResultNoRows, nil
 }
 
-func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement, options ExecOptions) (driver.Result, error) {
+func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement, options ExecOptions) (SpannerResult, error) {
 	if len(statements) == 0 {
 		return &result{}, nil
 	}
@@ -585,7 +589,7 @@ func (c *conn) execBatchDML(ctx context.Context, statements []spanner.Statement,
 			return err
 		}, options.TransactionOptions)
 	}
-	return &result{rowsAffected: sum(affected)}, err
+	return &result{rowsAffected: sum(affected), batchUpdateCounts: affected}, err
 }
 
 func sum(affected []int64) int64 {

driver.go

@@ -112,6 +112,18 @@ func determineDefaultStatementCacheSize() {
 	}
 }
 
+// SpannerResult is the result type returned by Spanner connections for
+// DML batches. This interface extends the standard sql.Result interface
+// and adds a BatchRowsAffected function that returns the affected rows
+// per statement.
+type SpannerResult interface {
+	driver.Result
+
+	// BatchRowsAffected returns the affected rows per statement in a DML batch.
+	// It returns an error if the statement was not a DML batch.
+	BatchRowsAffected() ([]int64, error)
+}
+
 // ExecOptions can be passed in as an argument to the Query, QueryContext,
 // Exec, and ExecContext functions to specify additional execution options
 // for a statement.
@@ -1033,6 +1045,94 @@ func clearTempReadOnlyTransactionOptions(conn *sql.Conn) {
 	_ = conn.Close()
 }
 
+// DmlBatch is used to execute a batch of DML statements on Spanner in a single round-trip.
+type DmlBatch interface {
+	// ExecContext buffers the given statement for execution on Spanner.
+	// All buffered statements are sent to Spanner as a single request when the DmlBatch
+	// function returns successfully.
+	ExecContext(ctx context.Context, dml string, args ...any) error
+}
+
+var _ DmlBatch = &dmlBatch{}
+
+type dmlBatch struct {
+	conn *sql.Conn
+}
+
+// ExecuteBatchDml executes a batch of DML statements in a single round-trip to Spanner.
+func ExecuteBatchDml(ctx context.Context, db *sql.DB, f func(ctx context.Context, batch DmlBatch) error) (SpannerResult, error) {
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return ExecuteBatchDmlOnConn(ctx, conn, f)
+}
+
+// ExecuteBatchDmlOnConn executes a batch of DML statements on a specific connection in a single round-trip to Spanner.
+func ExecuteBatchDmlOnConn(ctx context.Context, connection *sql.Conn, f func(ctx context.Context, batch DmlBatch) error) (SpannerResult, error) {
+	// Start the DML batch.
+	if err := connection.Raw(func(driverConn any) error {
+		c, ok := driverConn.(*conn)
+		if !ok {
+			return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "connection is not a Spanner connection"))
+		}
+		if _, err := c.startBatchDML(false); err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	// Let the callback execute the statements on the batch.
+	b := &dmlBatch{conn: connection}
+	if err := f(ctx, b); err != nil {
+		// The callback returned an error, abort the batch.
+		_ = connection.Raw(func(driverConn any) error {
+			c, _ := driverConn.(*conn)
+			_ = c.AbortBatch()
+			return nil
+		})
+		return nil, err
+	}
+
+	// Send the batch to Spanner.
+	var res SpannerResult
+	if err := connection.Raw(func(driverConn any) error {
+		// We know that the connection is a Spanner connection, so we don't bother to check that again here.
+		c, _ := driverConn.(*conn)
+		var err error
+		res, err = c.runDMLBatch(ctx)
+		if err != nil {
+			// Make sure the batch is removed from the connection/transaction.
+			_ = c.AbortBatch()
+			return err
+		}
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+func (b *dmlBatch) ExecContext(ctx context.Context, dml string, args ...any) error {
+	if err := b.conn.Raw(func(driverConn any) error {
+		c, ok := driverConn.(*conn)
+		if !ok {
+			return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "connection is not a Spanner connection"))
+		}
+		if !c.InDMLBatch() {
+			return spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "this batch is no longer active"))
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+	_, err := b.conn.ExecContext(ctx, dml, args...)
+	return err
+}
+
 // AutocommitDMLMode indicates whether a single DML statement should be executed
 // in a normal atomic transaction or as a Partitioned DML statement.
 // See https://cloud.google.com/spanner/docs/dml-partitioned for more information.

driver_with_mockserver_test.go

@@ -2755,6 +2755,133 @@ func TestAutocommitBatchDml(t *testing.T) {
 	}
 }
 
+func TestExecuteBatchDml(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	_ = server.TestSpanner.PutStatementResult("INSERT INTO Foo (Id, Val) VALUES (1, 'One')", &testutil.StatementResult{
+		Type:        testutil.StatementResultUpdateCount,
+		UpdateCount: 1,
+	})
+	_ = server.TestSpanner.PutStatementResult("INSERT INTO Foo (Id, Val) VALUES (2, 'Two')", &testutil.StatementResult{
+		Type:        testutil.StatementResultUpdateCount,
+		UpdateCount: 1,
+	})
+
+	res, err := ExecuteBatchDml(ctx, db, func(ctx context.Context, batch DmlBatch) error {
+		if err := batch.ExecContext(ctx, "INSERT INTO Foo (Id, Val) VALUES (1, 'One')"); err != nil {
+			return err
+		}
+		if err := batch.ExecContext(ctx, "INSERT INTO Foo (Id, Val) VALUES (2, 'Two')"); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to execute dml batch: %v", err)
+	}
+	affected, err := res.RowsAffected()
+	if err != nil {
+		t.Fatalf("could not get rows affected from batch: %v", err)
+	}
+	if g, w := affected, int64(2); g != w {
+		t.Fatalf("affected rows mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	batchAffected, err := res.BatchRowsAffected()
+	if err != nil {
+		t.Fatalf("could not get batch rows affected from batch: %v", err)
+	}
+	if g, w := batchAffected, []int64{1, 1}; !cmp.Equal(g, w) {
+		t.Fatalf("affected batch rows mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	// There should be no ExecuteSqlRequests on the server.
+	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 0; g != w {
+		t.Fatalf("sql requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	batchRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteBatchDmlRequest{}))
+	if g, w := len(batchRequests), 1; g != w {
+		t.Fatalf("BatchDML requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if !batchRequests[0].(*sppb.ExecuteBatchDmlRequest).LastStatements {
+		t.Fatal("last statements flag not set")
+	}
+	// The transaction should have been committed.
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 1; g != w {
+		t.Fatalf("Commit requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
+func TestExecuteBatchDmlError(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	_ = server.TestSpanner.PutStatementResult("INSERT INTO Foo (Id, Val) VALUES (1, 'One')", &testutil.StatementResult{
+		Type:        testutil.StatementResultUpdateCount,
+		UpdateCount: 1,
+	})
+	c, err := db.Conn(ctx)
+	defer func() { _ = c.Close() }()
+	if err != nil {
+		t.Fatalf("failed to obtain connection: %v", err)
+	}
+
+	_, err = ExecuteBatchDmlOnConn(ctx, c, func(ctx context.Context, batch DmlBatch) error {
+		if err := batch.ExecContext(ctx, "INSERT INTO Foo (Id, Val) VALUES (1, 'One')"); err != nil {
+			return err
+		}
+		return fmt.Errorf("test error")
+	})
+	if err == nil {
+		t.Fatalf("failed to execute dml batch: %v", err)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	// There should be no requests on the server.
+	batchRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteBatchDmlRequest{}))
+	if g, w := len(batchRequests), 0; g != w {
+		t.Fatalf("BatchDML requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 0; g != w {
+		t.Fatalf("Commit requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	// Verify that the connection is not in a batch, and that it can be used for other statements.
+	if err := c.Raw(func(driverConn any) error {
+		spannerConn, ok := driverConn.(SpannerConn)
+		if !ok {
+			return fmt.Errorf("driver connection is not a SpannerConn")
+		}
+		if spannerConn.InDMLBatch() {
+			return fmt.Errorf("connection is still in a batch")
+		}
+		return nil
+	}); err != nil {
+		t.Fatalf("check if connection is in a batch failed: %v", err)
+	}
+	res, err := c.ExecContext(ctx, `INSERT INTO Foo (Id, Val) VALUES (1, 'One')`)
+	if err != nil {
+		t.Fatalf("failed to execute dml statement: %v", err)
+	}
+	if affected, err := res.RowsAffected(); err != nil {
+		t.Fatalf("failed to obtain rows affected: %v", err)
+	} else {
+		if g, w := affected, int64(1); g != w {
+			t.Fatalf("affected rows mismatch\n Got: %v\nWant: %v", g, w)
+		}
+	}
+}
+
 func TestTransactionBatchDml(t *testing.T) {
 	t.Parallel()
 
@@ -2872,6 +2999,80 @@ func TestTransactionBatchDml(t *testing.T) {
 	}
 }
 
+func TestExecuteBatchDmlTransaction(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	db, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	_ = server.TestSpanner.PutStatementResult("INSERT INTO Foo (Id, Val) VALUES (1, 'One')", &testutil.StatementResult{
+		Type:        testutil.StatementResultUpdateCount,
+		UpdateCount: 1,
+	})
+	_ = server.TestSpanner.PutStatementResult("INSERT INTO Foo (Id, Val) VALUES (2, 'Two')", &testutil.StatementResult{
+		Type:        testutil.StatementResultUpdateCount,
+		UpdateCount: 1,
+	})
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatalf("failed to obtain connection: %v", err)
+	}
+	tx, err := conn.BeginTx(ctx, &sql.TxOptions{})
+	if err != nil {
+		t.Fatalf("failed to begin transaction: %v", err)
+	}
+	res, err := ExecuteBatchDmlOnConn(ctx, conn, func(ctx context.Context, batch DmlBatch) error {
+		if err := batch.ExecContext(ctx, "INSERT INTO Foo (Id, Val) VALUES (1, 'One')"); err != nil {
+			return err
+		}
+		if err := batch.ExecContext(ctx, "INSERT INTO Foo (Id, Val) VALUES (2, 'Two')"); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to execute dml batch: %v", err)
+	}
+	affected, err := res.RowsAffected()
+	if err != nil {
+		t.Fatalf("could not get rows affected from batch: %v", err)
+	}
+	if g, w := affected, int64(2); g != w {
+		t.Fatalf("affected rows mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	batchAffected, err := res.BatchRowsAffected()
+	if err != nil {
+		t.Fatalf("could not get batch rows affected from batch: %v", err)
+	}
+	if g, w := batchAffected, []int64{1, 1}; !cmp.Equal(g, w) {
+		t.Fatalf("affected batch rows mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	if err := tx.Commit(); err != nil {
+		t.Fatalf("failed to commit transaction after batch: %v", err)
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	// There should be no ExecuteSqlRequests on the server.
+	sqlRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteSqlRequest{}))
+	if g, w := len(sqlRequests), 0; g != w {
+		t.Fatalf("sql requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	batchRequests := requestsOfType(requests, reflect.TypeOf(&sppb.ExecuteBatchDmlRequest{}))
+	if g, w := len(batchRequests), 1; g != w {
+		t.Fatalf("BatchDML requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if batchRequests[0].(*sppb.ExecuteBatchDmlRequest).LastStatements {
+		t.Fatal("last statements flag was set, this should not happen for batches in a transaction")
+	}
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 1; g != w {
+		t.Fatalf("Commit requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
 func TestCommitTimestamp(t *testing.T) {
 	t.Parallel()
 

stmt.go

@@ -223,6 +223,8 @@ func convertParam(v driver.Value) driver.Value {
 	}
 }
 
+var _ SpannerResult = &result{}
+
 type result struct {
 	rowsAffected      int64
 	lastInsertId      int64
@@ -246,3 +248,12 @@ func (r *result) LastInsertId() (int64, error) {
 func (r *result) RowsAffected() (int64, error) {
 	return r.rowsAffected, nil
 }
+
+var errNoBatchRowsAffected = spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "BatchRowsAffected is only supported for batch DML results"))
+
+func (r *result) BatchRowsAffected() ([]int64, error) {
+	if r.batchUpdateCounts == nil {
+		return nil, errNoBatchRowsAffected
+	}
+	return r.batchUpdateCounts, nil
+}