Java Code Examples for org.apache.pig.PigServer#registerScript()

@Test
public void testImport() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "import");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("import 'piggybank.jar';")
                .append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("run evil.pig;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "Error during parsing. IMPORT command is not permitted. ");
    }
}
 

@Test
public void testRun() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "run");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("run evil.pig;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "RUN command is not permitted. ");
    }
}
 

/**
 * Verify that ASSERT operator works in a Pig script
 * See PIG-3670
 * @throws Exception
 */
@Test
public void testInScript() throws Exception {
    PigServer pigServer = new PigServer(ExecType.LOCAL);
    Data data = resetData(pigServer);

    data.set("foo",
            tuple(1),
            tuple(2),
            tuple(3)
            );

    StringBuffer query = new StringBuffer();
    query.append("A = LOAD 'foo' USING mock.Storage() AS (i:int);\n");
    query.append("ASSERT A BY i > 0;\n");
    query.append("STORE A INTO 'bar' USING mock.Storage();");

    InputStream is = new ByteArrayInputStream(query.toString().getBytes());
    pigServer.registerScript(is);

    List<Tuple> out = data.get("bar");
    assertEquals(3, out.size());
    assertEquals(tuple(1), out.get(0));
    assertEquals(tuple(2), out.get(1));
    assertEquals(tuple(3), out.get(2));
}
 

/**
 * Tests the blacklist filter. We blacklist "set" and make sure this test
 * throws a {@link FrontendException}
 * 
 * @throws Exception
 */
@Test
public void testBlacklist() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "set");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("set io.sort.mb 1000;")
                .append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "SET command is not permitted. ");
    }
}
 

@Test
public void testExec() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "exec");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("exec evil.pig;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "EXEC command is not permitted. ");
    }
}
 

@Test
public void testPreprocessorCommands2() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "dEfaUlt");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("set io.sort.mb 1000;")
        .append("%Default input 'foo';")
                .append("A = LOAD '$input' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        // We check RuntimeException here and not FrontendException as Pig wraps the error from Preprocessor
        // within RuntimeException
        Util.assertExceptionAndMessage(RuntimeException.class, e,
                "DEFAULT command is not permitted. ");
    }
}
 

@Test
public void testPreprocessorCommand3() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "Define");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("set io.sort.mb 1000;")
        .append("DEFINE UrlDecode InvokeForString('java.net.URLDecoder.decode', 'String String');  ")
                .append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "Error during parsing. DEFINE command is not permitted. ");
    }
}
 

@Test
public void testExplain() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "explain");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("A = LOAD 'foo' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("EXPLAIN B;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        Util.assertExceptionAndMessage(FrontendException.class, e,
                "EXPLAIN command is not permitted. ");
    }
}
 

@Test
public void testUDFReturnMap_LocalMode() throws Exception {
    PigServer pig = new PigServer(ExecType.LOCAL);
    pig.registerScript(TempScriptFile.getAbsolutePath());

    Iterator<Tuple> iterator = pig.openIterator("B");
    while (iterator.hasNext()) {
        Tuple tuple = iterator.next();
        @SuppressWarnings("unchecked")
        Map<Object, Object> result = (Map<Object, Object>) tuple.get(0);
        assertEquals(result, MyUDFReturnMap.map);
    }
}
 

private PigServer setup(String script, String endpoint) throws Exception {
    Configuration conf = new HdfsConfiguration();
    Map<String, String> parameters = new HashMap<>();
    parameters.put("ENDPOINT", endpoint);

    PigServer ps = new PigServer(ExecType.LOCAL, conf);
    ps.setBatchOn();
    ps.registerScript(script, parameters);

    return ps;
}
 

