package com.tinkerpop.blueprints.impls.neo4j;
import com.tinkerpop.blueprints.CloseableIterable;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Index;
import com.tinkerpop.blueprints.Parameter;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.util.StringFactory;
import org.neo4j.graphdb.Node;
import org.neo4j.graphdb.PropertyContainer;
import org.neo4j.graphdb.Relationship;
import org.neo4j.graphdb.index.IndexHits;
import org.neo4j.graphdb.index.IndexManager;
import java.util.HashMap;
import java.util.Map;
/**
* @author Marko A. Rodriguez (http://markorodriguez.com)
*/
public class Neo4jIndex<T extends Neo4jElement, S extends PropertyContainer> implements Index<T> {
private final Class<T> indexClass;
protected final Neo4jGraph graph;
private final String indexName;
protected org.neo4j.graphdb.index.Index<S> rawIndex;
protected Neo4jIndex(final String indexName, final Class<T> indexClass, final Neo4jGraph graph, final Parameter... indexParameters) {
this.indexClass = indexClass;
this.graph = graph;
this.indexName = indexName;
this.generateIndex(indexParameters);
}
public Class<T> getIndexClass() {
if (Vertex.class.isAssignableFrom(this.indexClass))
return (Class) Vertex.class;
else
return (Class) Edge.class;
}
public String getIndexName() {
return this.indexName;
}
public void put(final String key, final Object value, final T element) {
try {
this.graph.autoStartTransaction();
this.rawIndex.add((S) element.getRawElement(), key, value);
} catch (Exception e) {
throw new RuntimeException(e.getMessage(), e);
}
}
/**
* {@inheritDoc}
* <p/>
* The underlying Neo4j graph does not natively support this method within a transaction.
* If the graph is not currently in a transaction, then the operation runs efficiently.
* If the graph is in a transaction, then, for every element, a try/catch is used to determine if its in the current transaction.
*/
public CloseableIterable<T> get(final String key, final Object value) {
final IndexHits<S> itty = this.rawIndex.get(key, value);
if (this.indexClass.isAssignableFrom(Neo4jVertex.class))
return new Neo4jVertexIterable((Iterable<Node>) itty, this.graph, this.graph.checkElementsInTransaction());
else
return new Neo4jEdgeIterable((Iterable<Relationship>) itty, this.graph, this.graph.checkElementsInTransaction());
}
/**
* {@inheritDoc}
* <p/>
* The underlying Neo4j graph does not natively support this method within a transaction.
* If the graph is not currently in a transaction, then the operation runs efficiently.
* If the graph is in a transaction, then, for every element, a try/catch is used to determine if its in the current transaction.
*/
public CloseableIterable<T> query(final String key, final Object query) {
final IndexHits<S> itty = this.rawIndex.query(key, query);
if (this.indexClass.isAssignableFrom(Neo4jVertex.class))
return new Neo4jVertexIterable((Iterable<Node>) itty, this.graph, this.graph.checkElementsInTransaction());
else
return new Neo4jEdgeIterable((Iterable<Relationship>) itty, this.graph, this.graph.checkElementsInTransaction());
}
/**
* {@inheritDoc}
* <p/>
* The underlying Neo4j graph does not natively support this method within a transaction.
* If the graph is not currently in a transaction, then the operation runs efficiently.
* If the graph is in a transaction, then, for every element, a try/catch is used to determine if its in the current transaction.
*/
public long count(final String key, final Object value) {
if (!this.graph.checkElementsInTransaction()) {
final IndexHits hits = this.rawIndex.get(key, value);
final long count = hits.size();
hits.close();
return count;
} else {
final CloseableIterable<T> hits = this.get(key, value);
long count = 0;
for (final T t : hits) {
count++;
}
hits.close();
return count;
}
}
public void remove(final String key, final Object value, final T element) {
try {
this.graph.autoStartTransaction();
this.rawIndex.remove((S) element.getRawElement(), key, value);
} catch (Exception e) {
throw new RuntimeException(e.getMessage(), e);
}
}
private void generateIndex(final Parameter<Object, Object>... indexParameters) {
final IndexManager manager = this.graph.getRawGraph().index();
if (Vertex.class.isAssignableFrom(this.indexClass)) {
if (indexParameters.length > 0)
this.rawIndex = (org.neo4j.graphdb.index.Index<S>) manager.forNodes(this.indexName, generateParameterMap(indexParameters));
else
this.rawIndex = (org.neo4j.graphdb.index.Index<S>) manager.forNodes(this.indexName);
} else {
if (indexParameters.length > 0)
this.rawIndex = (org.neo4j.graphdb.index.Index<S>) manager.forRelationships(this.indexName, generateParameterMap(indexParameters));
else
this.rawIndex = (org.neo4j.graphdb.index.Index<S>) manager.forRelationships(this.indexName);
}
}
public String toString() {
return StringFactory.indexString(this);
}
private static Map<String, String> generateParameterMap(final Parameter<Object, Object>... indexParameters) {
final Map<String, String> map = new HashMap<String, String>();
for (final Parameter<Object, Object> parameter : indexParameters) {
map.put(parameter.getKey().toString(), parameter.getValue().toString());
}
return map;
}
public org.neo4j.graphdb.index.Index<S> getRawIndex() {
return this.rawIndex;
}
}