23
23
import com .mongodb .internal .bulk .WriteRequestWithIndex ;
24
24
import org .bson .BsonDocument ;
25
25
import org .bson .BsonDocumentWrapper ;
26
+ import org .bson .BsonObjectId ;
26
27
import org .bson .BsonValue ;
27
28
import org .bson .BsonWriter ;
28
29
import org .bson .codecs .BsonValueCodecProvider ;
37
38
import java .util .stream .Collectors ;
38
39
39
40
import static com .mongodb .assertions .Assertions .assertNotNull ;
41
+ import static com .mongodb .assertions .Assertions .assertTrue ;
40
42
import static com .mongodb .assertions .Assertions .isTrue ;
41
43
import static com .mongodb .assertions .Assertions .notNull ;
42
44
import static com .mongodb .internal .connection .SplittablePayload .Type .INSERT ;
@@ -57,6 +59,7 @@ public final class SplittablePayload {
57
59
private final WriteRequestEncoder writeRequestEncoder = new WriteRequestEncoder ();
58
60
private final Type payloadType ;
59
61
private final List <WriteRequestWithIndex > writeRequestWithIndexes ;
62
+ private final boolean ordered ;
60
63
private final Map <Integer , BsonValue > insertedIds = new HashMap <>();
61
64
private int position = 0 ;
62
65
@@ -91,9 +94,10 @@ public enum Type {
91
94
* @param payloadType the payload type
92
95
* @param writeRequestWithIndexes the writeRequests
93
96
*/
94
- public SplittablePayload (final Type payloadType , final List <WriteRequestWithIndex > writeRequestWithIndexes ) {
97
+ public SplittablePayload (final Type payloadType , final List <WriteRequestWithIndex > writeRequestWithIndexes , final boolean ordered ) {
95
98
this .payloadType = notNull ("batchType" , payloadType );
96
99
this .writeRequestWithIndexes = notNull ("writeRequests" , writeRequestWithIndexes );
100
+ this .ordered = ordered ;
97
101
}
98
102
99
103
/**
@@ -117,7 +121,7 @@ public String getPayloadName() {
117
121
}
118
122
119
123
boolean hasPayload () {
120
- return writeRequestWithIndexes .size () > 0 ;
124
+ return ! writeRequestWithIndexes .isEmpty () ;
121
125
}
122
126
123
127
public int size () {
@@ -137,10 +141,6 @@ public List<BsonDocument> getPayload() {
137
141
.collect (Collectors .toList ());
138
142
}
139
143
140
- public List <WriteRequestWithIndex > getWriteRequestWithIndexes () {
141
- return writeRequestWithIndexes ;
142
- }
143
-
144
144
/**
145
145
* @return the current position in the payload
146
146
*/
@@ -160,16 +160,22 @@ public void setPosition(final int position) {
160
160
* @return true if there are more values after the current position
161
161
*/
162
162
public boolean hasAnotherSplit () {
163
+ // this method must be not called before this payload having been encoded
164
+ assertTrue (position > 0 );
163
165
return writeRequestWithIndexes .size () > position ;
164
166
}
165
167
168
+ boolean isOrdered () {
169
+ return ordered ;
170
+ }
171
+
166
172
/**
167
173
* @return a new SplittablePayload containing only the values after the current position.
168
174
*/
169
175
public SplittablePayload getNextSplit () {
170
176
isTrue ("hasAnotherSplit" , hasAnotherSplit ());
171
177
List <WriteRequestWithIndex > nextPayLoad = writeRequestWithIndexes .subList (position , writeRequestWithIndexes .size ());
172
- return new SplittablePayload (payloadType , nextPayLoad );
178
+ return new SplittablePayload (payloadType , nextPayLoad , ordered );
173
179
}
174
180
175
181
/**
@@ -191,10 +197,23 @@ public void encode(final BsonWriter writer, final WriteRequestWithIndex writeReq
191
197
InsertRequest insertRequest = (InsertRequest ) writeRequestWithIndex .getWriteRequest ();
192
198
BsonDocument document = insertRequest .getDocument ();
193
199
194
- IdHoldingBsonWriter idHoldingBsonWriter = new IdHoldingBsonWriter (writer );
195
- getCodec (document ).encode (idHoldingBsonWriter , document ,
196
- EncoderContext .builder ().isEncodingCollectibleDocument (true ).build ());
197
- insertedIds .put (writeRequestWithIndex .getIndex (), idHoldingBsonWriter .getId ());
200
+ BsonValue documentId = insertedIds .compute (
201
+ writeRequestWithIndex .getIndex (),
202
+ (writeRequestIndex , writeRequestDocumentId ) -> {
203
+ IdHoldingBsonWriter idHoldingBsonWriter = new IdHoldingBsonWriter (
204
+ writer ,
205
+ // Reuse `writeRequestDocumentId` if it may have been generated
206
+ // by `IdHoldingBsonWriter` in a previous attempt.
207
+ // If its type is not `BsonObjectId`, we know it could not have been generated.
208
+ writeRequestDocumentId instanceof BsonObjectId ? writeRequestDocumentId .asObjectId () : null );
209
+ getCodec (document ).encode (idHoldingBsonWriter , document ,
210
+ EncoderContext .builder ().isEncodingCollectibleDocument (true ).build ());
211
+ return idHoldingBsonWriter .getId ();
212
+ });
213
+ if (documentId == null ) {
214
+ // we must add an entry anyway because we rely on all the indexes being present
215
+ insertedIds .put (writeRequestWithIndex .getIndex (), null );
216
+ }
198
217
} else if (writeRequestWithIndex .getType () == WriteRequest .Type .UPDATE
199
218
|| writeRequestWithIndex .getType () == WriteRequest .Type .REPLACE ) {
200
219
UpdateRequest update = (UpdateRequest ) writeRequestWithIndex .getWriteRequest ();
0 commit comments