Skip to content

Commit e4851af

Browse files
jyeminjyemin
committed
JAVA-2703: Ensure that the ElementExtendingBsonWriter appends the extra elements when a reader is piped into the writer at the top level
1 parent 8c1b465 commit e4851af

File tree

12 files changed

+323
-122
lines changed

12 files changed

+323
-122
lines changed

‎bson/src/main/org/bson/AbstractBsonWriter.java

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import java.io.Closeable;
2323
import java.util.Arrays;
24+
import java.util.List;
2425
import java.util.Map;
2526
import java.util.Stack;
2627

@@ -752,36 +753,58 @@ public void close() {
752753
@Override
753754
public void pipe(final BsonReader reader) {
754755
notNull("reader", reader);
755-
reader.readStartDocument();
756-
writeStartDocument();
757-
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
758-
writeName(reader.readName());
759-
pipeValue(reader);
756+
pipeDocument(reader, null);
757+
}
758+
759+
/**
760+
* Reads a single document from the given BsonReader and writes it to this, appending the given extra elements to the document.
761+
*
762+
* @param reader the source of the document
763+
* @param extraElements the extra elements to append to the document
764+
* @since 3.6
765+
*/
766+
public void pipe(final BsonReader reader, final List<BsonElement> extraElements) {
767+
notNull("reader", reader);
768+
notNull("extraElements", extraElements);
769+
pipeDocument(reader, extraElements);
770+
}
771+
772+
/**
773+
* Pipe a list of extra element to this writer
774+
*
775+
* @param extraElements the extra elements
776+
*/
777+
protected void pipeExtraElements(final List<BsonElement> extraElements) {
778+
notNull("extraElements", extraElements);
779+
for (BsonElement cur : extraElements) {
780+
writeName(cur.getName());
781+
pipeValue(cur.getValue());
760782
}
761-
reader.readEndDocument();
762-
writeEndDocument();
763783
}
764784

765-
private void pipeDocument(final BsonReader reader) {
785+
private void pipeDocument(final BsonReader reader, final List<BsonElement> extraElements) {
766786
reader.readStartDocument();
767787
writeStartDocument();
768788
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
769789
writeName(reader.readName());
770790
pipeValue(reader);
771791
}
772792
reader.readEndDocument();
793+
if (extraElements != null) {
794+
pipeExtraElements(extraElements);
795+
}
773796
writeEndDocument();
774797
}
775798

776799
private void pipeJavascriptWithScope(final BsonReader reader) {
777800
writeJavaScriptWithScope(reader.readJavaScriptWithScope());
778-
pipeDocument(reader);
801+
pipeDocument(reader, null);
779802
}
780803

781804
private void pipeValue(final BsonReader reader) {
782805
switch (reader.getCurrentBsonType()) {
783806
case DOCUMENT:
784-
pipeDocument(reader);
807+
pipeDocument(reader, null);
785808
break;
786809
case ARRAY:
787810
pipeArray(reader);

‎bson/src/main/org/bson/BsonBinaryWriter.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.bson.types.Decimal128;
2222
import org.bson.types.ObjectId;
2323

24+
import java.util.List;
2425
import java.util.Stack;
2526

2627
import static java.lang.String.format;
@@ -303,6 +304,17 @@ public void doWriteUndefined() {
303304
@Override
304305
public void pipe(final BsonReader reader) {
305306
notNull("reader", reader);
307+
pipeDocument(reader, null);
308+
}
309+
310+
@Override
311+
public void pipe(final BsonReader reader, final List<BsonElement> extraElements) {
312+
notNull("reader", reader);
313+
notNull("extraElements", extraElements);
314+
pipeDocument(reader, extraElements);
315+
}
316+
317+
private void pipeDocument(final BsonReader reader, final List<BsonElement> extraElements) {
306318
if (reader instanceof BsonBinaryReader) {
307319
BsonBinaryReader binaryReader = (BsonBinaryReader) reader;
308320
if (getState() == State.VALUE) {
@@ -322,6 +334,16 @@ public void pipe(final BsonReader reader) {
322334

323335
binaryReader.setState(AbstractBsonReader.State.TYPE);
324336

337+
if (extraElements != null) {
338+
bsonOutput.truncateToPosition(bsonOutput.getPosition() - 1);
339+
setContext(new Context(getContext(), BsonContextType.DOCUMENT, pipedDocumentStartPosition));
340+
setState(State.NAME);
341+
pipeExtraElements(extraElements);
342+
bsonOutput.writeByte(0);
343+
bsonOutput.writeInt32(pipedDocumentStartPosition, bsonOutput.getPosition() - pipedDocumentStartPosition);
344+
setContext(getContext().getParentContext());
345+
}
346+
325347
if (getContext() == null) {
326348
setState(State.DONE);
327349
} else {
@@ -333,6 +355,8 @@ public void pipe(final BsonReader reader) {
333355
}
334356

335357
validateSize(bsonOutput.getPosition() - pipedDocumentStartPosition);
358+
} else if (extraElements != null) {
359+
super.pipe(reader, extraElements);
336360
} else {
337361
super.pipe(reader);
338362
}

‎bson/src/test/unit/org/bson/BsonBinaryWriterTest.java

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@
2626
import java.io.ByteArrayOutputStream;
2727
import java.io.IOException;
2828
import java.nio.ByteBuffer;
29+
import java.util.List;
2930

31+
import static java.util.Arrays.asList;
3032
import static org.hamcrest.CoreMatchers.is;
3133
import static org.junit.Assert.assertArrayEquals;
3234
import static org.junit.Assert.assertEquals;
@@ -633,6 +635,98 @@ public void testPipeDocumentIntoScopeDocument() {
633635
reader2.readEndDocument();
634636
}
635637

638+
@Test
639+
public void testPipeWithExtraElements() {
640+
writer.writeStartDocument();
641+
writer.writeBoolean("a", true);
642+
writer.writeString("$db", "test");
643+
writer.writeStartDocument("$readPreference");
644+
writer.writeString("mode", "primary");
645+
writer.writeEndDocument();
646+
writer.writeEndDocument();
647+
648+
byte[] bytes = buffer.toByteArray();
649+
650+
BasicOutputBuffer pipedBuffer = new BasicOutputBuffer();
651+
BsonBinaryWriter pipedWriter = new BsonBinaryWriter(new BsonWriterSettings(100),
652+
new BsonBinaryWriterSettings(1024), pipedBuffer);
653+
654+
pipedWriter.writeStartDocument();
655+
pipedWriter.writeBoolean("a", true);
656+
pipedWriter.writeEndDocument();
657+
658+
List<BsonElement> extraElements = asList(
659+
new BsonElement("$db", new BsonString("test")),
660+
new BsonElement("$readPreference", new BsonDocument("mode", new BsonString("primary")))
661+
);
662+
663+
BasicOutputBuffer newBuffer = new BasicOutputBuffer();
664+
BsonBinaryWriter newWriter = new BsonBinaryWriter(newBuffer);
665+
try {
666+
BsonBinaryReader reader =
667+
new BsonBinaryReader(new ByteBufferBsonInput(new ByteBufNIO(ByteBuffer.wrap(pipedBuffer.toByteArray()))));
668+
try {
669+
newWriter.pipe(reader, extraElements);
670+
} finally {
671+
reader.close();
672+
}
673+
} finally {
674+
newWriter.close();
675+
}
676+
assertArrayEquals(bytes, newBuffer.toByteArray());
677+
}
678+
679+
@Test
680+
public void testPipeOfNestedDocumentWithExtraElements() {
681+
writer.writeStartDocument();
682+
writer.writeStartDocument("nested");
683+
684+
writer.writeBoolean("a", true);
685+
writer.writeString("$db", "test");
686+
writer.writeStartDocument("$readPreference");
687+
writer.writeString("mode", "primary");
688+
writer.writeEndDocument();
689+
writer.writeEndDocument();
690+
691+
writer.writeBoolean("b", true);
692+
writer.writeEndDocument();
693+
694+
byte[] bytes = buffer.toByteArray();
695+
696+
BasicOutputBuffer pipedBuffer = new BasicOutputBuffer();
697+
BsonBinaryWriter pipedWriter = new BsonBinaryWriter(new BsonWriterSettings(100),
698+
new BsonBinaryWriterSettings(1024), pipedBuffer);
699+
700+
pipedWriter.writeStartDocument();
701+
pipedWriter.writeBoolean("a", true);
702+
pipedWriter.writeEndDocument();
703+
704+
List<BsonElement> extraElements = asList(
705+
new BsonElement("$db", new BsonString("test")),
706+
new BsonElement("$readPreference", new BsonDocument("mode", new BsonString("primary")))
707+
);
708+
709+
BasicOutputBuffer newBuffer = new BasicOutputBuffer();
710+
BsonBinaryWriter newWriter = new BsonBinaryWriter(newBuffer);
711+
try {
712+
BsonBinaryReader reader =
713+
new BsonBinaryReader(new ByteBufferBsonInput(new ByteBufNIO(ByteBuffer.wrap(pipedBuffer.toByteArray()))));
714+
try {
715+
newWriter.writeStartDocument();
716+
newWriter.writeName("nested");
717+
newWriter.pipe(reader, extraElements);
718+
newWriter.writeBoolean("b", true);
719+
newWriter.writeEndDocument();
720+
} finally {
721+
reader.close();
722+
}
723+
} finally {
724+
newWriter.close();
725+
}
726+
byte[] actualBytes = newBuffer.toByteArray();
727+
assertArrayEquals(bytes, actualBytes);
728+
}
729+
636730
@Test
637731
public void testPipeOfDocumentWithInvalidSize() {
638732
byte[] bytes = {4, 0, 0, 0}; // minimum document size is 5;

‎bson/src/test/unit/org/bson/BsonDocumentWriterSpecification.groovy

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,4 +46,22 @@ class BsonDocumentWriterSpecification extends Specification {
4646
then:
4747
document == documentWithValuesOfEveryType()
4848
}
49+
50+
def 'should pipe all types with extra elements'() {
51+
given:
52+
def document = new BsonDocument()
53+
def reader = new BsonDocumentReader(new BsonDocument())
54+
def writer = new BsonDocumentWriter(document)
55+
56+
def extraElements = []
57+
for (def entry : documentWithValuesOfEveryType()) {
58+
extraElements.add(new BsonElement(entry.getKey(), entry.getValue()))
59+
}
60+
61+
when:
62+
writer.pipe(reader, extraElements)
63+
64+
then:
65+
document == documentWithValuesOfEveryType()
66+
}
4967
}

‎bson/src/test/unit/org/bson/BsonHelper.java

Lines changed: 41 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -31,36 +31,41 @@
3131
import static java.util.Arrays.asList;
3232

3333
public final class BsonHelper {
34-
private static final List<BsonValue> BSON_VALUES = asList(
35-
new BsonNull(),
36-
new BsonInt32(42),
37-
new BsonInt64(52L),
38-
new BsonDecimal128(Decimal128.parse("4.00")),
39-
new BsonBoolean(true),
40-
new BsonDateTime(new Date().getTime()),
41-
new BsonDouble(62.0),
42-
new BsonString("the fox ..."),
43-
new BsonMinKey(),
44-
new BsonMaxKey(),
45-
new BsonDbPointer("test.test", new ObjectId()),
46-
new BsonJavaScript("int i = 0;"),
47-
new BsonJavaScriptWithScope("x", new BsonDocument("x", new BsonInt32(1))),
48-
new BsonObjectId(new ObjectId()),
49-
new BsonRegularExpression("^test.*regex.*xyz$", "i"),
50-
new BsonSymbol("ruby stuff"),
51-
new BsonTimestamp(0x12345678, 5),
52-
new BsonUndefined(),
53-
new BsonBinary((byte) 80, new byte[]{5, 4, 3, 2, 1}),
54-
new BsonArray(asList(new BsonInt32(1), new BsonInt64(2L), new BsonBoolean(true),
55-
new BsonArray(asList(new BsonInt32(1), new BsonInt32(2), new BsonInt32(3), new BsonDocument("a", new BsonInt64(2L)))))),
56-
new BsonDocument("a", new BsonInt32(1)));
5734

58-
private static final BsonDocument BSON_DOCUMENT = new BsonDocument();
35+
private static final Date DATE = new Date();
36+
private static final ObjectId OBJECT_ID = new ObjectId();
5937

60-
static {
61-
for (int i = 0; i < BSON_VALUES.size(); i++) {
62-
BSON_DOCUMENT.append(Integer.toString(i), BSON_VALUES.get(i));
63-
}
38+
private static List<BsonValue> getBsonValues() {
39+
return asList(
40+
new BsonNull(),
41+
new BsonInt32(42),
42+
new BsonInt64(52L),
43+
new BsonDecimal128(Decimal128.parse("4.00")),
44+
new BsonBoolean(true),
45+
new BsonDateTime(DATE.getTime()),
46+
new BsonDouble(62.0),
47+
new BsonString("the fox ..."),
48+
new BsonMinKey(),
49+
new BsonMaxKey(),
50+
new BsonDbPointer("test.test", OBJECT_ID),
51+
new BsonJavaScript("int i = 0;"),
52+
new BsonJavaScriptWithScope("x", new BsonDocument("x", new BsonInt32(1))),
53+
new BsonObjectId(OBJECT_ID),
54+
new BsonRegularExpression("^test.*regex.*xyz$", "i"),
55+
new BsonSymbol("ruby stuff"),
56+
new BsonTimestamp(0x12345678, 5),
57+
new BsonUndefined(),
58+
new BsonBinary((byte) 80, new byte[]{5, 4, 3, 2, 1}),
59+
new BsonArray(asList(
60+
new BsonInt32(1),
61+
new BsonInt64(2L),
62+
new BsonBoolean(true),
63+
new BsonArray(asList(
64+
new BsonInt32(1),
65+
new BsonInt32(2),
66+
new BsonInt32(3),
67+
new BsonDocument("a", new BsonInt64(2L)))))),
68+
new BsonDocument("a", new BsonInt32(1)));
6469
}
6570

6671
// fail class loading if any BSON types are not represented in BSON_VALUES.
@@ -71,7 +76,7 @@ public final class BsonHelper {
7176
}
7277

7378
boolean found = false;
74-
for (BsonValue curBsonValue : BSON_VALUES) {
79+
for (BsonValue curBsonValue : getBsonValues()) {
7580
if (curBsonValue.getBsonType() == curBsonType) {
7681
found = true;
7782
break;
@@ -85,11 +90,16 @@ public final class BsonHelper {
8590
}
8691

8792
public static List<BsonValue> valuesOfEveryType() {
88-
return BSON_VALUES;
93+
return getBsonValues();
8994
}
9095

9196
public static BsonDocument documentWithValuesOfEveryType() {
92-
return BSON_DOCUMENT;
97+
BsonDocument document = new BsonDocument();
98+
List<BsonValue> bsonValues = getBsonValues();
99+
for (int i = 0; i < bsonValues.size(); i++) {
100+
document.append(Integer.toString(i), bsonValues.get(i));
101+
}
102+
return document;
93103
}
94104

95105
public static ByteBuffer toBson(final BsonDocument document) {

‎driver-core/src/main/com/mongodb/connection/CommandMessage.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -158,12 +158,12 @@ private FieldNameValidator getPayloadArrayFieldNameValidator() {
158158
}
159159

160160
private void addDocumentWithPayload(final BsonOutput bsonOutput) {
161-
BsonWriter writer = new BsonBinaryWriter(bsonOutput, getPayloadArrayFieldNameValidator());
162-
if (payload != null) {
163-
writer = new SplittablePayloadBsonWriter(writer, bsonOutput, getSettings(), payload);
164-
}
161+
BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(bsonOutput, getPayloadArrayFieldNameValidator());
162+
BsonWriter bsonWriter = payload == null
163+
? bsonBinaryWriter
164+
: new SplittablePayloadBsonWriter(bsonBinaryWriter, bsonOutput, getSettings(), payload);
165165
BsonDocument commandToEncode = getCommandToEncode();
166-
getCodec(commandToEncode).encode(writer, commandToEncode, EncoderContext.builder().build());
166+
getCodec(commandToEncode).encode(bsonWriter, commandToEncode, EncoderContext.builder().build());
167167
}
168168

169169
private int getFlagBits() {

0 commit comments

Comments
 (0)