Skip to content

Commit 814b9f8

Browse files
skriptbleskriptble
committed
Add support for Top Level Document reads
Adds support for reading the bytes of a top level document from a bsonrw.ValueReader. GODRIVER-673 Change-Id: Ic395e599d07ba214ef2964a186c76d9c4b4ff3b3
1 parent a4f9fc5 commit 814b9f8

File tree

4 files changed

+106
-33
lines changed

4 files changed

+106
-33
lines changed

‎bson/bsonrw/copier.go

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -118,15 +118,9 @@ func (c Copier) CopyDocumentToBytes(src ValueReader) ([]byte, error) {
118118
// AppendDocumentBytes functions the same as CopyDocumentToBytes, but will
119119
// append the result to dst.
120120
func (c Copier) AppendDocumentBytes(dst []byte, src ValueReader) ([]byte, error) {
121-
if vr, ok := src.(*valueReader); ok {
122-
length, err := vr.peakLength()
123-
if err != nil {
124-
return dst, err
125-
}
126-
dst = append(dst, vr.d[vr.offset:vr.offset+int64(length)]...)
127-
vr.offset += int64(length)
128-
vr.pop()
129-
return dst, nil
121+
if br, ok := src.(BytesReader); ok {
122+
_, dst, err := br.ReadValueBytes(dst)
123+
return dst, err
130124
}
131125

132126
vw := vwPool.Get().(*valueWriter)

‎bson/bsonrw/value_reader.go

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -260,14 +260,14 @@ func (vr *valueReader) nextElementLength() (int32, error) {
260260
var err error
261261
switch vr.stack[vr.frame].vType {
262262
case bsontype.Array, bsontype.EmbeddedDocument, bsontype.CodeWithScope:
263-
length, err = vr.peakLength()
263+
length, err = vr.peekLength()
264264
case bsontype.Binary:
265-
length, err = vr.peakLength()
265+
length, err = vr.peekLength()
266266
length += 4 + 1 // binary length + subtype byte
267267
case bsontype.Boolean:
268268
length = 1
269269
case bsontype.DBPointer:
270-
length, err = vr.peakLength()
270+
length, err = vr.peekLength()
271271
length += 4 + 12 // string length + ObjectID length
272272
case bsontype.DateTime, bsontype.Double, bsontype.Int64, bsontype.Timestamp:
273273
length = 8
@@ -276,7 +276,7 @@ func (vr *valueReader) nextElementLength() (int32, error) {
276276
case bsontype.Int32:
277277
length = 4
278278
case bsontype.JavaScript, bsontype.String, bsontype.Symbol:
279-
length, err = vr.peakLength()
279+
length, err = vr.peekLength()
280280
length += 4
281281
case bsontype.MaxKey, bsontype.MinKey, bsontype.Null, bsontype.Undefined:
282282
length = 0
@@ -303,20 +303,29 @@ func (vr *valueReader) nextElementLength() (int32, error) {
303303

304304
func (vr *valueReader) ReadValueBytes(dst []byte) (bsontype.Type, []byte, error) {
305305
switch vr.stack[vr.frame].mode {
306+
case mTopLevel:
307+
length, err := vr.peekLength()
308+
if err != nil {
309+
return bsontype.Type(0), nil, err
310+
}
311+
dst, err = vr.appendBytes(dst, length)
312+
if err != nil {
313+
return bsontype.Type(0), nil, err
314+
}
315+
return bsontype.Type(0), dst, nil
306316
case mElement, mValue:
317+
length, err := vr.nextElementLength()
318+
if err != nil {
319+
return bsontype.Type(0), dst, err
320+
}
321+
322+
dst, err = vr.appendBytes(dst, length)
323+
t := vr.stack[vr.frame].vType
324+
vr.pop()
325+
return t, dst, err
307326
default:
308327
return bsontype.Type(0), nil, vr.invalidTransitionErr(0, "ReadValueBytes", []mode{mElement, mValue})
309328
}
310-
311-
length, err := vr.nextElementLength()
312-
if err != nil {
313-
return bsontype.Type(0), dst, err
314-
}
315-
316-
dst, err = vr.appendBytes(dst, length)
317-
t := vr.stack[vr.frame].vType
318-
vr.pop()
319-
return t, dst, err
320329
}
321330

322331
func (vr *valueReader) Skip() error {
@@ -819,7 +828,7 @@ func (vr *valueReader) readString() (string, error) {
819828
return string(vr.d[start : start+int64(length)-1]), nil
820829
}
821830

822-
func (vr *valueReader) peakLength() (int32, error) {
831+
func (vr *valueReader) peekLength() (int32, error) {
823832
if vr.offset+4 > int64(len(vr.d)) {
824833
return 0, io.EOF
825834
}

‎bson/bsonrw/value_reader_test.go

Lines changed: 60 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1428,6 +1428,57 @@ func TestValueReader(t *testing.T) {
14281428
})
14291429
})
14301430
}
1431+
t.Run("ReadValueBytes/Top Level Doc", func(t *testing.T) {
1432+
testCases := []struct {
1433+
name string
1434+
want []byte
1435+
wantType bsontype.Type
1436+
wantErr error
1437+
}{
1438+
{
1439+
"success",
1440+
bsoncore.BuildDocument(nil, bsoncore.AppendDoubleElement(nil, "pi", 3.14159)),
1441+
bsontype.Type(0),
1442+
nil,
1443+
},
1444+
{
1445+
"wrong length",
1446+
[]byte{0x01, 0x02, 0x03},
1447+
bsontype.Type(0),
1448+
io.EOF,
1449+
},
1450+
{
1451+
"append bytes",
1452+
[]byte{0x01, 0x02, 0x03, 0x04},
1453+
bsontype.Type(0),
1454+
io.EOF,
1455+
},
1456+
}
1457+
1458+
for _, tc := range testCases {
1459+
tc := tc
1460+
t.Run(tc.name, func(t *testing.T) {
1461+
t.Parallel()
1462+
vr := &valueReader{
1463+
d: tc.want,
1464+
stack: []vrState{
1465+
{mode: mTopLevel},
1466+
},
1467+
frame: 0,
1468+
}
1469+
gotType, got, gotErr := vr.ReadValueBytes(nil)
1470+
if gotErr != tc.wantErr {
1471+
t.Errorf("Did not receive expected error. got %v; want %v", gotErr, tc.wantErr)
1472+
}
1473+
if tc.wantErr == nil && gotType != tc.wantType {
1474+
t.Errorf("Did not receive expected type. got %v; want %v", gotType, tc.wantType)
1475+
}
1476+
if tc.wantErr == nil && !bytes.Equal(got, tc.want) {
1477+
t.Errorf("Did not receive expected bytes. got %v; want %v", got, tc.want)
1478+
}
1479+
})
1480+
}
1481+
})
14311482
})
14321483

14331484
t.Run("invalid transition", func(t *testing.T) {
@@ -1439,14 +1490,15 @@ func TestValueReader(t *testing.T) {
14391490
t.Errorf("Expected correct invalid transition error. got %v; want %v", goterr, wanterr)
14401491
}
14411492
})
1442-
t.Run("ReadBytes", func(t *testing.T) {
1443-
vr := &valueReader{stack: []vrState{{mode: mTopLevel}}}
1444-
wanterr := (&valueReader{stack: []vrState{{mode: mTopLevel}}}).invalidTransitionErr(0, "ReadValueBytes", []mode{mElement, mValue})
1445-
_, _, goterr := vr.ReadValueBytes(nil)
1446-
if !cmp.Equal(goterr, wanterr, cmp.Comparer(compareErrors)) {
1447-
t.Errorf("Expected correct invalid transition error. got %v; want %v", goterr, wanterr)
1448-
}
1449-
})
1493+
})
1494+
t.Run("ReadBytes", func(t *testing.T) {
1495+
vr := &valueReader{stack: []vrState{{mode: mTopLevel}, {mode: mDocument}}, frame: 1}
1496+
wanterr := (&valueReader{stack: []vrState{{mode: mTopLevel}, {mode: mDocument}}, frame: 1}).
1497+
invalidTransitionErr(0, "ReadValueBytes", []mode{mElement, mValue})
1498+
_, _, goterr := vr.ReadValueBytes(nil)
1499+
if !cmp.Equal(goterr, wanterr, cmp.Comparer(compareErrors)) {
1500+
t.Errorf("Expected correct invalid transition error. got %v; want %v", goterr, wanterr)
1501+
}
14501502
})
14511503
}
14521504

‎bson/decoder_test.go

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package bson
88

99
import (
10+
"bytes"
1011
"errors"
1112
"reflect"
1213
"testing"
@@ -17,6 +18,7 @@ import (
1718
"github.com/mongodb/mongo-go-driver/bson/bsonrw/bsonrwtest"
1819
"github.com/mongodb/mongo-go-driver/bson/bsontype"
1920
"github.com/mongodb/mongo-go-driver/x/bsonx"
21+
"github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
2022
)
2123

2224
func TestBasicDecode(t *testing.T) {
@@ -116,6 +118,20 @@ func TestDecoderv2(t *testing.T) {
116118
}
117119
})
118120
}
121+
122+
t.Run("Unmarshaler/success bsonrw.ValueReader", func(t *testing.T) {
123+
want := bsoncore.BuildDocument(nil, bsoncore.AppendDoubleElement(nil, "pi", 3.14159))
124+
unmarshaler := &testUnmarshaler{}
125+
vr := bsonrw.NewBSONDocumentReader(want)
126+
dec, err := NewDecoder(DefaultRegistry, vr)
127+
noerr(t, err)
128+
err = dec.Decode(unmarshaler)
129+
noerr(t, err)
130+
got := unmarshaler.data
131+
if !bytes.Equal(got, want) {
132+
t.Errorf("Did not unmarshal properly. got %v; want %v", got, want)
133+
}
134+
})
119135
})
120136
})
121137
t.Run("NewDecoderv2", func(t *testing.T) {
@@ -205,9 +221,11 @@ func (tdc *testDecoderCodec) DecodeValue(bsoncodec.DecodeContext, bsonrw.ValueRe
205221
type testUnmarshaler struct {
206222
invoked bool
207223
err error
224+
data []byte
208225
}
209226

210-
func (tu *testUnmarshaler) UnmarshalBSON(_ []byte) error {
227+
func (tu *testUnmarshaler) UnmarshalBSON(d []byte) error {
211228
tu.invoked = true
229+
tu.data = d
212230
return tu.err
213231
}

0 commit comments

Comments
 (0)