src/main/java/net/onrc/onos/util/GraphDBConnection.java

package net.onrc.onos.util;

import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.thinkaurelius.titan.core.TitanFactory;
import com.thinkaurelius.titan.core.TitanGraph;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.TransactionalGraph;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.util.wrappers.event.EventTransactionalGraph;
import com.tinkerpop.frames.FramedGraph;

public class GraphDBConnection {
	public enum Transaction {
		COMMIT,
		ROLLBACK
	}
	public enum GenerateEvent {
		TRUE,
		FALSE
	}
	class TransactionHandle {
		protected TransactionalGraph tr;
		public void create() {
			tr = graph.newTransaction();			
		}
	}
	protected static Logger log = LoggerFactory.getLogger(GraphDBConnection.class);
	private static GraphDBConnection singleton = new GraphDBConnection( );
	private static TitanGraph graph;
	private static EventTransactionalGraph<TitanGraph> eg;
	private static GraphDBUtils utils;
	private static String configFile;

	   
	   /* A private Constructor prevents any other 
	    * class from instantiating.
	    */
	   private GraphDBConnection(){ }
	   
	   /* Static 'instance' method */
	   public static synchronized GraphDBConnection getInstance(final String conf) {
		   if (GraphDBConnection.configFile == null || GraphDBConnection.configFile.isEmpty()) {
			   GraphDBConnection.configFile = conf;
			   log.debug("GraphDBConnection::Setting Config File {}", GraphDBConnection.configFile);
		   }
		   if (!GraphDBConnection.configFile.isEmpty() && 
				   (graph == null||graph.isOpen() == Boolean.FALSE)) {
		        graph = TitanFactory.open(GraphDBConnection.configFile);		        
		        // FIXME: Creation on Indexes should be done only once
		        Set<String> s_vertices = graph.getIndexedKeys(Vertex.class);
		        if (!s_vertices.contains("dpid")) {
		           graph.createKeyIndex("dpid", Vertex.class);
		        }
		        if (!s_vertices.contains("type")) {
		        	graph.createKeyIndex("type", Vertex.class);
		        }
		        if (!s_vertices.contains("dl_address")) {
		        	graph.createKeyIndex("dl_address", Vertex.class);
		        }
		        if (!s_vertices.contains("flow_id")) {
		        	graph.createKeyIndex("flow_id", Vertex.class);
		        }
		        if (!s_vertices.contains("flow_entry_id")) {
		        	graph.createKeyIndex("flow_entry_id",
						     Vertex.class);
		        }
		        if (!s_vertices.contains("switch_state")) {
		        	graph.createKeyIndex("switch_state",
						     Vertex.class);
		        }
		        
		        // BEGIN: trial code
		        /*
		         * Trial to store link state information as properties of port.
		         * Currently no need to be implemented. Just reference.
		         */
		        // These keys are assigned to port vertex with "OUT" direction link
//		        if (!s_vertices.contains("first_seen_time")) {
//		        	graph.createKeyIndex("first_seen_time", Vertex.class);
//		        }
//		        if (!s_vertices.contains("last_lldp_received_time")) {
//		        	graph.createKeyIndex("last_lldp_received_time", Vertex.class);
//		        }
//		        if (!s_vertices.contains("last_bddp_received_time")) {
//		        	graph.createKeyIndex("last_bddp_received_time", Vertex.class);
//		        }
//		        
//		        
//		        Set<String> s_edges = graph.getIndexedKeys(Edge.class);
//		        if (!s_edges.contains("src_port_state")) {
//		        	graph.createKeyIndex("src_port_state",
//						     Edge.class);
//		        }
//		        if (!s_edges.contains("dst_port_state")) {
//		        	graph.createKeyIndex("dst_port_state",
//						     Edge.class);
//		        }
//		        if (!s_edges.contains("first_seen_time")) {
//		        	graph.createKeyIndex("first_seen_time", Edge.class);
//		        }
//		        if (!s_edges.contains("last_lldp_received_time")) {
//		        	graph.createKeyIndex("last_lldp_received_time", Edge.class);
//		        }
//		        if (!s_edges.contains("last_bddp_received_time")) {
//		        	graph.createKeyIndex("last_bddp_received_time", Edge.class);
//		        }
		        // END: trial code
		        
		        graph.commit();
		        eg = new EventTransactionalGraph<TitanGraph>(graph);
		   }		   
		   if (utils == null) {
			   utils = new GraphDBUtils();
		   }
	      return singleton;
	   }
	   
	   public IDBUtils utils() {
		   return utils;
	   }
	   
	   public FramedGraph<TitanGraph> getFramedGraph() {
	   
		   	if (isValid()) {
		   		FramedGraph<TitanGraph> fg = new FramedGraph<TitanGraph>(graph);
		   		return fg;
		   	} else {
		   		log.error("new FramedGraph failed");
		   		return null;
		   	}
	   }
	   
	   protected EventTransactionalGraph<TitanGraph> getEventGraph() {
		   
		   	if (isValid()) {		   		
		   		return eg;
		   	} else {
		   		return null;
		   	}
	   }
	   
	   public void addEventListener(final LocalGraphChangedListener listener) {		   
		   EventTransactionalGraph<TitanGraph> eg = this.getEventGraph();
		   eg.addListener(listener);
		   log.debug("Registered listener {}",listener.getClass());
	   }
	   
	public Boolean isValid() {
		   
		   return (graph != null||graph.isOpen());
	   }
	   
	public void startTx() {
		   
		   
	   }
	   
	public void endTx(Transaction tx) {
		   try {
			   switch (tx) {
			   case COMMIT:
				   graph.commit();
			   case ROLLBACK:
				   graph.rollback();
			   }
		   } catch (Exception e) {
			   // TODO Auto-generated catch block
			   log.error("{}",e.toString());
		   }
	   }
	   
	   public void endTx(TransactionHandle tr, Transaction tx) {
		   switch (tx) {
		   case COMMIT:
			   if (tr != null && tr.tr != null) {
				   tr.tr.commit();
			   } else {
				   graph.commit();
			   }
		   case ROLLBACK:
			   if (tr != null && tr.tr != null) {
				   tr.tr.rollback();
			   } else {
				   graph.rollback();
			   }
		   }
	   }   
	   
	   public void endTx(Transaction tx, GenerateEvent fire) {

		   try {
			if (fire.equals(GenerateEvent.TRUE)) {
				   switch (tx) {
				   case COMMIT:
					   eg.commit();
				   case ROLLBACK:
					   eg.rollback();
				   }
			   } else {
					endTx(tx);   			   
			   }
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	   }
	   
	public void close() {
		   endTx(Transaction.COMMIT);
	   }	   
}