org.apache.lens.lib.query.TestAbstractFileFormatter.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.lens.lib.query.TestAbstractFileFormatter.java

Source

/**
 * 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 org.apache.lens.lib.query;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;

import java.util.ArrayList;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

import org.apache.lens.api.LensConf;
import org.apache.lens.server.api.LensConfConstants;
import org.apache.lens.server.api.driver.LensDriver;
import org.apache.lens.server.api.driver.LensResultSetMetadata;
import org.apache.lens.server.api.driver.MockDriver;
import org.apache.lens.server.api.error.LensException;
import org.apache.lens.server.api.query.QueryContext;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.CompressionCodecFactory;

import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.Test;

/**
 * The Class TestAbstractFileFormatter.
 */
public abstract class TestAbstractFileFormatter {

    /**
     * The formatter.
     */
    protected WrappedFileFormatter formatter;

    /**
     * Cleanup.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @AfterMethod
    public void cleanup() throws IOException {
        if (formatter != null) {
            FileSystem fs = new Path(formatter.getFinalOutputPath()).getFileSystem(new Configuration());
            fs.delete(new Path(formatter.getFinalOutputPath()), true);
        }
    }

    /**
     * Test formatter.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testFormatter() throws IOException {
        Configuration conf = new Configuration();
        setConf(conf);
        testFormatter(conf, "UTF8", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv", getMockedResultSet());
        // validate rows
        Assert.assertEquals(readFinalOutputFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-8"),
                getExpectedCSVRows());
    }

    /**
     * Test compression.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testCompression() throws IOException {
        Configuration conf = new Configuration();
        conf.setBoolean(LensConfConstants.QUERY_OUTPUT_ENABLE_COMPRESSION, true);
        setConf(conf);
        testFormatter(conf, "UTF8", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv.gz",
                getMockedResultSet());
        // validate rows
        Assert.assertEquals(readCompressedFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-8"),
                getExpectedCSVRows());
    }

    /**
     * Test custom compression.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testCustomCompression() throws IOException {
        Configuration conf = new Configuration();
        conf.setBoolean(LensConfConstants.QUERY_OUTPUT_ENABLE_COMPRESSION, true);
        conf.set(LensConfConstants.QUERY_OUTPUT_COMPRESSION_CODEC,
                org.apache.hadoop.io.compress.DefaultCodec.class.getCanonicalName());
        setConf(conf);
        testFormatter(conf, "UTF8", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv.deflate",
                getMockedResultSet());
        // validate rows
        Assert.assertEquals(readCompressedFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-8"),
                getExpectedCSVRows());
    }

    /**
     * Test encoding.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testEncoding() throws IOException {
        Configuration conf = new Configuration();
        conf.set(LensConfConstants.QUERY_OUTPUT_CHARSET_ENCODING, "UTF-16LE");
        setConf(conf);
        testFormatter(conf, "UnicodeLittleUnmarked", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv",
                getMockedResultSet());
        // validate rows
        Assert.assertEquals(readFinalOutputFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-16LE"),
                getExpectedCSVRows());
    }

    /**
     * Test compression and encoding.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testCompressionAndEncoding() throws IOException {
        Configuration conf = new Configuration();
        conf.set(LensConfConstants.QUERY_OUTPUT_CHARSET_ENCODING, "UTF-16LE");
        conf.setBoolean(LensConfConstants.QUERY_OUTPUT_ENABLE_COMPRESSION, true);
        setConf(conf);
        testFormatter(conf, "UnicodeLittleUnmarked", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv.gz",
                getMockedResultSet());
        // validate rows
        Assert.assertEquals(readCompressedFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-16LE"),
                getExpectedCSVRows());
    }

    /**
     * Test output path.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testOutputPath() throws IOException {
        Configuration conf = new Configuration();
        String outputParent = "target/" + getClass().getSimpleName();
        conf.set(LensConfConstants.RESULT_SET_PARENT_DIR, outputParent);
        setConf(conf);
        testFormatter(conf, "UTF8", outputParent, ".csv", getMockedResultSet());
        // validate rows
        Assert.assertEquals(readFinalOutputFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-8"),
                getExpectedCSVRows());
    }

    /**
     * Test compression with custom output path.
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testCompressionWithCustomOutputPath() throws IOException {
        Configuration conf = new Configuration();
        String outputParent = "target/" + getClass().getSimpleName();
        conf.set(LensConfConstants.RESULT_SET_PARENT_DIR, outputParent);
        conf.setBoolean(LensConfConstants.QUERY_OUTPUT_ENABLE_COMPRESSION, true);
        setConf(conf);
        testFormatter(conf, "UTF8", outputParent, ".csv.gz", getMockedResultSet());
        // validate rows
        Assert.assertEquals(readCompressedFile(new Path(formatter.getFinalOutputPath()), conf, "UTF-8"),
                getExpectedCSVRows());
    }

    /**
     * Test formatter persistence
     *
     * @throws IOException Signals that an I/O exception has occurred.
     */
    @Test
    public void testFormatterPersistence() throws IOException, ClassNotFoundException {
        Configuration conf = new Configuration();
        setConf(conf);
        testFormatter(conf, "UTF8", LensConfConstants.RESULT_SET_PARENT_DIR_DEFAULT, ".csv", getMockedResultSet());

        // Write formatter to stream
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try {
            new ObjectOutputStream(outputStream).writeObject(formatter);
        } finally {
            outputStream.close();
        }
        // Create another formatter from the stream
        ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray());
        WrappedFileFormatter newFormatter = createFormatter();
        try {
            newFormatter = (WrappedFileFormatter) new ObjectInputStream(inputStream).readObject();
        } finally {
            inputStream.close();
        }