@Test
public void testUDFContext() throws Exception {
    File a = Util.createLocalInputFile("a.txt", new String[] { "dumb" });
    File b = Util.createLocalInputFile("b.txt", new String[] { "dumber" });
    FileLocalizer.deleteTempFiles();
    PigServer pig = new PigServer(ExecType.LOCAL, new Properties());
    String[] statement = { "A = LOAD '" + Util.encodeEscape(a.getAbsolutePath()) +
            "' USING org.apache.pig.test.utils.UDFContextTestLoader('joe');",
        "B = LOAD '" + Util.encodeEscape(b.getAbsolutePath()) +
        "' USING org.apache.pig.test.utils.UDFContextTestLoader('jane');",
        "C = union A, B;",
        "D = FOREACH C GENERATE $0, $1, org.apache.pig.test.utils.UDFContextTestEvalFunc($0), " +
        "org.apache.pig.test.utils.UDFContextTestEvalFunc2($0);" };

    File tmpFile = File.createTempFile("temp_jira_851", ".pig");
    FileWriter writer = new FileWriter(tmpFile);
    for (String line : statement) {
        writer.write(line + "\n");
    }
    writer.close();

    pig.registerScript(tmpFile.getAbsolutePath());
    Iterator<Tuple> iterator = pig.openIterator("D");
    while (iterator.hasNext()) {
        Tuple tuple = iterator.next();
        if ("dumb".equals(tuple.get(0).toString())) {
            assertEquals(tuple.get(1).toString(), "joe");
        } else if ("dumber".equals(tuple.get(0).toString())) {
            assertEquals(tuple.get(1).toString(), "jane");
        }
    	assertEquals(Integer.valueOf(tuple.get(2).toString()), new Integer(5));
    	assertEquals(tuple.get(3).toString(), "five");
    }
}
 

public void testRegisteredJarVisibility(PigServer pigServer, String inputPath) throws IOException {
    String query = "register " + jarFile.getAbsolutePath() + ";\n"
            + "a = load '" + Util.generateURI(inputPath, pigServer.getPigContext())
            + "' using org.apache.pig.test.RegisteredJarVisibilityLoader();\n"
            // register again to test classloader consistency
            + "register " +  jarFile.getAbsolutePath() + ";\n"
            + "b = load 'non_existent' "
            + "using org.apache.pig.test.RegisteredJarVisibilityLoader();";

    LOG.info("Running pig script:\n" + query);
    pigServer.registerScript(new ByteArrayInputStream(query.getBytes()));

    pigServer.openIterator("a");
    pigServer.shutdown();
}
 

/**
 * A few commands such as DECLARE, DEFAULT go via
 * {@link PreprocessorContext}. This step basically parses commands and
 * substitutes parameters. The parameters can be evaluated using shell
 * commands, which need to validated if specified in the white or blacklist.
 * This test handles that scenario
 * 
 * @throws Exception
 */
@Test
public void testPreprocessorCommands() throws Exception {
    try {
        ctx.getProperties().setProperty(PigConfiguration.PIG_BLACKLIST, "dEclAre");
        PigServer pigServer = new PigServer(ctx);
        Data data = resetData(pigServer);

        data.set("foo", tuple("a", 1, "b"), tuple("b", 2, "c"),
                tuple("c", 3, "d"));

        StringBuilder script = new StringBuilder();
        script.append("set io.sort.mb 1000;")
        .append("%declare X `echo`; ")
        .append("%default input 'foo';")
                .append("A = LOAD '$input' USING mock.Storage() AS (f1:chararray,f2:int,f3:chararray);")
                .append("B = order A by f1,f2,f3 DESC;")
                .append("STORE B INTO 'bar' USING mock.Storage();");

        pigServer.registerScript(IOUtils.toInputStream(script));
        fail();
    } catch (Exception e) {
        // We check RuntimeException here and not FrontendException as Pig wraps the error from Preprocessor
        // within RuntimeException
        Util.assertExceptionAndMessage(RuntimeException.class, e,
                "DECLARE command is not permitted. ");
    }
}
 

@Test
public void testShellCommandOrder() throws Throwable {
    PigServer server = new PigServer(ExecType.LOCAL, new Properties());

    String strRemove = "rm";

    if (Util.WINDOWS)
    {
        strRemove = "del";
    }

    File inputFile = File.createTempFile("testInputFile", ".txt");
    PrintWriter pwInput = new PrintWriter(new FileWriter(inputFile));
    pwInput.println("1");
    pwInput.close();

    File inputScript = File.createTempFile("testInputScript", "");
    File outputFile = File.createTempFile("testOutputFile", ".txt");
    outputFile.delete();
    PrintWriter pwScript = new PrintWriter(new FileWriter(inputScript));
    pwScript.println("a = load '" + Util.encodeEscape(inputFile.getAbsolutePath()) + "';");
    pwScript.println("store a into '" + Util.encodeEscape(outputFile.getAbsolutePath()) + "';");
    pwScript.println("sh " + strRemove + " " + Util.encodeEscape(inputFile.getAbsolutePath()));
    pwScript.close();

    InputStream inputStream = new FileInputStream(inputScript.getAbsoluteFile());
    server.setBatchOn();
    server.registerScript(inputStream);
    List<ExecJob> execJobs = server.executeBatch();
    assertTrue(execJobs.get(0).getStatus() == JOB_STATUS.COMPLETED);
}
 

