Skip to content

Commit

Permalink
fix(spanner): support custom encoding and decoding of protos (#10799)
Browse files Browse the repository at this point in the history
* fix(spanner): add custom encoder decoder check for protos

* fix(spanner): add header
  • Loading branch information
harshachinta authored Sep 2, 2024
1 parent 633dc86 commit d410907
Show file tree
Hide file tree
Showing 6 changed files with 221 additions and 11 deletions.
41 changes: 41 additions & 0 deletions spanner/testdata/protos/custom_singer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
Copyright 2024 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package protos

import (
"errors"
)

func (c *CustomSingerInfo) EncodeSpanner() (interface{}, error) {
if c == nil {
return nil, nil
}
return c.SingerName, nil
}

func (c CustomGenre) EncodeSpanner() (interface{}, error) {
return c.String(), nil
}

func (c *CustomSingerInfo) DecodeSpanner(input interface{}) error {
str, ok := input.(string)
if !ok {
return errors.New("the interface does not contain a string")
}
c.SingerName = &str
return nil
}
Binary file modified spanner/testdata/protos/descriptors.pb
Binary file not shown.
146 changes: 135 additions & 11 deletions spanner/testdata/protos/singer.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 11 additions & 0 deletions spanner/testdata/protos/singer.proto
Original file line number Diff line number Diff line change
Expand Up @@ -32,3 +32,14 @@ enum Genre {
FOLK = 2;
ROCK = 3;
}

message CustomSingerInfo {
optional string singer_name = 1;
}

enum CustomGenre {
CUSTOM_POP = 0;
CUSTOM_JAZZ = 1;
CUSTOM_FOLK = 2;
CUSTOM_ROCK = 3;
}
29 changes: 29 additions & 0 deletions spanner/value.go
Original file line number Diff line number Diff line change
Expand Up @@ -2279,6 +2279,15 @@ func decodeValue(v *proto3.Value, t *sppb.Type, ptr interface{}, opts ...DecodeO
reflect.ValueOf(p.ProtoEnumVal).Elem().SetInt(y)
p.Valid = true
case proto.Message:
// Check if the pointer is a custom type that implements spanner.Decoder
// interface.
if decodedVal, ok := ptr.(Decoder); ok {
x, err := getGenericValue(t, v)
if err != nil {
return err
}
return decodedVal.DecodeSpanner(x)
}
if p == nil {
return errNilDst(p)
}
Expand Down Expand Up @@ -4429,6 +4438,16 @@ func encodeValue(v interface{}) (*proto3.Value, *sppb.Type, error) {
case []GenericColumnValue:
return nil, nil, errEncoderUnsupportedType(v)
case protoreflect.Enum:
// Check if the value is of protoreflect.Enum type that implements spanner.Encoder
// interface.
if encodedVal, ok := v.(Encoder); ok {
nv, err := encodedVal.EncodeSpanner()
if err != nil {
return nil, nil, err
}
return encodeValue(nv)
}

if v != nil {
var protoEnumfqn string
rv := reflect.ValueOf(v)
Expand All @@ -4447,6 +4466,16 @@ func encodeValue(v interface{}) (*proto3.Value, *sppb.Type, error) {
}
return nil, nil, errNotValidSrc(v)
case proto.Message:
// Check if the value is of proto.Message type that implements spanner.Encoder
// interface.
if encodedVal, ok := v.(Encoder); ok {
nv, err := encodedVal.EncodeSpanner()
if err != nil {
return nil, nil, err
}
return encodeValue(nv)
}

if v != nil {
if v.ProtoReflect().IsValid() {
bytes, err := proto.Marshal(v)
Expand Down
5 changes: 5 additions & 0 deletions spanner/value_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -563,6 +563,9 @@ func TestEncodeValue(t *testing.T) {
{[]*pb.Genre{nil, (*pb.Genre)(nil)}, listProto(nullProto(), nullProto()), listType(tProtoEnum), "Array of Proto Enum with nil values"},
{[]*pb.SingerInfo{singer1ProtoMsg, singer2ProtoMsg, nil, (*pb.SingerInfo)(nil)}, listProto(protoMessageProto(singer1ProtoMsg), protoMessageProto(singer2ProtoMsg), nullProto(), nullProto()), listType(tProtoMessage), "Array of Proto Message with non-nil and nil values"},
{[]*pb.Genre{&singer1ProtoEnum, &singer2ProtoEnum, nil, (*pb.Genre)(nil)}, listProto(protoEnumProto(singer1ProtoEnum), protoEnumProto(singer2ProtoEnum), nullProto(), nullProto()), listType(tProtoEnum), "Array of Proto Enum with non-nil and nil values"},
// PROTO MESSAGE AND ENUM WITH CUSTOM ENCODER
{&pb.CustomSingerInfo{SingerName: &sValue}, stringProto("abc"), tString, "Proto message with encoder interface to string"},
{pb.CustomGenre_CUSTOM_ROCK, stringProto("CUSTOM_ROCK"), tString, "Proto Enum with encoder interface to string"},
} {
got, gotType, err := encodeValue(test.in)
if err != nil {
Expand Down Expand Up @@ -1996,6 +1999,8 @@ func TestDecodeValue(t *testing.T) {
{desc: "decode all NULL elements in ARRAY<PROTO<>> to []*pb.SingerInfo", proto: listProto(nullProto(), nullProto()), protoType: listType(protoMessageType(protoMessagefqn)), want: []*pb.SingerInfo{nil, nil}},
{desc: "decode ARRAY<ENUM<>> to []*pb.Genre", proto: listProto(nullProto(), protoEnumProto(pb.Genre_ROCK), protoEnumProto(pb.Genre_FOLK)), protoType: listType(protoEnumType(protoEnumfqn)), want: []*pb.Genre{nil, &singerEnumValue, &singer2ProtoEnum}},
{desc: "decode all NULL elements in ARRAY<ENUM<>> to []*pb.Genre", proto: listProto(nullProto(), nullProto()), protoType: listType(protoEnumType(protoEnumfqn)), want: []*pb.Genre{nil, nil}},
// PROTO MESSAGE WITH CUSTOM DECODER
{desc: "decode STRING to Proto message", proto: stringProto("abc"), protoType: stringType(), want: pb.CustomSingerInfo{SingerName: proto.String("abc")}},
} {
gotp := reflect.New(reflect.TypeOf(test.want))
v := gotp.Interface()
Expand Down

0 comments on commit d410907

Please sign in to comment.