added support for reading multiple top-level documents.

Change-Id: I36ebc5b18c1ed83cfb02716a50bf81318263cd32

Author: craiggwilson

bson/bsonrw/extjson_parser.go

@@ -163,10 +163,12 @@ func (ejp *extJSONParser) readKey() (string, bsontype.Type, error) {
 	case jpsSawValue, jpsSawEndObject, jpsSawEndArray:
 		ejp.advanceState()
 		switch ejp.s {
-		case jpsSawComma:
+		case jpsSawBeginObject, jpsSawComma:
 			ejp.advanceState()
-		case jpsSawEndObject, jpsDoneState:
+		case jpsSawEndObject:
 			return "", 0, ErrEOD
+		case jpsDoneState:
+			return "", 0, io.EOF
 		case jpsInvalidState:
 			return "", 0, ejp.err
 		default:
@@ -552,6 +554,12 @@ var jpsValidTransitionTokens = map[jsonParseState]map[jsonTokenType]bool{
 
 func (ejp *extJSONParser) validateToken(jtt jsonTokenType) bool {
 	switch ejp.s {
+	case jpsSawEndObject:
+		// if we are at depth zero and the next token is a '{',
+		// we can consider it valid only if we are not in array mode.
+		if jtt == jttBeginObject && ejp.depth == 0 {
+			return ejp.peekMode() != jpmArrayMode
+		}
 	case jpsSawComma:
 		switch ejp.peekMode() {
 		// the only valid next token after a comma inside a document is a string (a key)
@@ -560,13 +568,10 @@ func (ejp *extJSONParser) validateToken(jtt jsonTokenType) bool {
 		case jpmInvalidMode:
 			return false
 		}
-
-		// fallthrough for commas in arrays
-		fallthrough
-	default:
-		_, ok := jpsValidTransitionTokens[ejp.s][jtt]
-		return ok
 	}
+
+	_, ok := jpsValidTransitionTokens[ejp.s][jtt]
+	return ok
 }
 
 // ensureExtValueType returns true if the current value has the expected

bson/bsonrw/extjson_parser_test.go

@@ -7,6 +7,7 @@
 package bsonrw
 
 import (
+	"io"
 	"strings"
 	"testing"
 
@@ -19,6 +20,7 @@ var (
 	typDiff = specificDiff("type")
 	valDiff = specificDiff("value")
 
+	expectErrEOF = expectSpecificError(io.EOF)
 	expectErrEOD = expectSpecificError(ErrEOD)
 	expectErrEOA = expectSpecificError(ErrEOA)
 )
@@ -694,7 +696,7 @@ func TestExtJSONParserAllTypes(t *testing.T) {
 
 	// expect end of whole document: read EOF
 	k, typ, err = ejp.readKey()
-	readKeyDiff(t, "", k, bsontype.Type(0), typ, err, expectErrEOD, "")
+	readKeyDiff(t, "", k, bsontype.Type(0), typ, err, expectErrEOF, "")
 	if diff := cmp.Diff(jpsDoneState, ejp.s); diff != "" {
 		t.Errorf("expected parser to be in done state but instead is in %v\n", ejp.s)
 		t.FailNow()

bson/bsonrw/extjson_reader_test.go

@@ -7,6 +7,9 @@
 package bsonrw
 
 import (
+	"fmt"
+	"io"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -45,3 +48,118 @@ func TestExtJSONReader(t *testing.T) {
 		})
 	})
 }
+
+func TestReadMultipleTopLevelDocuments(t *testing.T) {
+	testCases := []struct {
+		name     string
+		input    string
+		expected [][]byte
+	}{
+		{
+			"single top-level document",
+			"{\"foo\":1}",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+		{
+			"single top-level document with leading and trailing whitespace",
+			"\n\n   {\"foo\":1}   \n",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+		{
+			"two top-level documents",
+			"{\"foo\":1}{\"foo\":2}",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x02, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+		{
+			"two top-level documents with leading and trailing whitespace and whitespace separation ",
+			"\n\n  {\"foo\":1}\n{\"foo\":2}\n  ",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x02, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+		{
+			"top-level array with single document",
+			"[{\"foo\":1}]",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+		{
+			"top-level array with 2 documents",
+			"[{\"foo\":1},{\"foo\":2}]",
+			[][]byte{
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x01, 0x00, 0x00, 0x00, 0x00},
+				{0x0E, 0x00, 0x00, 0x00, 0x10, 'f', 'o', 'o', 0x00, 0x02, 0x00, 0x00, 0x00, 0x00},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			r := strings.NewReader(tc.input)
+			vr := NewExtJSONValueReader(r, false)
+
+			actual, err := readAllDocuments(vr)
+			if err != nil {
+				t.Fatalf("expected no error, but got %v", err)
+			}
+
+			if diff := cmp.Diff(tc.expected, actual); diff != "" {
+				t.Fatalf("expected does not match actual: %v", diff)
+			}
+		})
+	}
+}
+
+func readAllDocuments(vr ValueReader) ([][]byte, error) {
+	c := NewCopier()
+	var actual [][]byte
+
+	switch vr.Type() {
+	case bsontype.EmbeddedDocument:
+		for {
+			result, err := c.CopyDocumentToBytes(vr)
+			if err != nil {
+				if err == io.EOF {
+					break
+				}
+				return nil, err
+			}
+
+			actual = append(actual, result)
+		}
+	case bsontype.Array:
+		ar, err := vr.ReadArray()
+		if err != nil {
+			return nil, err
+		}
+		for {
+			evr, err := ar.ReadValue()
+			if err != nil {
+				if err == ErrEOA {
+					break
+				}
+				return nil, err
+			}
+
+			result, err := c.CopyDocumentToBytes(evr)
+			if err != nil {
+				return nil, err
+			}
+
+			actual = append(actual, result)
+		}
+	default:
+		return nil, fmt.Errorf("expected an array or a document, but got %s", vr.Type())
+	}
+
+	return actual, nil
+}