feat: add mutations support

Author: olavloite

aborted_transactions_test.go

@@ -60,6 +60,51 @@ func TestCommitAborted(t *testing.T) {
 	}
 }
 
+func TestCommitWithMutationsAborted(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnectionWithParams(t, "minSessions=1;maxSessions=1")
+	defer teardown()
+	ctx := context.Background()
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatalf("failed to open connection: %v", err)
+	}
+	defer func() { _ = conn.Close() }()
+
+	tx, err := conn.BeginTx(ctx, &sql.TxOptions{})
+	if err != nil {
+		t.Fatalf("begin failed: %v", err)
+	}
+	if err := conn.Raw(func(driverConn interface{}) error {
+		spannerConn, _ := driverConn.(SpannerConn)
+		mutation := spanner.Insert("foo", []string{}, []interface{}{})
+		return spannerConn.BufferWrite([]*spanner.Mutation{mutation})
+	}); err != nil {
+		t.Fatalf("failed to buffer mutations: %v", err)
+	}
+	// Abort the transaction on the first commit attempt.
+	server.TestSpanner.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
+		Errors: []error{status.Error(codes.Aborted, "Aborted")},
+	})
+	err = tx.Commit()
+	if err != nil {
+		t.Fatalf("commit failed: %v", err)
+	}
+	reqs := drainRequestsFromServer(server.TestSpanner)
+	commitReqs := requestsOfType(reqs, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitReqs), 2; g != w {
+		t.Fatalf("commit request count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	for _, req := range commitReqs {
+		commitReq := req.(*sppb.CommitRequest)
+		if g, w := len(commitReq.Mutations), 1; g != w {
+			t.Fatalf("mutation count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+	}
+}
+
 func TestCommitAbortedWithInternalRetriesDisabled(t *testing.T) {
 	t.Parallel()
 

go.mod

@@ -4,6 +4,8 @@ go 1.24
 
 toolchain go1.24.4
 
+replace cloud.google.com/go/spanner => /Users/loite/GolandProjects/google-cloud-go/spanner
+
 require (
 	cloud.google.com/go v0.121.2
 	cloud.google.com/go/longrunning v0.6.7
@@ -26,7 +28,7 @@ require (
 	cloud.google.com/go/compute/metadata v0.7.0 // indirect
 	cloud.google.com/go/iam v1.5.2 // indirect
 	cloud.google.com/go/monitoring v1.24.2 // indirect
-	github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp v1.5.2 // indirect
+	github.com/GoogleCloudPlatform/grpc-gcp-go/grpcgcp v1.5.3 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f // indirect
@@ -44,7 +46,7 @@ require (
 	github.com/zeebo/errs v1.4.0 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
-	go.opentelemetry.io/contrib/detectors/gcp v1.35.0 // indirect
+	go.opentelemetry.io/contrib/detectors/gcp v1.36.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 // indirect
 	go.opentelemetry.io/otel v1.36.0 // indirect
@@ -60,5 +62,5 @@ require (
 	golang.org/x/text v0.26.0 // indirect
 	golang.org/x/time v0.12.0 // indirect
 	google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20250505200425-f936aa4a68b2 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822 // indirect
 )

go.sum

@@ -7,13 +7,15 @@ import (
 	"fmt"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"cloud.google.com/go/spanner"
 	"cloud.google.com/go/spanner/apiv1/spannerpb"
 	spannerdriver "github.com/googleapis/go-sql-spanner"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"spannerlib/backend"
 )
 
@@ -25,6 +27,18 @@ func CloseConnection(poolId, connId int64) *Message {
 	return conn.close()
 }
 
+func Apply(poolId, connId int64, mutationBytes []byte) *Message {
+	mutations := spannerpb.BatchWriteRequest_MutationGroup{}
+	if err := proto.Unmarshal(mutationBytes, &mutations); err != nil {
+		return errMessage(err)
+	}
+	conn, err := findConnection(poolId, connId)
+	if err != nil {
+		return errMessage(err)
+	}
+	return conn.apply(&mutations)
+}
+
 func BeginTransaction(poolId, connId int64, txOptsBytes []byte) *Message {
 	txOpts := spannerpb.TransactionOptions{}
 	if err := proto.Unmarshal(txOptsBytes, &txOpts); err != nil {
@@ -94,6 +108,34 @@ func (conn *Connection) close() *Message {
 	return &Message{}
 }
 
+func (conn *Connection) apply(mutation *spannerpb.BatchWriteRequest_MutationGroup) *Message {
+	ctx := context.Background()
+	mutations := make([]*spanner.Mutation, 0, len(mutation.Mutations))
+	for _, m := range mutation.Mutations {
+		spannerMutation, err := spanner.WrapMutation(m)
+		if err != nil {
+			return errMessage(err)
+		}
+		mutations = append(mutations, spannerMutation)
+	}
+	var commitTimestamp time.Time
+	if err := conn.backend.Conn.Raw(func(driverConn any) (err error) {
+		spannerConn, _ := driverConn.(spannerdriver.SpannerConn)
+		commitTimestamp, err = spannerConn.Apply(ctx, mutations)
+		return err
+	}); err != nil {
+		return errMessage(err)
+	}
+	response := spannerpb.CommitResponse{
+		CommitTimestamp: timestamppb.New(commitTimestamp),
+	}
+	res, err := proto.Marshal(&response)
+	if err != nil {
+		return errMessage(err)
+	}
+	return &Message{Res: res}
+}
+
 func (conn *Connection) BeginTransaction(txOpts *spannerpb.TransactionOptions) *Message {
 	var tx *sql.Tx
 	var err error

spannerlib/exported/connection.go

@@ -177,6 +177,243 @@ func TestDisableInternalRetries(t *testing.T) {
 	CloseRows(pool.ObjectId, conn.ObjectId, results.ObjectId)
 }
 
+func TestApply(t *testing.T) {
+	t.Parallel()
+
+	dsn, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	pool := CreatePool(dsn)
+	defer ClosePool(pool.ObjectId)
+	conn := CreateConnection(pool.ObjectId)
+	defer CloseConnection(pool.ObjectId, conn.ObjectId)
+
+	mutations := sppb.BatchWriteRequest_MutationGroup{
+		Mutations: []*sppb.Mutation{
+			{Operation: &sppb.Mutation_Insert{Insert: &sppb.Mutation_Write{
+				Table:   "foo",
+				Columns: []string{"id", "value"},
+				Values: []*structpb.ListValue{
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "1"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "One"}},
+					}},
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "2"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "Two"}},
+					}},
+				},
+			}}},
+		},
+	}
+	mutationsBytes, _ := proto.Marshal(&mutations)
+	response := Apply(pool.ObjectId, conn.ObjectId, mutationsBytes)
+	if response.Code != 0 {
+		t.Fatalf("failed to apply mutations: %v", response.Code)
+	}
+	commitResponse := sppb.CommitResponse{}
+	_ = proto.Unmarshal(response.Res, &commitResponse)
+	if commitResponse.CommitTimestamp == nil {
+		t.Fatal("commit timestamp missing")
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	beginRequests := requestsOfType(requests, reflect.TypeOf(&sppb.BeginTransactionRequest{}))
+	if g, w := len(beginRequests), 1; g != w {
+		t.Fatalf("begin requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	req := beginRequests[0].(*sppb.BeginTransactionRequest)
+	if req.Options == nil {
+		t.Fatalf("missing tx opts")
+	}
+	if req.Options.GetReadWrite() == nil {
+		t.Fatalf("missing tx read write")
+	}
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 1; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	commitReq := commitRequests[0].(*sppb.CommitRequest)
+	if g, w := len(commitReq.Mutations), 1; g != w {
+		t.Fatalf("mutation count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if g, w := len(commitReq.Mutations[0].GetInsert().Values), 2; g != w {
+		t.Fatalf("mutation values count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+}
+
+func TestBufferWrite(t *testing.T) {
+	t.Parallel()
+
+	dsn, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	pool := CreatePool(dsn)
+	defer ClosePool(pool.ObjectId)
+	conn := CreateConnection(pool.ObjectId)
+	defer CloseConnection(pool.ObjectId, conn.ObjectId)
+
+	txOpts := &sppb.TransactionOptions{}
+	txOptsBytes, _ := proto.Marshal(txOpts)
+	tx := BeginTransaction(pool.ObjectId, conn.ObjectId, txOptsBytes)
+
+	mutations := sppb.BatchWriteRequest_MutationGroup{
+		Mutations: []*sppb.Mutation{
+			{Operation: &sppb.Mutation_Insert{Insert: &sppb.Mutation_Write{
+				Table:   "foo",
+				Columns: []string{"id", "value"},
+				Values: []*structpb.ListValue{
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "1"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "One"}},
+					}},
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "2"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "Two"}},
+					}},
+				},
+			}}},
+		},
+	}
+	mutationsBytes, _ := proto.Marshal(&mutations)
+	response := BufferWrite(pool.ObjectId, conn.ObjectId, tx.ObjectId, mutationsBytes)
+	if response.Code != 0 {
+		t.Fatalf("failed to apply mutations: %v", response.Code)
+	}
+	if response.Length() > 0 {
+		t.Fatal("response length mismatch")
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	beginRequests := requestsOfType(requests, reflect.TypeOf(&sppb.BeginTransactionRequest{}))
+	if g, w := len(beginRequests), 1; g != w {
+		t.Fatalf("begin requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	req := beginRequests[0].(*sppb.BeginTransactionRequest)
+	if req.Options == nil {
+		t.Fatalf("missing tx opts")
+	}
+	if req.Options.GetReadWrite() == nil {
+		t.Fatalf("missing tx read write")
+	}
+
+	// There should not be any commit requests yet.
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 0; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+
+	// Commit the transaction with the mutation.
+	res := Commit(pool.ObjectId, conn.ObjectId, tx.ObjectId)
+	if res.Code != 0 {
+		t.Fatalf("failed to commit: %v", res.Code)
+	}
+
+	// Verify that we have a commit request on the server.
+	requests = drainRequestsFromServer(server.TestSpanner)
+	commitRequests = requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 1; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	commitReq := commitRequests[0].(*sppb.CommitRequest)
+	if g, w := len(commitReq.Mutations), 1; g != w {
+		t.Fatalf("mutation count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if g, w := len(commitReq.Mutations[0].GetInsert().Values), 2; g != w {
+		t.Fatalf("mutation values count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+}
+
+func TestBufferWrite_RetryAborted(t *testing.T) {
+	t.Parallel()
+
+	dsn, server, teardown := setupTestDBConnection(t)
+	defer teardown()
+
+	pool := CreatePool(dsn)
+	defer ClosePool(pool.ObjectId)
+	conn := CreateConnection(pool.ObjectId)
+	defer CloseConnection(pool.ObjectId, conn.ObjectId)
+
+	txOpts := &sppb.TransactionOptions{}
+	txOptsBytes, _ := proto.Marshal(txOpts)
+	tx := BeginTransaction(pool.ObjectId, conn.ObjectId, txOptsBytes)
+
+	mutations := sppb.BatchWriteRequest_MutationGroup{
+		Mutations: []*sppb.Mutation{
+			{Operation: &sppb.Mutation_Insert{Insert: &sppb.Mutation_Write{
+				Table:   "foo",
+				Columns: []string{"id", "value"},
+				Values: []*structpb.ListValue{
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "1"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "One"}},
+					}},
+					{Values: []*structpb.Value{
+						{Kind: &structpb.Value_StringValue{StringValue: "2"}},
+						{Kind: &structpb.Value_StringValue{StringValue: "Two"}},
+					}},
+				},
+			}}},
+		},
+	}
+	mutationsBytes, _ := proto.Marshal(&mutations)
+	response := BufferWrite(pool.ObjectId, conn.ObjectId, tx.ObjectId, mutationsBytes)
+	if response.Code != 0 {
+		t.Fatalf("failed to apply mutations: %v", response.Code)
+	}
+	if response.Length() > 0 {
+		t.Fatal("response length mismatch")
+	}
+
+	requests := drainRequestsFromServer(server.TestSpanner)
+	beginRequests := requestsOfType(requests, reflect.TypeOf(&sppb.BeginTransactionRequest{}))
+	if g, w := len(beginRequests), 1; g != w {
+		t.Fatalf("begin requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	req := beginRequests[0].(*sppb.BeginTransactionRequest)
+	if req.Options == nil {
+		t.Fatalf("missing tx opts")
+	}
+	if req.Options.GetReadWrite() == nil {
+		t.Fatalf("missing tx read write")
+	}
+
+	// There should not be any commit requests yet.
+	commitRequests := requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 0; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+
+	// Instruct the mock server to abort the transaction.
+	server.TestSpanner.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
+		Errors: []error{status.Error(codes.Aborted, "Aborted")},
+	})
+
+	// Commit the transaction with the mutation.
+	res := Commit(pool.ObjectId, conn.ObjectId, tx.ObjectId)
+	if res.Code != 0 {
+		t.Fatalf("failed to commit: %v", res.Code)
+	}
+
+	// Verify that we have a commit request on the server.
+	requests = drainRequestsFromServer(server.TestSpanner)
+	commitRequests = requestsOfType(requests, reflect.TypeOf(&sppb.CommitRequest{}))
+	if g, w := len(commitRequests), 2; g != w {
+		t.Fatalf("commit requests count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	for _, req := range commitRequests {
+		commitReq := req.(*sppb.CommitRequest)
+		if g, w := len(commitReq.Mutations), 1; g != w {
+			t.Fatalf("mutation count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+		if g, w := len(commitReq.Mutations[0].GetInsert().Values), 2; g != w {
+			t.Fatalf("mutation values count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+	}
+}
+
 func setupTestDBConnection(t *testing.T) (dsn string, server *testutil.MockedSpannerInMemTestServer, teardown func()) {
 	return setupTestDBConnectionWithParams(t, "")
 }