package org.bson.codecs;
import org.bson.BsonDocument;
import org.bson.BsonElement;
import org.bson.BsonObjectId;
import org.bson.BsonReader;
import org.bson.BsonType;
import org.bson.BsonValue;
import org.bson.BsonWriter;
import org.bson.codecs.configuration.CodecRegistry;
import org.bson.types.ObjectId;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import static org.bson.codecs.BsonValueCodecProvider.getBsonTypeClassMap;
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
public class BsonDocumentCodec implements CollectibleCodec<BsonDocument> {
private static final String ID_FIELD_NAME = "_id";
private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
private final CodecRegistry codecRegistry;
private final BsonTypeCodecMap bsonTypeCodecMap;
public BsonDocumentCodec() {
this(DEFAULT_REGISTRY);
}
public BsonDocumentCodec(final CodecRegistry codecRegistry) {
if (codecRegistry == null) {
throw new IllegalArgumentException("Codec registry can not be null");
}
this.codecRegistry = codecRegistry;
this.bsonTypeCodecMap = new BsonTypeCodecMap(getBsonTypeClassMap(), codecRegistry);
}
public CodecRegistry getCodecRegistry() {
return codecRegistry;
}
@Override
public BsonDocument decode(final BsonReader reader, final DecoderContext decoderContext) {
List<BsonElement> keyValuePairs = new ArrayList<BsonElement>();
reader.readStartDocument();
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
String fieldName = reader.readName();
keyValuePairs.add(new BsonElement(fieldName, readValue(reader, decoderContext)));
}
reader.readEndDocument();
return new BsonDocument(keyValuePairs);
}
protected BsonValue readValue(final BsonReader reader, final DecoderContext decoderContext) {
return (BsonValue) bsonTypeCodecMap.get(reader.getCurrentBsonType()).decode(reader, decoderContext);
}
@Override
public void encode(final BsonWriter writer, final BsonDocument value, final EncoderContext encoderContext) {
writer.writeStartDocument();
beforeFields(writer, encoderContext, value);
for (Map.Entry<String, BsonValue> entry : value.entrySet()) {
if (skipField(encoderContext, entry.getKey())) {
continue;
}
writer.writeName(entry.getKey());
writeValue(writer, encoderContext, entry.getValue());
}
writer.writeEndDocument();
}
private void beforeFields(final BsonWriter bsonWriter, final EncoderContext encoderContext, final BsonDocument value) {
if (encoderContext.isEncodingCollectibleDocument() && value.containsKey(ID_FIELD_NAME)) {
bsonWriter.writeName(ID_FIELD_NAME);
writeValue(bsonWriter, encoderContext, value.get(ID_FIELD_NAME));
}
}
private boolean skipField(final EncoderContext encoderContext, final String key) {
return encoderContext.isEncodingCollectibleDocument() && key.equals(ID_FIELD_NAME);
}
@SuppressWarnings({"unchecked", "rawtypes"})
private void writeValue(final BsonWriter writer, final EncoderContext encoderContext, final BsonValue value) {
Codec codec = codecRegistry.get(value.getClass());
encoderContext.encodeWithChildContext(codec, writer, value);
}
@Override
public Class<BsonDocument> getEncoderClass() {
return BsonDocument.class;
}
@Override
public BsonDocument generateIdIfAbsentFromDocument(final BsonDocument document) {
if (!documentHasId(document)) {
document.put(ID_FIELD_NAME, new BsonObjectId(new ObjectId()));
}
return document;
}
@Override
public boolean documentHasId(final BsonDocument document) {
return document.containsKey(ID_FIELD_NAME);
}
@Override
public BsonValue getDocumentId(final BsonDocument document) {
return document.get(ID_FIELD_NAME);
}
}