From 9f346f26805724d5361d96ea8f3fb6b9cacfeb48 Mon Sep 17 00:00:00 2001
From: Motonori Iwata <121048+iwata@users.noreply.github.com>
Date: Mon, 28 Feb 2022 10:57:14 +0900
Subject: [PATCH] use BatchUpdateWithOptions instead of Update to apply DML
 (#49)

---
 pkg/spanner/client.go | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/pkg/spanner/client.go b/pkg/spanner/client.go
index 5fe5f20..81e058c 100644
--- a/pkg/spanner/client.go
+++ b/pkg/spanner/client.go
@@ -239,20 +239,30 @@ func (c *Client) ApplyDMLFile(ctx context.Context, ddl []byte, partitioned bool,
 func (c *Client) ApplyDML(ctx context.Context, statements []string, priority PriorityType) (int64, error) {
 	p := priorityPBOf(priority)
 	numAffectedRows := int64(0)
-	_, err := c.spannerClient.ReadWriteTransactionWithOptions(ctx, func(ctx context.Context, tx *spanner.ReadWriteTransaction) error {
-		for _, s := range statements {
-			num, err := tx.Update(ctx, spanner.Statement{
-				SQL: s,
+	_, err := c.spannerClient.ReadWriteTransactionWithOptions(
+		ctx,
+		func(ctx context.Context, tx *spanner.ReadWriteTransaction) error {
+			stmts := make([]spanner.Statement, len(statements))
+			for i, s := range statements {
+				stmts[i] = spanner.Statement{SQL: s}
+			}
+			counts, err := tx.BatchUpdateWithOptions(ctx, stmts, spanner.QueryOptions{
+				Priority: p,
 			})
 			if err != nil {
 				return err
 			}
-			numAffectedRows += num
-		}
-		return nil
-	}, spanner.TransactionOptions{
-		CommitPriority: p,
-	})
+
+			for _, num := range counts {
+				numAffectedRows += num
+			}
+
+			return nil
+		},
+		spanner.TransactionOptions{
+			CommitPriority: p,
+		},
+	)
 	if err != nil {
 		return 0, &Error{
 			Code: ErrorCodeUpdateDML,