/** * 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.fasterxml.jackson.core.JsonProcessingException; import com.mongodb.BasicDBObject; import com.mongodb.DBObject; import org.apache.commons.io.FilenameUtils; import org.bson.BSONDecoder; import org.bson.BSONObject; import org.bson.BasicBSONDecoder; import javax.annotation.Nullable; import java.io.ByteArrayInputStream; import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; import java.net.URI; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; public class BsonReader implements DumpReader { private final Map<String, List<DBObject>> collectionMap; public BsonReader(URI location) { final File dir = new File(location); collectionMap = readBsonDirectory(dir); } protected List<DBObject> readBsonFile(String filename) { Path filePath = Paths.get(filename); List<DBObject> dataset = new ArrayList<>(); try { ByteArrayInputStream fileBytes = new ByteArrayInputStream(Files.readAllBytes(filePath)); BSONDecoder decoder = new BasicBSONDecoder(); BSONObject obj; while ((obj = decoder.readObject(fileBytes)) != null) { final DBObject mongoDocument = new BasicDBObject(obj.toMap()); dataset.add(mongoDocument); } } catch (FileNotFoundException e) { e.printStackTrace(); throw new RuntimeException("Can not open BSON input file.", e); } catch (JsonProcessingException e) { e.printStackTrace(); throw new RuntimeException("Can not parse BSON data.", e); } catch (IOException e) { //EOF } return dataset; } @Nullable protected Map<String, List<DBObject>> readBsonDirectory(File directory) { final Map<String, List<DBObject>> collections = new HashMap<>(); final File[] collectionListing = directory.listFiles((dir, name) -> (name.endsWith(".bson") && !name.startsWith("system.indexes."))); if (collectionListing != null) { for (File collection : collectionListing) { List<DBObject> collectionData = readBsonFile(collection.getAbsolutePath()); collections.put(FilenameUtils.removeExtension(collection.getName()), collectionData); } } return collections; } @Override public Map<String, List<DBObject>> toMap() { return collectionMap; } }