        Assert.assertEquals(formatter.getFinalOutputPath(), newFormatter.getFinalOutputPath());
        Assert.assertEquals(formatter.getFileSize(), newFormatter.getFileSize());
        Assert.assertEquals(formatter.getNumRows(), newFormatter.getNumRows());
        Assert.assertEquals(formatter.getMetadata().toJson(), newFormatter.getMetadata().toJson());
    }

    /**
     * Creates the formatter.
     *
     * @return the wrapped file formatter
     */
    protected abstract WrappedFileFormatter createFormatter();

    /**
     * Write all rows.
     *
     * @param conf the conf
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected abstract void writeAllRows(Configuration conf) throws IOException;

    protected void setConf(Configuration conf) {
    }

    /**
     * Creates the query context
     * @param conf      the conf
     * @param queryName the name of query
     * @return the query context
     */
    protected QueryContext createContext(Configuration conf, String queryName) {
        final LensDriver mockDriver = new MockDriver();
        try {
            mockDriver.configure(conf, null, null);
        } catch (LensException e) {
            Assert.fail(e.getMessage());
        }
        QueryContext ctx = QueryContext.createContextWithSingleDriver("test writer query", "testuser",
                new LensConf(), conf, mockDriver, null, false);

        ctx.setSelectedDriver(mockDriver);
        ctx.setQueryName(queryName);
        return ctx;
    }

    /**
     * Validates the formatter
     * @param conf              the conf
     * @param charsetEncoding   the charset encoding
     * @param outputParentDir   the output parent dir
     * @param fileExtn          the file extn
     * @param columnNames       the column names
     * @param ctx               the query context
     * @param expectedFinalPath the final path of output
     * @throws IOException Signals that an I/O exception has occurred.
     */
    public void validateFormatter(Configuration conf, String charsetEncoding, String outputParentDir,
            String fileExtn, LensResultSetMetadata columnNames, QueryContext ctx, Path expectedFinalPath)
            throws IOException {
        formatter = createFormatter();
        formatter.init(ctx, columnNames);

        // check output spec
        Assert.assertEquals(formatter.getEncoding(), charsetEncoding);
        Path tmpPath = formatter.getTmpPath();
        Path expectedTmpPath = new Path(outputParentDir, ctx.getQueryHandle() + ".tmp" + fileExtn);
        Assert.assertEquals(tmpPath, expectedTmpPath);

        // write header, rows and footer;
        formatter.writeHeader();
        writeAllRows(conf);
        formatter.writeFooter();
        FileSystem fs = expectedTmpPath.getFileSystem(conf);
        Assert.assertTrue(fs.exists(tmpPath));

        // commit and close
        formatter.commit();
        formatter.close();
        Assert.assertFalse(fs.exists(tmpPath));
        Path finalPath = new Path(formatter.getFinalOutputPath());
        Assert.assertEquals(finalPath, expectedFinalPath);
        Assert.assertTrue(fs.exists(finalPath));
    }

    /**
     * Test formatter.
     *
     * @param conf            the conf
     * @param charsetEncoding the charset encoding
     * @param outputParentDir the output parent dir
     * @param fileExtn        the file extn
     * @param columnNames     the column names
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected void testFormatter(Configuration conf, String charsetEncoding, String outputParentDir,
            String fileExtn, LensResultSetMetadata columnNames) throws IOException {

        QueryContext ctx = createContext(conf, null);

        Path expectedFinalPath = new Path(outputParentDir, ctx.getQueryHandle() + fileExtn);
        FileSystem fs = expectedFinalPath.getFileSystem(conf);
        expectedFinalPath = expectedFinalPath.makeQualified(fs);
        validateFormatter(conf, charsetEncoding, outputParentDir, fileExtn, columnNames, ctx, expectedFinalPath);
    }

    /**
     * Test Formatter with a different final path
     * @param conf              the conf
     * @param charsetEncoding   the charset encoding
     * @param outputParentDir   the output parent dir
     * @param fileExtn          the file extn
     * @param columnNames       the column names
     * @param queryName         the name of the query
     * @param expectedFinalPath Final path of the output
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected void testFormatterWithFinalPath(Configuration conf, String charsetEncoding, String outputParentDir,
            String fileExtn, LensResultSetMetadata columnNames, String queryName, Path expectedFinalPath)
            throws IOException {
        QueryContext ctx = createContext(conf, queryName);

        validateFormatter(conf, charsetEncoding, outputParentDir, fileExtn, columnNames, ctx, expectedFinalPath);
    }

    /**
     * Read final output file.
     *
     * @param finalPath the final path
     * @param conf      the conf
     * @param encoding  the encoding
     * @return the list
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected List<String> readFinalOutputFile(Path finalPath, Configuration conf, String encoding)
            throws IOException {
        FileSystem fs = finalPath.getFileSystem(conf);
        return readFromStream(new InputStreamReader(fs.open(finalPath), encoding));
    }

    /**
     * Read compressed file.
     *
     * @param finalPath the final path
     * @param conf      the conf
     * @param encoding  the encoding
     * @return the list
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected List<String> readCompressedFile(Path finalPath, Configuration conf, String encoding)
            throws IOException {
        CompressionCodecFactory compressionCodecs = new CompressionCodecFactory(conf);
        compressionCodecs = new CompressionCodecFactory(conf);
        final CompressionCodec codec = compressionCodecs.getCodec(finalPath);
        FileSystem fs = finalPath.getFileSystem(conf);
        return readFromStream(new InputStreamReader(codec.createInputStream(fs.open(finalPath)), encoding));
    }

    /**
     * Read from stream.
     *
     * @param ir the ir
     * @return the list
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected List<String> readFromStream(InputStreamReader ir) throws IOException {
        List<String> result = new ArrayList<String>();
        BufferedReader reader = new BufferedReader(ir);
        String line = reader.readLine();
        while (line != null) {
            result.add(line);
            line = reader.readLine();
        }
        reader.close();
        return result;

    }

    protected LensResultSetMetadata getMockedResultSet() {
        return MockLensResultSetMetadata.createMockedResultSet();
    }

    protected LensResultSetMetadata getMockedResultSetWithoutComma() {
        return MockLensResultSetMetadata.createMockedResultSetWithoutComma();
    }

    /**
     * Read zip output file.
     *
     * @param finalPath the final path
     * @param conf      the conf
     * @param encoding  the encoding
     * @return the list
     * @throws IOException Signals that an I/O exception has occurred.
     */
    protected List<String> readZipOutputFile(Path finalPath, Configuration conf, String encoding)
            throws IOException {
        FileSystem fs = finalPath.getFileSystem(conf);
        List<String> result = new ArrayList<String>();
        ZipEntry ze = null;
        ZipInputStream zin = new ZipInputStream(fs.open(finalPath));
        while ((ze = zin.getNextEntry()) != null) {
            BufferedReader reader = new BufferedReader(new InputStreamReader(zin, encoding));
            String line = reader.readLine();
            while (line != null) {
                result.add(line);
                line = reader.readLine();
            }
            zin.closeEntry();
        }
        zin.close();
        return result;
    }

    protected abstract List<String> getExpectedCSVRows();

    protected abstract List<String> getExpectedTextRows();

    protected abstract List<String> getExpectedCSVRowsWithoutComma();

    protected abstract List<String> getExpectedTextRowsWithoutComma();

    protected abstract List<String> getExpectedCSVRowsWithMultiple();

    protected abstract List<String> getExpectedTextRowsWithMultiple();

    protected abstract List<String> getExpectedCSVRowsWithMultipleWithoutComma();

    protected abstract List<String> getExpectedTextRowsWithMultipleWithoutComma();
}