Java tutorial
/** * This file is part of Graylog. * * Graylog is free software: you can redistribute it and/or modify * it under the terms of the GNU General Public License as published by * the Free Software Foundation, either version 3 of the License, or * (at your option) any later version. * * Graylog is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU General Public License for more details. * * You should have received a copy of the GNU General Public License * along with Graylog. If not, see <http://www.gnu.org/licenses/>. */ package integration.util.mongodb; import com.mongodb.BasicDBObject; import com.mongodb.DB; import com.mongodb.DBCollection; import com.mongodb.DBObject; import com.mongodb.MongoClient; import integration.IntegrationTestsConfig; import java.io.IOException; import java.net.URI; import java.net.URL; import java.net.UnknownHostException; import java.util.List; import java.util.Map; public class MongodbSeed { private final MongoClient mongoClient; private final DB mongoDatabase; public MongodbSeed(String dbName) throws UnknownHostException { mongoClient = new MongoClient(IntegrationTestsConfig.getMongodbHost(), IntegrationTestsConfig.getMongodbPort()); mongoDatabase = mongoClient.getDB(dbName); mongoDatabase.dropDatabase(); } private Map<String, List<DBObject>> parseDatabaseDump(URL seedUrl) throws IOException { final DumpReader dumpReader; if (seedUrl.getPath().endsWith(".json")) { dumpReader = new JsonReader(seedUrl); } else { dumpReader = new BsonReader(seedUrl); } return dumpReader.toMap(); } private Map<String, List<DBObject>> updateNodeIdFirstNode(Map<String, List<DBObject>> collections, String nodeId) { final List<DBObject> nodes = collections.get("nodes"); if (nodes == null || nodes.isEmpty()) return collections; final DBObject firstNode = nodes.get(0); firstNode.put("node_id", nodeId); nodes.set(0, firstNode); collections.remove("nodes"); collections.put("nodes", nodes); return collections; } private Map<String, List<DBObject>> updateNodeIdInputs(Map<String, List<DBObject>> collections, String nodeId) { List<DBObject> inputs = collections.get("inputs"); if (inputs == null) { return collections; } for (DBObject input : inputs) { input.put("node_id", nodeId); } collections.remove("inputs"); collections.put("inputs", inputs); return collections; } public void loadDataset(URL dbPath, String nodeId) throws IOException { Map<String, List<DBObject>> collections = parseDatabaseDump(dbPath); collections = updateNodeIdFirstNode(collections, nodeId); collections = updateNodeIdInputs(collections, nodeId); for (Map.Entry<String, List<DBObject>> collection : collections.entrySet()) { final String collectionName = collection.getKey(); if (mongoDatabase.getCollection(collectionName) == null) { mongoDatabase.createCollection(collectionName, new BasicDBObject()); } final DBCollection mongoCollection = mongoDatabase.getCollection(collectionName); if (!collection.getValue().isEmpty()) { mongoCollection.insert(collection.getValue()); } } } }