Skip to content

Commit e655889

Browse files
olavloiteolavloite
authored
feat(spanner): wrap proto mutation (#12497)
Add an API for wrapping an existing protobuf mutation in a Spanner mutation. This allows applications or libraries that use protobuf mutations to use these directly with the Spanner client library, without first having to translate these to a spanner.Mutation struct, and then having the client library convert it back to a protobuf definition.
1 parent 2d66d4f commit e655889

4 files changed

Lines changed: 120 additions & 58 deletions

File tree

‎spanner/client_test.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5937,7 +5937,7 @@ func TestClient_BatchWrite(t *testing.T) {
59375937
defer teardown()
59385938
mutationGroups := []*MutationGroup{
59395939
{[]*Mutation{
5940-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
5940+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
59415941
}},
59425942
}
59435943
iter := client.BatchWrite(context.Background(), mutationGroups)
@@ -5972,7 +5972,7 @@ func TestClient_BatchWrite_SessionNotFound(t *testing.T) {
59725972
)
59735973
mutationGroups := []*MutationGroup{
59745974
{[]*Mutation{
5975-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
5975+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
59765976
}},
59775977
}
59785978
iter := client.BatchWrite(context.Background(), mutationGroups)
@@ -6010,7 +6010,7 @@ func TestClient_BatchWrite_Error(t *testing.T) {
60106010
)
60116011
mutationGroups := []*MutationGroup{
60126012
{[]*Mutation{
6013-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
6013+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
60146014
}},
60156015
}
60166016
iter := client.BatchWrite(context.Background(), mutationGroups)
@@ -6085,7 +6085,7 @@ func TestClient_BatchWrite_Options(t *testing.T) {
60856085

60866086
mutationGroups := []*MutationGroup{
60876087
{[]*Mutation{
6088-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
6088+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
60896089
}},
60906090
}
60916091
iter := client.BatchWriteWithOptions(context.Background(), mutationGroups, tt.write)
@@ -6131,7 +6131,7 @@ func checkBatchWriteSpan(t *testing.T, errors []error, code codes.Code) {
61316131
)
61326132
mutationGroups := []*MutationGroup{
61336133
{[]*Mutation{
6134-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
6134+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
61356135
}},
61366136
}
61376137
iter := client.BatchWrite(context.Background(), mutationGroups)
@@ -6683,7 +6683,7 @@ func TestClient_BatchWriteExcludeTxnFromChangeStreams(t *testing.T) {
66836683

66846684
mutationGroups := []*MutationGroup{
66856685
{[]*Mutation{
6686-
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}},
6686+
{opInsertOrUpdate, "t_test", nil, []string{"key", "val"}, []interface{}{"foo1", 1}, nil},
66876687
}},
66886688
}
66896689
iter := client.BatchWriteWithOptions(context.Background(), mutationGroups, BatchWriteOptions{ExcludeTxnFromChangeStreams: true})

‎spanner/mutation.go‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ limitations under the License.
1717
package spanner
1818

1919
import (
20+
"fmt"
2021
"math/rand"
2122
"reflect"
2223
"time"
@@ -141,6 +142,10 @@ type Mutation struct {
141142
// values specifies the new values for the target columns
142143
// named by Columns.
143144
values []interface{}
145+
146+
// wrapped is the protobuf mutation that is the source for this Mutation.
147+
// This is only set if the [WrapMutation] function was used to create the Mutation.
148+
wrapped *sppb.Mutation
144149
}
145150

146151
// A MutationGroup is a list of Mutation to be committed atomically.
@@ -356,6 +361,37 @@ func Delete(table string, ks KeySet) *Mutation {
356361
}
357362
}
358363

364+
// WrapMutation creates a mutation that wraps an existing protobuf mutation object.
365+
func WrapMutation(proto *sppb.Mutation) (*Mutation, error) {
366+
if proto == nil {
367+
return nil, fmt.Errorf("protobuf mutation may not be nil")
368+
}
369+
op, table, err := getTableAndSpannerOperation(proto)
370+
if err != nil {
371+
return nil, err
372+
}
373+
return &Mutation{
374+
op: op,
375+
table: table,
376+
wrapped: proto,
377+
}, nil
378+
}
379+
380+
func getTableAndSpannerOperation(proto *sppb.Mutation) (op, string, error) {
381+
if proto.GetInsert() != nil {
382+
return opInsert, proto.GetInsert().Table, nil
383+
} else if proto.GetUpdate() != nil {
384+
return opUpdate, proto.GetUpdate().Table, nil
385+
} else if proto.GetReplace() != nil {
386+
return opReplace, proto.GetReplace().Table, nil
387+
} else if proto.GetDelete() != nil {
388+
return opDelete, proto.GetDelete().Table, nil
389+
} else if proto.GetInsertOrUpdate() != nil {
390+
return opInsertOrUpdate, proto.GetInsertOrUpdate().Table, nil
391+
}
392+
return 0, "", spannerErrorf(codes.InvalidArgument, "unknown op type: %d", proto.Operation)
393+
}
394+
359395
// prepareWrite generates sppb.Mutation_Write from table name, column names
360396
// and new column values.
361397
func prepareWrite(table string, columns []string, vals []interface{}) (*sppb.Mutation_Write, error) {
@@ -375,9 +411,13 @@ func errInvdMutationOp(m Mutation) error {
375411
return spannerErrorf(codes.InvalidArgument, "Unknown op type: %d", m.op)
376412
}
377413

378-
// proto converts spanner.Mutation to sppb.Mutation, in preparation to send
414+
// proto converts a spanner.Mutation to sppb.Mutation, in preparation to send
379415
// RPCs.
380416
func (m Mutation) proto() (*sppb.Mutation, error) {
417+
if m.wrapped != nil {
418+
return m.wrapped, nil
419+
}
420+
381421
var pb *sppb.Mutation
382422
switch m.op {
383423
case opDelete:

0 commit comments

Comments
 (0)