TTextProtocol.java

// =================================================================================================
// 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
 *
 * @author Alex Roetter
 */
public class TTextProtocol extends TProtocol {
  private static final Logger LOG = Logger.getLogger(
      TTextProtocol.class.getName());
  private static final String OBJECT_AS_KEY_ILLEGAL =
      "A JsonObject (map or struct) can't be a key in a map!";
  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 JsonWriter writer;
  private JsonStreamParser parser;

  /**
   * 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);

    ByteArrayOutputStream mybaos = new TTransportOutputStream();
    OutputStreamWriter osw;
    osw = new OutputStreamWriter(mybaos, Charsets.UTF_8);
    writer = new JsonWriter(osw);
    writer.setIndent("  ");  // two spaces

    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 {
      writer.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()) {
      throw new TException(OBJECT_AS_KEY_ILLEGAL);
    }
    pushContext(context);
    try {
      writer.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 {
      writer.endObject();
      popContext();

      // flush at the end of the final struct.
      if (1 == contextStack.size()) {
        writer.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 {
      writer.beginArray();
    } catch (IOException ex) {
      throw new TTransportException(ex);
    }
  }

  /**
   * Helper shared by write{List/Set}End
   * @throws TException
   */
  private void writeSequenceEnd() throws TException {
    try {
      writer.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()) {
        writer.name(val.toString());
      } else {
        helper.writeValue(writer, 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();
    if (getCurrentContext().isMapKey()) {
      throw new TException(OBJECT_AS_KEY_ILLEGAL);
    }

    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 (!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();
    if (getCurrentContext().isMapKey()) {
      throw new TException(OBJECT_AS_KEY_ILLEGAL);
    }

    JsonElement curElem = getCurrentContext().getCurrentChild();
    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
   * @return
   * @throws TException
   */
  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 JsonPrimitve 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();
  }

  /** 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 tranport
    @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();
    }
  }
}