/*
* Copyright (c) 2007-2013 The Broad Institute, Inc.
* SOFTWARE COPYRIGHT NOTICE
* This software and its documentation are the copyright of the Broad Institute, Inc. All rights are reserved.
*
* This software is supplied without any warranty or guaranteed support whatsoever. The Broad Institute is not responsible for its use, misuse, or functionality.
*
* This software is licensed under the terms of the GNU Lesser General Public License (LGPL),
* Version 2.1 which is available at http://www.opensource.org/licenses/lgpl-2.1.php.
*/
package org.broad.igv.plugin.mongocollab;
import com.mongodb.*;
import org.apache.log4j.Logger;
import org.broad.igv.dev.api.IGVPlugin;
import org.broad.igv.dev.api.LoadHandler;
import org.broad.igv.feature.Locus;
import org.broad.igv.feature.genome.GenomeManager;
import org.broad.igv.feature.tribble.CodecFactory;
import org.broad.igv.track.Track;
import org.broad.igv.track.TrackLoader;
import org.broad.igv.ui.color.ColorUtilities;
import org.broad.igv.ui.util.FileDialogUtils;
import org.broad.igv.ui.util.MessageUtils;
import org.broad.igv.util.ParsingUtils;
import htsjdk.tribble.AbstractFeatureReader;
import htsjdk.tribble.Feature;
import htsjdk.tribble.FeatureCodec;
import org.bson.BSON;
import org.bson.Transformer;
import java.awt.*;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.UnknownHostException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
/**
* @author jacob
* @date 2013-Sep-06
*/
public class MongoCollabPlugin implements IGVPlugin {
private static Logger log = Logger.getLogger(MongoCollabPlugin.class);
static{
BSON.addEncodingHook(Color.class, new Transformer() {
@Override
public Object transform(Object o) {
if(o instanceof Color){
return ColorUtilities.colorToString((Color) o);
}else{
return o;
}
}
});
}
@Override
public void init() {
TrackLoader.registerHandler("db.spec", new TrackLoadHandler());
}
static void insertFeaturesFromFile(DBCollection collection){
File featFile = FileDialogUtils.chooseFile("Select feature file");
if (featFile != null) {
String path = featFile.getAbsolutePath();
insertFeaturesFromFile(collection, path);
}
}
/**
*
* @param collection
* @param path
* @return The number of features inserted
*/
static int insertFeaturesFromFile(DBCollection collection, String path) {
FeatureCodec codec = CodecFactory.getCodec(path, GenomeManager.getInstance().getCurrentGenome());
if (codec != null) {
AbstractFeatureReader<Feature, ?> bfs = AbstractFeatureReader.getFeatureReader(path, codec, false);
Iterable<Feature> iter = null;
try {
iter = bfs.iterator();
} catch (IOException ex) {
log.error(ex.getMessage(), ex);
throw new RuntimeException("Error reading file: " + path, ex);
}
int count = 0;
for (Feature feat : iter) {
String err = saveFeature(collection, DBFeature.create(feat));
if(err == null){
count += 1;
}else{
log.error("Error inserting feature: " + err);
}
}
return count;
} else {
throw new RuntimeException("Cannot load features from file of this type");
}
}
/**
* Save the specified FeatDBObject to the specified collection
* Does either an insert or update
*
* @param collection
* @param dbFeat
* @return
*/
static String saveFeature(DBCollection collection, DBFeature dbFeat) {
String errorMessage = "";
try {
if(log.isDebugEnabled()){
log.debug("Saving feature " + Locus.getFormattedLocusString(dbFeat.getChr(), dbFeat.getStart(), dbFeat.getEnd()));
}
WriteResult wr = collection.save(dbFeat);
errorMessage = wr.getError();
} catch (Exception ex) {
errorMessage = ex.getMessage();
log.error(errorMessage, ex);
if (errorMessage == null) errorMessage = "" + ex;
}
return errorMessage;
}
public static void removeFeature(DBCollection collection, DBFeature featDBObject) {
collection.remove(featDBObject);
}
private static Map<String, Mongo> connections = new HashMap<String, Mongo>();
static Mongo getMongo(String host, int port){
String key = createConnString(host, port);
Mongo connection = connections.get(key);
if(connection == null){
try {
log.info("Connecting to MongoDB host=" + host + " port=" + port);
connection = new MongoClient(host, port);
} catch (UnknownHostException e) {
log.error(e.getMessage(), e);
throw new RuntimeException(e.getMessage(), e);
}
connections.put(key, connection);
}
return connection;
}
static void closeMongo(String host, int port) {
String key = createConnString(host, port);
Mongo connection = connections.get(key);
if(connection != null){
log.info("Closing connection to MongoDB host=" + host + " port=" + port);
connection.close();
connections.remove(key);
}
}
private static String createConnString(String host, int port){
return host + ":" + port;
}
static DBCollection getCollection(Locator locator) {
Mongo mongo = getMongo(locator.host, locator.port);
DB mongoDB = mongo.getDB(locator.dbName);
return mongoDB.getCollection(locator.collectionName);
}
public static class Locator {
public final String host;
public final int port;
public final String dbName;
public final String collectionName;
public final boolean buildLocusIndex;
public Locator(String path) throws IOException{
this(ParsingUtils.openInputStream(path));
}
public Locator(InputStream is){
Map<String, String> fields = ParsingUtils.loadMap(is);
this.host = fields.get("host");
this.port = Integer.parseInt(fields.get("port"));
this.dbName = fields.get("dbName");
this.collectionName = fields.get("collectionName");
boolean tmpBuildIndex = false;
if(fields.containsKey("buildLocusIndex")){
tmpBuildIndex = Boolean.parseBoolean(fields.get("buildLocusIndex"));
}
this.buildLocusIndex = tmpBuildIndex;
}
}
public static final int DB_EXISTS = 0x1;
public static final int COLLECTION_EXISTS = 0x2;
/**
* Check whether the database/collection specified by the given {@code locator}
* exists.
* @param locator
* @return An integer consisting of 2 flags:
* DB_EXISTS
* COLLECTION_EXISTS
*
* with each set if the DB/COLLECTION exists. !DB_EXISTS && COLLECTION_EXISTS should never happen
*/
public static int checkDestinationExists(Locator locator){
Mongo mongo = getMongo(locator.host, locator.port);
List<String> dbNames = mongo.getDatabaseNames();
boolean dbExists = dbNames.indexOf(locator.dbName) >= 0;
if(!dbExists){
return 0;
}
int result = DB_EXISTS;
DB db = mongo.getDB(locator.dbName);
Set<String> collections = db.getCollectionNames();
result |= collections.contains(locator.collectionName) ? COLLECTION_EXISTS : 0;
return result;
}
public static class TrackLoadHandler implements LoadHandler {
@Override
public void load(String path, List<Track> newTracks) throws IOException{
Locator locator = new Locator(path);
int destExists = checkDestinationExists(locator);
String toAsk = null;
boolean doLoadTrack = false;
if( (destExists & COLLECTION_EXISTS) > 0){
//Both DB and collection exist
assert (destExists & DB_EXISTS) > 0;
doLoadTrack = true;
}else if(destExists == 0){
//Neither collection nor track exist
toAsk = String.format("Host '%s' does not contain database '%s'. Do you wish to create it?\n", locator.host, locator.dbName);
toAsk += String.format("If you select yes, collection '%s' will be created as well.", locator.collectionName);
}else{
//Collection doesn't exist but DB does
toAsk = String.format("Host '%s', database '%s', does not contain collection '%s'.\nDo you wish to create it?",
locator.host, locator.dbName, locator.collectionName);
}
if(toAsk != null){
doLoadTrack = MessageUtils.confirm(toAsk);
}
if(doLoadTrack) MongoFeatureSource.loadFeatureTrack(locator, newTracks);
}
}
}