importclass.py

# -*- coding: utf-8 -*-
# Class implementation for importing stuff to Neo4j
# based on Lightsocket examples by 
#James Thornton (http://jamesthornton.com)

from java.lang import Long, Double, Integer, String

from com.tinkerpop.rexster.util import ElementHelper
from com.tinkerpop.blueprints.pgm.impls.neo4j import Neo4jGraph
from com.tinkerpop.blueprints.pgm.impls.neo4jbatch import Neo4jBatchGraph

from com.tinkerpop.blueprints.pgm import Vertex, Edge
#from com.tinkerpop.pipes.util import FluentPipeline


from lightsocket.server import Resource, Response, Router

# all proxies are inherited from Resource, see server.py for base class
class VertexProxy(Resource):

	_type = "vertex"
	
	def __init__(self,graph):
		self.router = Router(self)
		self.graph = graph
		self.method_map = dict(create=self.create, 
			createInfo=self.createInfo,
			initializeElements=self.initializeElements)
		self.nodeTypes = dict()
		self.indices = dict()

		# which attributes to include to index for searching purposes?
		self.indexAttributes = [ 'label', 'alias', 'source', 'chr', 'start', 'end' ]

	def initializeElements(self, request):
		self.nodeTypes = request.data['nodeTypes']
		self.indices = request.data['indices']
		return Response(200)

	# create an info node containing barcodes and later possibly other stuff
	def createInfo(self,request):
		values = dict( request.data['values'] )
		nodeTypes = dict( request.data['types'] )

		vertex = self.graph.addVertex(None)
		for key,val in values.iteritems():
			valtype = nodeTypes.get( key )
			if( valtype is 'varchar' or valtype is 'longtext' ):
				vertex.setProperty( key, String( val ) )
				index.put( key, String( val ), vertex )
			elif( valtype is 'int' ):
				vertex.setProperty( key, Integer( val ) )
				index.put( key, String( val ), vertex )
			elif( valtype is 'double' ):
				vertex.setProperty( key, Double( val ) )
				index.put( key, String( val ), vertex )
		return Response(201)


	def create(self,request):
		values = dict( request.data['values'] )
		source = values.get('source')
		indexName = self.indices.get( source )
		index = self.graph.getIndex( indexName, Vertex )

		# let db determine suitable nodeID
		vertex = self.graph.addVertex(None)
		for key,val in values.iteritems():
			valtype = self.nodeTypes.get( key  )
			if( valtype is 'varchar' or valtype is 'longtext' ):
				vertex.setProperty( key, String( val ) )
			elif( valtype is 'int' ):
				vertex.setProperty( key, Integer( val ) )
			elif( valtype is 'double' ):
				vertex.setProperty( key, Double( val ) )

		# add only specific attributes to the indices
		for attribute in self.indexAttributes:
			valtype = self.nodeTypes.get( attribute )
			if( valtype is 'varchar' or valtype is 'longtext' ):
				index.put( attribute, String( values.get(attribute) ), vertex )
			if( valtype is 'int' ):
				index.put( attribute, Integer( values.get(attribute) ), vertex )
			if( valtype is 'double' ):
				index.put( attribute, Double( values.get(attribute) ), vertex )
		return Response(201, data={'id': vertex.getId() })

class EdgeProxy(Resource):

	_type = "edge"

	def __init__(self,graph):
		self.router = Router(self)
		self.graph = graph
		self.method_map = dict(create=self.create,
			initializeElements=self.initializeElements)

		# only one edge type
		self.label = 'ASSOCIATION'
		self.edgeTypes = dict()
		self.indices = dict()

		# default: index all attributes


	def initializeElements(self, request):
		self.edgeTypes = request.data['edgeTypes']
		self.indices = request.data['indices']
		return Response(200)

	def create(self,request):
		values = dict( request.data['values'] )
		sourceNode = self.graph.getVertex( Long( request.data['source'] ) )
		targetNode = self.graph.getVertex( Long( request.data['target'] ) )

		indexName = self.indices.get( 'ASSOCIATION' )
		indexNameDist = self.indices.get( 'DISTANCE' )
		index = self.graph.getIndex( indexName, Edge )
		indexDist = self.graph.getIndex( indexName, Edge )

		edge = self.graph.addEdge( None, sourceNode, targetNode, self.label )
		distance = values.get('distance')
		for key,val in values.iteritems():
			valtype = self.edgeTypes.get( key  )
			if( valtype is 'varchar' or valtype is 'longtext' ):
				edge.setProperty( key, String( val ) )
				index.put( key, String( val ), edge )
				if distance:
					indexDist.put( key, String( val ), edge )
			elif( valtype is 'int' ):
				edge.setProperty( key, Integer( val ) )
				index.put( key, Integer( val ), edge )
				if distance:
					indexDist.put( key, Integer( val ), edge )
			elif( valtype is 'double' ):
				edge.setProperty( key, Double( val ) )
				index.put( key, Double( val ), edge )
				if distance:
					indexDist.put( key, Double( val ), edge )
		return Response(201)

class IndexProxy(Resource):

	def __init__(self,graph):
		self.router = Router(self)
		self.graph = graph
		self.graph = graph
		self.method_map = dict(flush=self.flush,
							   create_manual=self.create_manual)

		# store created indexes
		self.indexes = dict()
		self.indexManager = self.graph.getRawGraph().index()

	def flush(self, request):
		# This is necessary prior to using indices to ensure that indexed data 
		# is available to index queries. This method is not part of the 
		# Blueprints Graph or IndexableGraph API.
		
		#self.graph.flushIndices()
		return Response(200)


	def create_manual(self, request):
		index_name = dict( request.data['name'] ).values()[0]
		index_config = { 'provider': 'lucene', 'type': 'fulltext' }
		index = None
		if request.data['type'] is 'node':
			index = self.indexManager.forNodes( index_name, index_config )
		elif request.data['type'] is 'edge':
			index = self.indexManager.forRelationships( index_name, index_config )

		self.indexes[ dict( request.data['name'] ).keys()[0] ] = index
		return Response(201) #,index.getPropertyMap())

class Gremlin(Resource):
	pass
	

class Blueprints(Resource):
	"""Abstract method"""

	def __init__(self,db_dir):
		self.db_dir = db_dir
		self._startup()
		self.router = Router(self)

		# provide this class to vertexproxy so it can ask the relevant indexes from indexproxy
		self.router.add("vertices", VertexProxy(self.graph))
		self.router.add("edges", EdgeProxy(self.graph))
		self.router.add("indices", IndexProxy(self.graph))
		#self.router.add("gremlin", Gremlin(self.graph))
		self.method_map = dict(startup=self.startup,
							   shutdown=self.shutdown)


	def getIndexProxy(self):
		return 

	def _startup(self):
		self.graph = self.graph_class(self.db_dir)

	def startup(self,request):
		self._startup()
		return Response(200)

	def shutdown(self,request):
		print "Shutting down graph..."
		self.graph.shutdown()
		return Response(200)
	#
	# Resource-specific methods
	#

	def clear(self):
		raise NotImplementedError


class Neo4jGraph(Blueprints):
	
	graph_class = Neo4jGraph


# use batch graph for faster import
class Neo4jBatchGraph(Blueprints):

	graph_class = Neo4jBatchGraph