From 8d0690577e4b62bc98b53b36743fefb3e4425269 Mon Sep 17 00:00:00 2001 From: Matt Burgess Date: Sat, 6 May 2017 11:43:38 -0400 Subject: [PATCH] Add Jython support --- build.gradle | 5 +- .../impl/JythonScriptEngineConfigurator.java | 63 +++++++++++++++++++ .../services/javax.script.ScriptEngineFactory | 2 + .../nifi.script.ScriptEngineConfigurator | 1 + src/test/java/nifi/ScriptRunnerTest.java | 8 ++- src/test/resources/test_json2json.py | 28 +++++++++ 6 files changed, 104 insertions(+), 3 deletions(-) create mode 100644 src/main/java/nifi/script/impl/JythonScriptEngineConfigurator.java create mode 100644 src/main/resources/META-INF/services/javax.script.ScriptEngineFactory create mode 100644 src/test/resources/test_json2json.py diff --git a/build.gradle b/build.gradle index b075f73..f817768 100644 --- a/build.gradle +++ b/build.gradle @@ -23,7 +23,7 @@ plugins { id 'com.github.johnrengelman.shadow' version '1.2.3' } -version = '1.1.1' +version = '1.1.2' description = "A project to create a stub/mock environment for testing ExecuteScript processors" group = 'mattyb149' status = 'RELEASE' @@ -45,6 +45,7 @@ dependencies { // script engine(s) runtime 'org.codehaus.groovy:groovy-all:2.4.6' + runtime 'org.python:jython-standalone:2.7.0' } mainClassName = 'nifi.ScriptRunner' @@ -114,5 +115,5 @@ bintray { } task wrapper(type: Wrapper) { - gradleVersion = "2.13" + gradleVersion = "3.5" } diff --git a/src/main/java/nifi/script/impl/JythonScriptEngineConfigurator.java b/src/main/java/nifi/script/impl/JythonScriptEngineConfigurator.java new file mode 100644 index 0000000..b86312e --- /dev/null +++ b/src/main/java/nifi/script/impl/JythonScriptEngineConfigurator.java @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package nifi.script.impl; + +import org.apache.nifi.logging.ComponentLog; +import nifi.script.ScriptEngineConfigurator; + +import javax.script.ScriptEngine; +import javax.script.ScriptException; +import java.net.URL; + +/** + * A helper class to configure the Jython engine with any specific requirements + */ +public class JythonScriptEngineConfigurator implements ScriptEngineConfigurator { + + @Override + public String getScriptEngineName() { + return "python"; + } + + @Override + public URL[] getModuleURLsForClasspath(String[] modulePaths, ComponentLog log) { + // We don't need to add the module paths to the classpath, they will be added via sys.path.append + return new URL[0]; + } + + @Override + public Object init(ScriptEngine engine, String[] modulePaths) throws ScriptException { + return null; + } + + @Override + public Object eval(ScriptEngine engine, String scriptBody, String[] modulePaths) throws ScriptException { + Object returnValue = null; + if (engine != null) { + // Need to import the module path inside the engine, in order to pick up + // other Python/Jython modules + engine.eval("import sys"); + if (modulePaths != null) { + for (String modulePath : modulePaths) { + engine.eval("sys.path.append('" + modulePath + "')"); + } + } + returnValue = engine.eval(scriptBody); + } + return returnValue; + } +} diff --git a/src/main/resources/META-INF/services/javax.script.ScriptEngineFactory b/src/main/resources/META-INF/services/javax.script.ScriptEngineFactory new file mode 100644 index 0000000..2bbec81 --- /dev/null +++ b/src/main/resources/META-INF/services/javax.script.ScriptEngineFactory @@ -0,0 +1,2 @@ +org.codehaus.groovy.jsr223.GroovyScriptEngineFactory +org.python.jsr223.PyScriptEngineFactory \ No newline at end of file diff --git a/src/main/resources/META-INF/services/nifi.script.ScriptEngineConfigurator b/src/main/resources/META-INF/services/nifi.script.ScriptEngineConfigurator index e1f2c33..9293121 100644 --- a/src/main/resources/META-INF/services/nifi.script.ScriptEngineConfigurator +++ b/src/main/resources/META-INF/services/nifi.script.ScriptEngineConfigurator @@ -15,3 +15,4 @@ nifi.script.impl.GroovyScriptEngineConfigurator nifi.script.impl.JavascriptScriptEngineConfigurator +nifi.script.impl.JythonScriptEngineConfigurator diff --git a/src/test/java/nifi/ScriptRunnerTest.java b/src/test/java/nifi/ScriptRunnerTest.java index 87e6ce8..d0768c2 100644 --- a/src/test/java/nifi/ScriptRunnerTest.java +++ b/src/test/java/nifi/ScriptRunnerTest.java @@ -76,9 +76,15 @@ public void testOutputAllWithInputDir() throws Exception { } @Test - public void testReadWrite() throws Exception { + public void testReadWriteGroovy() throws Exception { System.setIn(new FileInputStream("src/test/resources/input_files/jolt.json")); ScriptRunner.main(new String[]{"-all", "src/test/resources/test_json2json.groovy"}); } + @Test + public void testReadWriteJython() throws Exception { + System.setIn(new FileInputStream("src/test/resources/input_files/jolt.json")); + ScriptRunner.main(new String[]{"-all", "src/test/resources/test_json2json.py"}); + } + } \ No newline at end of file diff --git a/src/test/resources/test_json2json.py b/src/test/resources/test_json2json.py new file mode 100644 index 0000000..138ec83 --- /dev/null +++ b/src/test/resources/test_json2json.py @@ -0,0 +1,28 @@ +import json +import java.io +from org.apache.commons.io import IOUtils +from java.nio.charset import StandardCharsets +from org.apache.nifi.processor.io import StreamCallback + +class PyStreamCallback(StreamCallback): + def __init__(self): + pass + def process(self, inputStream, outputStream): + text = IOUtils.toString(inputStream, StandardCharsets.UTF_8) + obj = json.loads(text) + newObj = { + "Range": 5, + "Rating": obj['rating']['primary']['value'], + "SecondaryRatings": {} + } + for key, value in obj['rating'].iteritems(): + if key != "primary": + newObj['SecondaryRatings'][key] = {"Id": key, "Range": 5, "Value": value['value']} + + outputStream.write(bytearray(json.dumps(newObj, indent=4).encode('utf-8'))) + +flowFile = session.get() +if (flowFile != None): + flowFile = session.write(flowFile,PyStreamCallback()) + flowFile = session.putAttribute(flowFile, "filename", flowFile.getAttribute('filename').split('.')[0]+'_translated.json') + session.transfer(flowFile, REL_SUCCESS) \ No newline at end of file