com.xlrit.gears.plugin.mongo.MongoExportTarget Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of gears-runtime-plugin-mongo Show documentation
There is a newer version: 1.17.6
Show newest version
package com.xlrit.gears.plugin.mongo;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicLong;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.ReplaceOptions;
import com.xlrit.gears.engine.export.ExportOptions;
import com.xlrit.gears.engine.export.ExportTarget;
import org.bson.BsonDocument;
import org.bson.BsonDocumentWriter;
import org.bson.Document;
import org.bson.UuidRepresentation;
import org.bson.codecs.EncoderContext;
import org.bson.conversions.Bson;
import org.mongojack.internal.stream.JacksonEncoder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class MongoExportTarget implements ExportTarget {
	private static final Logger LOG = LoggerFactory.getLogger(MongoExportTarget.class);
	private static final ReplaceOptions UPSERT = new ReplaceOptions().upsert(true);
	private static final EncoderContext ENCODER_CONTEXT = EncoderContext.builder().isEncodingCollectibleDocument(true).build();

	private final MongoClient mongoClient;
	private final MongoDatabase database;
	private final ObjectMapper objectMapper;
	private final JacksonEncoder transcoder;
	private final Deque> collection = new ArrayDeque<>();
	private final AtomicLong counter = new AtomicLong(0L);

	public MongoExportTarget(MongoConfig config, ObjectMapper objectMapper) {
		this.mongoClient = MongoClients.create(config.connectionString());
		this.database = mongoClient.getDatabase(config.database());
		this.objectMapper = objectMapper;
		this.transcoder = new JacksonEncoder<>(JsonNode.class, null, objectMapper, UuidRepresentation.UNSPECIFIED);
	}

	@Override
	public ExportOptions defaultOptions() {
		return new ExportOptions("_id", null);
	}

	@Override
	public void startCollection(String name) {
		Objects.requireNonNull(name, "name must not be null");
		database.createCollection(name);
		collection.push(database.getCollection(name, BsonDocument.class));
	}

	@Override
	public void add(String id, JsonNode element) {
		MongoCollection current = collection.peek();
		Objects.requireNonNull(element, "element must not be null");
		Objects.requireNonNull(current, "no current collection set");
		LOG.info("Inserting '{}' element {}", current.getNamespace().getCollectionName(), element);

		BsonDocument document = toDocument(element);
		Bson filter = new Document("_id", element.get("_id").textValue());
		current.replaceOne(filter, document, UPSERT);

		counter.incrementAndGet();
	}

	@Override
	public void endCollection(String name) {
		Objects.requireNonNull(name, "name must not be null");
		MongoCollection current = collection.peek();
		if (!name.equals(current.getNamespace().getCollectionName()))
			throw new IllegalStateException("name mismatch");
	}

	@Override
	public JsonNode result() {
		ObjectNode result = objectMapper.createObjectNode();
		result.put("totalCount", counter.get());
		return result;
	}

	private BsonDocument toDocument(JsonNode jsonNode) {
		BsonDocument document = new BsonDocument();
		transcoder.encode(new BsonDocumentWriter(document), jsonNode, ENCODER_CONTEXT);
		return document;
	}

	@Override
	public void close() {
		mongoClient.close();
	}
}