com.twitter.common.thrift.text.TTextProtocol.java Source code

Java tutorial

Introduction

Here is the source code for com.twitter.common.thrift.text.TTextProtocol.java

Source

// =================================================================================================
// Copyright 2011 Twitter, Inc.
// -------------------------------------------------------------------------------------------------
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this work except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file, or 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 com.twitter.common.thrift.text;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStreamWriter;
import java.nio.ByteBuffer;
import java.util.Stack;
import java.util.logging.Logger;

import com.google.common.base.Charsets;
import com.google.common.io.ByteStreams;
import com.google.gson.JsonElement;
import com.google.gson.JsonStreamParser;
import com.google.gson.stream.JsonWriter;

import org.apache.commons.codec.binary.Base64;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TField;
import org.apache.thrift.protocol.TList;
import org.apache.thrift.protocol.TMap;
import org.apache.thrift.protocol.TMessage;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.protocol.TSet;
import org.apache.thrift.protocol.TStruct;
import org.apache.thrift.protocol.TType;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;

/**
 * A simple text format for serializing/deserializing thrift
 * messages. This format is inefficient in space.
 *
 * For an example, see:
 * tests/resources/com/twitter/common/thrift/text/TTextProtocol_TestData.txt
 *
 * which is a text encoding of the thrift message defined in:
 *
 * src/main/thrift/com/twitter/common/thrift/text/TTextProtocolTest.thrift
 *
 * Whitespace (including newlines) is not significant.
 *
 * No comments are allowed in the json.
 *
 * We support parsing structs and anything embedded in a struct,
 * but not messages (which are generated as part of thrift RPC
 * service definitions).
 *
 * TODO(Alex Roetter): write a wrapper that allows us to read in a file
 * of many structs (perhaps stored in a JsonArray), passing each struct to
 * this class for parsing.
 *
 * See thrift's @see org.apache.thrift.protocol.TJSONProtocol
 * for another example an implementation of the @see TProtocol
 * interface. This class is based on that.
 *
 * TODO(Alex Roetter): Also add a new TEXT_PROTOCOL field to ThriftCodec
 *
 * TODO(Alex Roetter): throw this up on my github. Seems generally useful.
 *
 * TODO(Alex Roetter): add support for enums
 *
 */
public class TTextProtocol extends TProtocol {
    private static final Logger LOG = Logger.getLogger(TTextProtocol.class.getName());
    private static final String SEQUENCE_AS_KEY_ILLEGAL = "Can't have a sequence (list or set) as a key in a map!";

    private static final TStruct ANONYMOUS_STRUCT = new TStruct();

    // how many bytes to read at once
    private static final int READ_BUFFER_SIZE = 1024;

    private static final byte UNUSED_TYPE = TType.STOP;

    private Stack<BaseContext> contextStack;
    private Base64 base64Encoder = new Base64();
    private final Stack<WriterByteArrayOutputStream> writers;
    private JsonStreamParser parser;

    private static class WriterByteArrayOutputStream {
        final JsonWriter writer;
        final ByteArrayOutputStream baos;

        public WriterByteArrayOutputStream(JsonWriter writer, ByteArrayOutputStream baos) {
            this.writer = writer;
            this.baos = baos;
        }
    }

    /**
     * Factory
     */
    public static class Factory implements TProtocolFactory {
        @Override
        public TProtocol getProtocol(TTransport trans) {
            return new TTextProtocol(trans);
        }
    }

    /**
     * Create a parser which can read from trans, and create the output writer
     * that can write to a TTransport
     */
    public TTextProtocol(TTransport trans) {
        super(trans);

        writers = new Stack<WriterByteArrayOutputStream>();

        ByteArrayOutputStream mybaos = new TTransportOutputStream();
        pushWriter(mybaos);

        reset();

        try {
            parser = createParser();
        } catch (IOException ex) {
            // This happens when we're created in write mode (i.e.
            // there is nothing to parse). Calls to read methods will fail.
        }
        contextStack = new Stack<BaseContext>();
        contextStack.push(new BaseContext());
    }

    @Override
    public final void reset() {
    }

    /**
     * I believe these two messages are called for a thrift service
     * interface. We don't plan on storing any text objects of that
     * type on disk.
     */
    @Override
    public void writeMessageBegin(TMessage message) throws TException {
        unsupportedOperation();
    }

