package com.tinkerpop.rexster.gremlin;
import com.codahale.metrics.Counter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Graph;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.util.io.graphson.GraphSONMode;
import com.tinkerpop.pipes.util.structures.Pair;
import com.tinkerpop.rexster.RexsterApplicationGraph;
import com.tinkerpop.rexster.RexsterResourceContext;
import com.tinkerpop.rexster.Tokens;
import com.tinkerpop.rexster.extension.AbstractRexsterExtension;
import com.tinkerpop.rexster.extension.ExtensionApi;
import com.tinkerpop.rexster.extension.ExtensionConfiguration;
import com.tinkerpop.rexster.extension.ExtensionDefinition;
import com.tinkerpop.rexster.extension.ExtensionDescriptor;
import com.tinkerpop.rexster.extension.ExtensionMethod;
import com.tinkerpop.rexster.extension.ExtensionNaming;
import com.tinkerpop.rexster.extension.ExtensionPoint;
import com.tinkerpop.rexster.extension.ExtensionRequestParameter;
import com.tinkerpop.rexster.extension.ExtensionResponse;
import com.tinkerpop.rexster.extension.HttpMethod;
import com.tinkerpop.rexster.extension.RexsterContext;
import com.tinkerpop.rexster.gremlin.converter.JSONResultConverter;
import com.tinkerpop.rexster.protocol.EngineController;
import com.tinkerpop.rexster.protocol.EngineHolder;
import com.tinkerpop.rexster.util.ElementHelper;
import com.tinkerpop.rexster.util.RequestObjectHelper;
import org.apache.commons.io.IOUtils;
import org.apache.log4j.Logger;
import org.codehaus.jettison.json.JSONArray;
import org.codehaus.jettison.json.JSONObject;
import javax.script.Bindings;
import javax.script.ScriptEngine;
import javax.script.ScriptException;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
/**
* @author Stephen Mallette (http://stephen.genoprime.com)
*/
@ExtensionNaming(namespace = GremlinExtension.EXTENSION_NAMESPACE, name = GremlinExtension.EXTENSION_NAME)
public class GremlinExtension extends AbstractRexsterExtension {
protected static final Logger logger = Logger.getLogger(GremlinExtension.class);
public static final String EXTENSION_NAMESPACE = "tp";
public static final String EXTENSION_NAME = "gremlin";
private static final ConcurrentMap<String, String> cachedScripts = new ConcurrentHashMap<String, String>();
private static final String GRAPH_VARIABLE = "g";
private static final String VERTEX_VARIABLE = "v";
private static final String EDGE_VARIABLE = "e";
private static final String WILDCARD = "*";
private static final String SCRIPT = "script";
private static final String LANGUAGE = "language";
private static final String PARAMS = "params";
private static final String LOAD = "load";
private static final String RETURN_TOTAL = "returnTotal";
private static final String API_SHOW_TYPES = "displays the properties of the elements with their native data type (default is false)";
private static final String API_SCRIPT = "the Gremlin script to be evaluated";
private static final String API_RETURN_KEYS = "an array of element property keys to return (default is to return all element properties)";
private static final String API_START_OFFSET = "start index for a paged set of data to be returned";
private static final String API_END_OFFSET = "end index for a paged set of data to be returned";
private static final String API_LANGUAGE = "the gremlin language flavor to use (default is groovy)";
private static final String API_PARAMS = "a map of parameters to bind to the script engine";
private static final String API_LOAD = "a list of 'stored procedures' to execute prior to the 'script' (if 'script' is not specified then the last script in this argument will return the values";
private static final String API_RETURN_TOTAL = "when set to true, the full result set will be iterated and the results returned (default is false)";
@ExtensionDefinition(extensionPoint = ExtensionPoint.EDGE, method = HttpMethod.GET)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for an edge.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluateGetOnEdge(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@RexsterContext Edge edge,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, null, edge, script);
}
@ExtensionDefinition(extensionPoint = ExtensionPoint.EDGE, method = HttpMethod.POST)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for an edge.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluatePostOnEdge(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@RexsterContext Edge edge,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, null, edge, script);
}
@ExtensionDefinition(extensionPoint = ExtensionPoint.VERTEX, method = HttpMethod.GET)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for a vertex.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluateGetOnVertex(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@RexsterContext Vertex vertex,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, vertex, null, script);
}
@ExtensionDefinition(extensionPoint = ExtensionPoint.VERTEX, method = HttpMethod.POST)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for a vertex.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluatePostOnVertex(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@RexsterContext Vertex vertex,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, vertex, null, script);
}
@ExtensionDefinition(extensionPoint = ExtensionPoint.GRAPH, method = HttpMethod.GET)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for a graph.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluateGetOnGraph(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, null, null, script);
}
@ExtensionDefinition(extensionPoint = ExtensionPoint.GRAPH, method = HttpMethod.POST)
@ExtensionDescriptor(description = "evaluate an ad-hoc Gremlin script for a graph.",
api = {
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.SHOW_TYPES, description = API_SHOW_TYPES),
@ExtensionApi(parameterName = LANGUAGE, description = API_LANGUAGE),
@ExtensionApi(parameterName = PARAMS, description = API_PARAMS),
@ExtensionApi(parameterName = LOAD, description = API_LOAD),
@ExtensionApi(parameterName = RETURN_TOTAL, description = API_RETURN_TOTAL),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.RETURN_KEYS, description = API_RETURN_KEYS),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_START, description = API_START_OFFSET),
@ExtensionApi(parameterName = Tokens.REXSTER + "." + Tokens.OFFSET_END, description = API_END_OFFSET)
})
public ExtensionResponse evaluatePostOnGraph(@RexsterContext RexsterResourceContext rexsterResourceContext,
@RexsterContext Graph graph,
@ExtensionRequestParameter(name = SCRIPT, description = API_SCRIPT, parseToJson = false) String script) {
return tryExecuteGremlinScript(rexsterResourceContext, graph, null, null, script);
}
private ExtensionResponse tryExecuteGremlinScript(final RexsterResourceContext rexsterResourceContext,
final Graph graph, final Vertex vertex, final Edge edge,
final String script) {
final MetricRegistry metricRegistry = rexsterResourceContext.getMetricRegistry();
final Timer scriptTimer = metricRegistry.timer(MetricRegistry.name("http", "script-engine"));
final Counter successfulExecutions = metricRegistry.counter(MetricRegistry.name("http", "script-engine", "success"));
final Counter failedExecutions = metricRegistry.counter(MetricRegistry.name("http", "script-engine", "fail"));
ExtensionResponse extensionResponse;
final JSONObject requestObject = rexsterResourceContext.getRequestObject();
// can't initialize this statically because the configure() method won't get called before it.
// need to think a bit on how to best initialized the controller.
final EngineController engineController = EngineController.getInstance();
final boolean showTypes = RequestObjectHelper.getShowTypes(requestObject);
final long offsetStart = RequestObjectHelper.getStartOffset(requestObject);
final long offsetEnd = RequestObjectHelper.getEndOffset(requestObject);
final boolean returnTotal = getReturnTotal(requestObject);
final GraphSONMode mode = showTypes ? GraphSONMode.EXTENDED : GraphSONMode.NORMAL;
final Set<String> returnKeys = RequestObjectHelper.getReturnKeys(requestObject, WILDCARD);
final String languageToExecuteWith = getLanguageToExecuteWith(requestObject);
final EngineHolder engineHolder;
final ScriptEngine scriptEngine;
try {
if (!engineController.isEngineAvailable(languageToExecuteWith)) {
return ExtensionResponse.error("language requested is not available on the server");
}
engineHolder = engineController.getEngineByLanguageName(languageToExecuteWith);
scriptEngine = engineHolder.getEngine();
} catch (ScriptException se) {
return ExtensionResponse.error("could not get request script engine");
}
final Bindings bindings = createBindings(graph, vertex, edge, scriptEngine);
// add all keys not defined by this request as bindings to the script engine
placeParametersOnBinding(requestObject, bindings, showTypes);
// get the list of "stored procedures" to run
final RexsterApplicationGraph rag = rexsterResourceContext.getRexsterApplicationGraph();
final ExtensionMethod extensionMethod = rexsterResourceContext.getExtensionMethod();
Map configurationMap = null;
Iterator<String> scriptsToRun = null;
try {
final ExtensionConfiguration extensionConfiguration = rag != null ? rag.findExtensionConfiguration(EXTENSION_NAMESPACE, EXTENSION_NAME) : null;
if (extensionConfiguration != null) {
configurationMap = extensionConfiguration.tryGetMapFromConfiguration();
scriptsToRun = getScriptsToRun(requestObject, configurationMap);
}
} catch (IOException ioe) {
return ExtensionResponse.error(ioe,
generateErrorJson(extensionMethod.getExtensionApiAsJson()));
}
if ((script == null || script.isEmpty()) && scriptsToRun == null) {
return ExtensionResponse.badRequest(
"no scripts provided",
generateErrorJson(extensionMethod.getExtensionApiAsJson()));
}
final Timer.Context context = scriptTimer.time();
try {
// result is either the ad-hoc script on the query string or the last "stored procedure"
Object result = null;
if (scriptsToRun != null) {
while (scriptsToRun.hasNext()) {
result = engineHolder.getEngine().eval(scriptsToRun.next(), bindings);
}
}
if (isClientScriptAllowed(configurationMap) && script != null && !script.isEmpty()) {
result = engineHolder.getEngine().eval(script, bindings);
}
final Pair<JSONArray, Long> convertedResults = new JSONResultConverter(mode, offsetStart, offsetEnd, returnKeys).convert(result, returnTotal);
final JSONArray results = convertedResults.getA();
final HashMap<String, Object> resultMap = new HashMap<String, Object>();
resultMap.put(Tokens.SUCCESS, true);
resultMap.put(Tokens.RESULTS, results);
if (returnTotal)
resultMap.put(Tokens.COUNT, convertedResults.getB());
final JSONObject resultObject = new JSONObject(resultMap);
extensionResponse = ExtensionResponse.ok(resultObject);
successfulExecutions.inc();
} catch (Exception e) {
logger.error(String.format("Gremlin Extension: %s", e.getMessage()), e);
extensionResponse = ExtensionResponse.error(e,
generateErrorJson(extensionMethod.getExtensionApiAsJson()));
failedExecutions.inc();
} finally {
context.stop();
}
return extensionResponse;
}
private static Bindings createBindings(final Graph graph, final Vertex vertex, final Edge edge,
final ScriptEngine scriptEngine) {
final Bindings bindings = scriptEngine.createBindings();
bindings.put(GRAPH_VARIABLE, graph);
if (vertex != null) {
bindings.put(VERTEX_VARIABLE, vertex);
}
if (edge != null) {
bindings.put(EDGE_VARIABLE, edge);
}
return bindings;
}
/*
private static JSONObject getBindingsAsJson(final Bindings bindings) throws Exception{
final HashMap<String, Object> bindingJsonValues = new HashMap<String, Object>();
for (String key : bindings.keySet()) {
if (!key.equals(Tokens.REXSTER) && !key.equals(LANGUAGE) && !key.equals(SCRIPT)
&& !key.equals(GRAPH_VARIABLE) && !key.equals(EDGE_VARIABLE) && !key.equals(VERTEX_VARIABLE)) {
bindingJsonValues.put(key, bindings.get(key));
}
}
JSONObject bindingJson = null;
if (!bindingJsonValues.isEmpty()) {
bindingJson = new JSONObject(bindingJsonValues);
}
return bindingJson;
}
*/
private static void placeParametersOnBinding(final JSONObject requestObject, final Bindings bindings, final boolean parseTypes) {
if (requestObject != null) {
JSONObject paramMap = requestObject.optJSONObject(PARAMS);
if (paramMap != null) {
final Iterator keyIterator = paramMap.keys();
while (keyIterator.hasNext()) {
final String key = (String) keyIterator.next();
bindings.put(key, ElementHelper.getTypedPropertyValue(paramMap.opt(key), parseTypes));
}
}
}
}
private static String getLanguageToExecuteWith(final JSONObject requestObject) {
final String language = requestObject != null ? requestObject.optString(LANGUAGE) : null;
String requestedLanguage = "groovy";
if (language != null && !language.equals("")) {
requestedLanguage = language;
}
return requestedLanguage;
}
private static boolean getReturnTotal(final JSONObject requestObject) {
final String retTotalString = requestObject != null ? requestObject.optString(RETURN_TOTAL) : null;
boolean returnTotal = false;
if (retTotalString != null && !retTotalString.equals("")) {
returnTotal = Boolean.parseBoolean(retTotalString);
}
return returnTotal;
}
private static Iterator<String> getScriptsToRun(final JSONObject requestObject, final Map configuration) throws IOException {
if (configuration == null) {
logger.warn("No scripts are configured for the Gremlin Extension so 'load' query string parameter will be ignored");
return null;
}
if (!configuration.containsKey("scripts")) {
logger.warn("The configuration suppled for the Gremlin Extension does not contain a 'scripts' key so 'load' query string parameter will be ignored");
return null;
}
boolean scriptsAreCached = areScriptsCached(configuration);
String scriptLocation = (String) configuration.get("scripts");
final JSONArray jsonArray = requestObject != null ? requestObject.optJSONArray(LOAD) : null;
Iterator<String> scripts = null;
if (jsonArray != null) {
List<String> scriptList = new ArrayList<String>();
for (int ix = 0; ix < jsonArray.length(); ix++) {
final String locationAndScriptFile = scriptLocation + File.separator + jsonArray.optString(ix) + ".gremlin";
String script = cachedScripts.get(locationAndScriptFile);
if (script == null) {
script = readFile(locationAndScriptFile);
if (scriptsAreCached) {
synchronized (GremlinExtension.class) {
cachedScripts.putIfAbsent(locationAndScriptFile, script);
}
}
}
scriptList.add(script);
}
scripts = scriptList.iterator();
}
return scripts;
}
private static synchronized String readFile(final String fileName) throws IOException {
final StringWriter stringWriter = new StringWriter();
IOUtils.copy(new FileInputStream(new File(fileName)), stringWriter);
return stringWriter.toString();
}
private static boolean isClientScriptAllowed(final Map configuration) {
boolean allowClientScript = true;
if (configuration != null && configuration.containsKey("allow-client-script")) {
final String configValue = (String) configuration.get("allow-client-script");
allowClientScript = configValue.toLowerCase().equals("true") ? true : false;
}
return allowClientScript;
}
private static boolean areScriptsCached(final Map configuration) {
boolean cacheScripts = true;
if (configuration != null && configuration.containsKey("cache-scripts")) {
final String configValue = (String) configuration.get("cache-scripts");
cacheScripts = configValue.toLowerCase().equals("true") ? true : false;
}
return cacheScripts;
}
}