public static void registerMultiLineQuery(PigServer pigServer, String query) throws IOException {
    File f = File.createTempFile("tmp", "");
    PrintWriter pw = new PrintWriter(f);
    pw.println(query);
    pw.close();
    pigServer.registerScript(f.getCanonicalPath());
}
 

@Test
public void testBzipStoreInMultiQuery3() throws Exception {
    String[] inputData = new String[] {
            "1\t2\r3\t4"
    };

    String inputFileName = "input3.txt";
    Util.createInputFile(cluster, inputFileName, inputData);

    String inputScript = "set mapred.output.compress true\n" +
            "set mapreduce.output.fileoutputformat.compress true\n" +
            "set mapred.output.compression.codec org.apache.hadoop.io.compress.BZip2Codec\n" +
            "set mapreduce.output.fileoutputformat.compress.codec org.apache.hadoop.io.compress.BZip2Codec\n" +
            "a = load '" + inputFileName + "';\n" +
            "store a into 'output3.bz2';\n" +
            "store a into 'output3';";

    String inputScriptName = "script3.txt";
    PrintWriter pw = new PrintWriter(new FileWriter(inputScriptName));
    pw.println(inputScript);
    pw.close();

    PigServer pig = new PigServer(cluster.getExecType(), properties);

    FileInputStream fis = new FileInputStream(inputScriptName);
    pig.registerScript(fis);

    FileSystem fs = FileSystem.get(ConfigurationUtil.toConfiguration(
            pig.getPigContext().getProperties()));
    FileStatus[] outputFiles = fs.listStatus(new Path("output3"),
            Util.getSuccessMarkerPathFilter());
    assertTrue(outputFiles[0].getLen() > 0);

    outputFiles = fs.listStatus(new Path("output3.bz2"),
            Util.getSuccessMarkerPathFilter());
    assertTrue(outputFiles[0].getLen() > 0);
}
 

@Test
public void testUDFWithoutParameter() throws Exception {
    PigServer pig = new PigServer(ExecType.LOCAL);
    pig.registerScript(TempScriptFile.getAbsolutePath());

    Iterator<Tuple> iterator=pig.openIterator("B");
    int index=0;
    while(iterator.hasNext()){
        Tuple tuple=iterator.next();
        index++;
        int result=(Integer)tuple.get(0);
        assertEquals(result, index);
    }
}
 

@Override
public void execute(Query query) {
    String queryName = query.name;
    String queryValue = query.value;
    Map<String, String> queryMetadata = query.getMetadata();
    String execType = Query.getKey(queryMetadata, METADATA_EXEC_TYPE_KEY).orElse(defaultExecType);
    String alias = Query.getKey(queryMetadata, METADATA_ALIAS_KEY).orElse(defaultOutputAlias);
    log.info("Running {} for alias {}: {}", queryName, alias, queryValue);
    try {
        PigServer server = getPigServer(execType);
        server.registerScript(new ByteArrayInputStream(queryValue.getBytes()));
        Iterator<Tuple> queryResults = server.openIterator(alias);
        Result result = query.createResults();
        // dumpSchema will also, unfortunately, print the schema to stdout.
        List<FieldDetail> metadata = getFieldDetails(server.dumpSchema(alias));
        populateColumns(metadata, result);
        while (queryResults.hasNext()) {
            populateRow(queryResults.next(), metadata, result);
        }
        server.shutdown();
    } catch (IOException ioe) {
        log.error("Problem with Pig query: {}\n{}", queryValue, ioe);
        query.setFailure(ioe.toString());
    } catch (Exception e) {
        log.error("Error occurred while processing Pig query: {}\n{}", queryValue, e);
        query.setFailure(e.toString());
    }
}