    @Override
    public void writeMessageEnd() throws TException {
        unsupportedOperation();
    }

    /**
     * Throws an exception for invoked operations that we don't support.
     */
    private void unsupportedOperation() {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    @Override
    public void writeStructBegin(TStruct struct) throws TException {
        writeJsonObjectBegin(new StructContext(null));
    }

    @Override
    public void writeStructEnd() throws TException {
        writeJsonObjectEnd();
    }

    @Override
    public void writeFieldBegin(TField field) throws TException {
        try {
            getCurrentWriter().name(field.name);
        } catch (IOException ex) {
            throw new TException(ex);
        }
    }

    @Override
    public void writeFieldEnd() throws TException {
    }

    @Override
    public void writeFieldStop() throws TException {
    }

    @Override
    public void writeMapBegin(TMap map) throws TException {
        writeJsonObjectBegin(new MapContext(null));
    }

    @Override
    public void writeMapEnd() throws TException {
        writeJsonObjectEnd();
    }

    /**
     * Helper to write out the beginning of a Thrift type (either struct or map),
     * both of which are written as JsonObjects.
     */
    private void writeJsonObjectBegin(BaseContext context) throws TException {
        getCurrentContext().write();
        if (getCurrentContext().isMapKey()) {
            pushWriter(new ByteArrayOutputStream());
        }
        pushContext(context);
        try {
            getCurrentWriter().beginObject();
        } catch (IOException ex) {
            throw new TException(ex);
        }
    }

    /**
     * Helper to write out the end of a Thrift type (either struct or map),
     * both of which are written as JsonObjects.
     * @throws TException
     */
    private void writeJsonObjectEnd() throws TException {
        try {
            getCurrentWriter().endObject();
            popContext();
            if (getCurrentContext().isMapKey()) {
                String writerString = getWriterString();
                popWriter();
                getCurrentWriter().name(writerString);
            }

            // flush at the end of the final struct.
            if (1 == contextStack.size()) {
                getCurrentWriter().flush();
            }
        } catch (IOException ex) {
            throw new TException(ex);
        }
    }

    @Override
    public void writeListBegin(TList list) throws TException {
        writeSequenceBegin(list.size);
    }

    @Override
    public void writeListEnd() throws TException {
        writeSequenceEnd();
    }

    @Override
    public void writeSetBegin(TSet set) throws TException {
        writeSequenceBegin(set.size);
    }

    @Override
    public void writeSetEnd() throws TException {
        writeListEnd();
    }

    /**
     * Helper shared by write{List/Set}Begin
     */
    private void writeSequenceBegin(int size) throws TException {
        getCurrentContext().write();
        if (getCurrentContext().isMapKey()) {
            throw new TException(SEQUENCE_AS_KEY_ILLEGAL);
        }
        pushContext(new SequenceContext(null));

        try {
            getCurrentWriter().beginArray();
        } catch (IOException ex) {
            throw new TTransportException(ex);
        }
    }

    /**
     * Helper shared by write{List/Set}End
     * @throws TException
     */
    private void writeSequenceEnd() throws TException {
        try {
            getCurrentWriter().endArray();
        } catch (IOException ex) {
            throw new TTransportException(ex);
        }
        popContext();
    }

    @Override
    public void writeBool(boolean b) throws TException {
        writeNameOrValue(TypedParser.BOOLEAN, b);
    }

    @Override
    public void writeByte(byte b) throws TException {
        writeNameOrValue(TypedParser.BYTE, b);
    }

    @Override
    public void writeI16(short i16) throws TException {
        writeNameOrValue(TypedParser.SHORT, i16);
    }

    @Override
    public void writeI32(int i32) throws TException {
        writeNameOrValue(TypedParser.INTEGER, i32);
    }

    @Override
    public void writeI64(long i64) throws TException {
        writeNameOrValue(TypedParser.LONG, i64);
    }

    @Override
    public void writeDouble(double dub) throws TException {
        writeNameOrValue(TypedParser.DOUBLE, dub);
    }

    @Override
    public void writeString(String str) throws TException {
        writeNameOrValue(TypedParser.STRING, str);
    }

    @Override
    public void writeBinary(ByteBuffer buf) throws TException {
        writeString(new String(base64Encoder.encode(buf.array())));
    }

    /**
     * Write out the given value, either as a JSON name (meaning it's
     * escaped by quotes), or a value. The TypedParser knows how to
     * handle the writing.
     */
    private <T> void writeNameOrValue(TypedParser<T> helper, T val) throws TException {
        getCurrentContext().write();
        try {
            if (getCurrentContext().isMapKey()) {
                getCurrentWriter().name(val.toString());
            } else {
                helper.writeValue(getCurrentWriter(), val);
            }
        } catch (IOException ex) {
            throw new TException(ex);
        }
    }

    /////////////////////////////////////////
    // Read methods
    /////////////////////////////////////////
    @Override
    public TMessage readMessageBegin() throws TException {
        unsupportedOperation();
        return null;
    }

    @Override
    public void readMessageEnd() throws TException {
        unsupportedOperation();
    }

    @Override
    public TStruct readStructBegin() throws TException {
        getCurrentContext().read();

        JsonElement structElem;
        // Reading a new top level struct if the only item on the stack
        // is the BaseContext
        if (1 == contextStack.size()) {
            structElem = parser.next();
            if (null == structElem) {
                throw new TException("parser.next() has nothing to parse!");
            }
        } else {
            structElem = getCurrentContext().getCurrentChild();
        }

        if (getCurrentContext().isMapKey()) {
            structElem = new JsonStreamParser(structElem.getAsString()).next();
        }

        if (!structElem.isJsonObject()) {
            throw new TException("Expected Json Object!");
        }

        pushContext(new StructContext(structElem.getAsJsonObject()));
        return ANONYMOUS_STRUCT;
    }

    @Override
    public void readStructEnd() throws TException {
        popContext();
    }

    @Override
    public TField readFieldBegin() throws TException {
        String name = null;

        if (!getCurrentContext().hasMoreChildren()) {
            return new TField("", UNUSED_TYPE, (short) 0);
        }

        getCurrentContext().read();

        JsonElement jsonName = getCurrentContext().getCurrentChild();

        if (!jsonName.getAsJsonPrimitive().isString()) {
            throw new RuntimeException("Expected String for a field name");
        }

        return getCurrentContext().getTFieldByName(jsonName.getAsJsonPrimitive().getAsString());
    }

    @Override
    public void readFieldEnd() throws TException {
    }

    @Override
    public TMap readMapBegin() throws TException {
        getCurrentContext().read();

        JsonElement curElem = getCurrentContext().getCurrentChild();

        if (getCurrentContext().isMapKey()) {
            curElem = new JsonStreamParser(curElem.getAsString()).next();
        }

        if (!curElem.isJsonObject()) {
            throw new TException("Expected JSON Object!");
        }

        pushContext(new MapContext(curElem.getAsJsonObject()));

        return new TMap(UNUSED_TYPE, UNUSED_TYPE, curElem.getAsJsonObject().entrySet().size());
    }

    @Override
    public void readMapEnd() throws TException {
        popContext();
    }

    @Override
    public TList readListBegin() throws TException {
        int size = readSequenceBegin();
        return new TList(UNUSED_TYPE, size);
    }

    @Override
    public void readListEnd() throws TException {
        readSequenceEnd();
    }

    @Override
    public TSet readSetBegin() throws TException {
        int size = readSequenceBegin();
        return new TSet(UNUSED_TYPE, size);
    }

    @Override
    public void readSetEnd() throws TException {
        readSequenceEnd();
    }

    /**
     * Helper shared by read{List/Set}Begin
     */
    private int readSequenceBegin() throws TException {
        getCurrentContext().read();
        if (getCurrentContext().isMapKey()) {
            throw new TException(SEQUENCE_AS_KEY_ILLEGAL);
        }

        JsonElement curElem = getCurrentContext().getCurrentChild();
        if (!curElem.isJsonArray()) {
            throw new TException("Expected JSON Array!");
        }

        pushContext(new SequenceContext(curElem.getAsJsonArray()));
        return curElem.getAsJsonArray().size();
    }

    /**
     * Helper shared by read{List/Set}End
     */
    private void readSequenceEnd() {
        popContext();
    }

    @Override
    public boolean readBool() throws TException {
        return readNameOrValue(TypedParser.BOOLEAN);
    }

    @Override
    public byte readByte() throws TException {
        return readNameOrValue(TypedParser.BYTE);
    }

    @Override
    public short readI16() throws TException {
        return readNameOrValue(TypedParser.SHORT);
    }

    @Override
    public int readI32() throws TException {
        return readNameOrValue(TypedParser.INTEGER);
    }

    @Override
    public long readI64() throws TException {
        return readNameOrValue(TypedParser.LONG);
    }

    @Override
    public double readDouble() throws TException {
        return readNameOrValue(TypedParser.DOUBLE);
    }

    @Override
    public String readString() throws TException {
        return readNameOrValue(TypedParser.STRING);
    }

    @Override
    public ByteBuffer readBinary() throws TException {
        return ByteBuffer.wrap(base64Encoder.decode(readString()));
    }

    /**
     * Read in a value of the given type, either as a name (meaning the
     * JSONElement is a string and we convert it), or as a value
     * (meaning the JSONElement has the type we expect).
     * Uses a TypedParser to do the real work.
     *
     * TODO(Alex Roetter): not sure TypedParser is a win for the number of
     * lines it saves. Consider expanding out all the readX() methods to
     * do what readNameOrValue does, calling the relevant methods from
     * the TypedParser directly.
     */
    private <T> T readNameOrValue(TypedParser<T> ch) {
        getCurrentContext().read();

        JsonElement elem = getCurrentContext().getCurrentChild();
        if (getCurrentContext().isMapKey()) {
            // Will throw a ClassCastException if this is not a JsonPrimitive string
            return ch.readFromString(elem.getAsString());
        } else {
            return ch.readFromJsonElement(elem);
        }
    }

    /**
     * Set up the stream parser to read from the trans_ TTransport
     * buffer.
     */
    private JsonStreamParser createParser() throws IOException {
        return new JsonStreamParser(new String(ByteStreams.toByteArray(new InputStream() {
            private int index;
            private int max;
            private final byte[] buffer = new byte[READ_BUFFER_SIZE];

            @Override
            public int read() throws IOException {
                if (max == -1) {
                    return -1;
                }
                if (max > 0 && index < max) {
                    return buffer[index++];
                }
                try {
                    max = trans_.read(buffer, 0, READ_BUFFER_SIZE);
                    index = 0;
                } catch (TTransportException e) {
                    if (TTransportException.END_OF_FILE != e.getType()) {
                        throw new IOException(e);
                    }
                    max = -1;
                }
                return read();
            }
        }), Charsets.UTF_8));
    }

    /**
     * Return the current parsing context
     */
    private BaseContext getCurrentContext() {
        return contextStack.peek();
    }

    /**
     * Add a new parsing context onto the parse context stack
     */
    private void pushContext(BaseContext c) {
        contextStack.push(c);
    }

    /**
     * Pop a parsing context from the parse context stack
     */
    private void popContext() {
        contextStack.pop();
    }

    /**
     * Return the current parsing context
     */
    private JsonWriter getCurrentWriter() {
        return writers.peek().writer;
    }

    private String getWriterString() throws TException {
        WriterByteArrayOutputStream wbaos = writers.peek();
        String ret;
        try {
            wbaos.writer.flush();
            ret = new String(wbaos.baos.toByteArray());
            wbaos.writer.close();
        } catch (IOException e) {
            throw new TException(e);
        }
        return ret;
    }

    private void pushWriter(ByteArrayOutputStream baos) {
        OutputStreamWriter osw = new OutputStreamWriter(baos, Charsets.UTF_8);

        JsonWriter writer = new JsonWriter(osw);
        writer.setIndent("  "); // two spaces

        WriterByteArrayOutputStream wbaos = new WriterByteArrayOutputStream(writer, baos);
        writers.push(wbaos);
    }

    private void popWriter() {
        writers.pop();
    }

    /** Just a byte array output stream that forwards all data to
     * a TTransport when it is flushed or closed
     */
    private class TTransportOutputStream extends ByteArrayOutputStream {
        // This isn't necessary, but a good idea to close the transport
        @Override
        public void close() throws IOException {
            flush();

            super.close();
            trans_.close();
        }

        @Override
        public void flush() throws IOException {
            try {
                super.flush();
                byte[] bytes = toByteArray();
                trans_.write(bytes);
                trans_.flush();

            } catch (TTransportException ex) {
                throw new IOException(ex);
            }
            // Clears the internal memory buffer, since we've already
            // written it out.
            super.reset();
        }
    }
}