core/src/main/java/org/apache/calcite/avatica/remote/ProtobufTranslationImpl.java

/*
 * 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.calcite.avatica.remote;

import org.apache.calcite.avatica.proto.Common.WireMessage;
import org.apache.calcite.avatica.proto.Requests.CatalogsRequest;
import org.apache.calcite.avatica.proto.Requests.CloseConnectionRequest;
import org.apache.calcite.avatica.proto.Requests.CloseStatementRequest;
import org.apache.calcite.avatica.proto.Requests.ColumnsRequest;
import org.apache.calcite.avatica.proto.Requests.CommitRequest;
import org.apache.calcite.avatica.proto.Requests.ConnectionSyncRequest;
import org.apache.calcite.avatica.proto.Requests.CreateStatementRequest;
import org.apache.calcite.avatica.proto.Requests.DatabasePropertyRequest;
import org.apache.calcite.avatica.proto.Requests.ExecuteBatchRequest;
import org.apache.calcite.avatica.proto.Requests.ExecuteRequest;
import org.apache.calcite.avatica.proto.Requests.FetchRequest;
import org.apache.calcite.avatica.proto.Requests.OpenConnectionRequest;
import org.apache.calcite.avatica.proto.Requests.PrepareAndExecuteBatchRequest;
import org.apache.calcite.avatica.proto.Requests.PrepareAndExecuteRequest;
import org.apache.calcite.avatica.proto.Requests.PrepareRequest;
import org.apache.calcite.avatica.proto.Requests.RollbackRequest;
import org.apache.calcite.avatica.proto.Requests.SchemasRequest;
import org.apache.calcite.avatica.proto.Requests.SyncResultsRequest;
import org.apache.calcite.avatica.proto.Requests.TableTypesRequest;
import org.apache.calcite.avatica.proto.Requests.TablesRequest;
import org.apache.calcite.avatica.proto.Requests.TypeInfoRequest;
import org.apache.calcite.avatica.proto.Responses.CloseConnectionResponse;
import org.apache.calcite.avatica.proto.Responses.CloseStatementResponse;
import org.apache.calcite.avatica.proto.Responses.CommitResponse;
import org.apache.calcite.avatica.proto.Responses.ConnectionSyncResponse;
import org.apache.calcite.avatica.proto.Responses.CreateStatementResponse;
import org.apache.calcite.avatica.proto.Responses.DatabasePropertyResponse;
import org.apache.calcite.avatica.proto.Responses.ErrorResponse;
import org.apache.calcite.avatica.proto.Responses.ExecuteBatchResponse;
import org.apache.calcite.avatica.proto.Responses.ExecuteResponse;
import org.apache.calcite.avatica.proto.Responses.FetchResponse;
import org.apache.calcite.avatica.proto.Responses.OpenConnectionResponse;
import org.apache.calcite.avatica.proto.Responses.PrepareResponse;
import org.apache.calcite.avatica.proto.Responses.ResultSetResponse;
import org.apache.calcite.avatica.proto.Responses.RollbackResponse;
import org.apache.calcite.avatica.proto.Responses.RpcMetadata;
import org.apache.calcite.avatica.proto.Responses.SyncResultsResponse;
import org.apache.calcite.avatica.remote.Service.Request;
import org.apache.calcite.avatica.remote.Service.Response;
import org.apache.calcite.avatica.remote.Service.RpcMetadataResponse;
import org.apache.calcite.avatica.util.UnsynchronizedBuffer;

import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import com.google.protobuf.Parser;
import com.google.protobuf.TextFormat;
import com.google.protobuf.UnsafeByteOperations;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static java.nio.charset.StandardCharsets.UTF_8;

/**
 * Implementation of {@link ProtobufTranslationImpl} that translates
 * protobuf requests to POJO requests.
 */
public class ProtobufTranslationImpl implements ProtobufTranslation {
  private static final Logger LOG = LoggerFactory.getLogger(ProtobufTranslationImpl.class);

  /**
   * Encapsulate the logic of transforming a protobuf Request message into the Avatica POJO request.
   */
  static class RequestTranslator {

    private final Parser<? extends Message> parser;
    private final Service.Request impl;

    RequestTranslator(Parser<? extends Message> parser, Service.Request impl) {
      this.parser = parser;
      this.impl = impl;
    }

    public Service.Request transform(ByteString serializedMessage) throws
        InvalidProtocolBufferException {
      // This should already be an aliased CodedInputStream from the WireMessage parsing.
      Message msg = parser.parseFrom(serializedMessage.newCodedInput());
      if (LOG.isTraceEnabled()) {
        LOG.trace(
            "Deserialized {} '{}'",
            msg.getClass().getSimpleName(),
            TextFormat.shortDebugString(msg)
        );
      }
      return impl.deserialize(msg);
    }
  }

  /**
   * Encapsulate the logic of transforming a protobuf Response message into the Avatica POJO
   * Response.
   */
  static class ResponseTranslator {

    private final Parser<? extends Message> parser;
    private final Service.Response impl;

    ResponseTranslator(Parser<? extends Message> parser, Service.Response impl) {
      this.parser = parser;
      this.impl = impl;
    }

    public Service.Response transform(ByteString serializedMessage) throws
        InvalidProtocolBufferException {
      Message msg = parser.parseFrom(serializedMessage);
      if (LOG.isTraceEnabled()) {
        LOG.trace(
            "Deserialized {} '{}'",
            msg.getClass().getSimpleName(),
            TextFormat.shortDebugString(msg)
        );
      }
      return impl.deserialize(msg);
    }
  }

  // Extremely ugly mapping of PB class name into a means to convert it to the POJO
  private static final Map<String, RequestTranslator> REQUEST_PARSERS;
  private static final Map<String, ResponseTranslator> RESPONSE_PARSERS;
  private static final Map<Class<?>, ByteString> MESSAGE_CLASSES;

  static {
    Map<String, RequestTranslator> reqParsers = new ConcurrentHashMap<>();
    reqParsers.put(CatalogsRequest.class.getName(),
        new RequestTranslator(CatalogsRequest.parser(), new Service.CatalogsRequest()));
    reqParsers.put(OpenConnectionRequest.class.getName(),
        new RequestTranslator(OpenConnectionRequest.parser(), new Service.OpenConnectionRequest()));
    reqParsers.put(CloseConnectionRequest.class.getName(),
        new RequestTranslator(CloseConnectionRequest.parser(),
          new Service.CloseConnectionRequest()));
    reqParsers.put(CloseStatementRequest.class.getName(),
        new RequestTranslator(CloseStatementRequest.parser(), new Service.CloseStatementRequest()));
    reqParsers.put(ColumnsRequest.class.getName(),
        new RequestTranslator(ColumnsRequest.parser(), new Service.ColumnsRequest()));
    reqParsers.put(ConnectionSyncRequest.class.getName(),
        new RequestTranslator(ConnectionSyncRequest.parser(), new Service.ConnectionSyncRequest()));
    reqParsers.put(CreateStatementRequest.class.getName(),
        new RequestTranslator(CreateStatementRequest.parser(),
          new Service.CreateStatementRequest()));
    reqParsers.put(DatabasePropertyRequest.class.getName(),
        new RequestTranslator(DatabasePropertyRequest.parser(),
            new Service.DatabasePropertyRequest()));
    reqParsers.put(FetchRequest.class.getName(),
        new RequestTranslator(FetchRequest.parser(), new Service.FetchRequest()));
    reqParsers.put(PrepareAndExecuteRequest.class.getName(),
        new RequestTranslator(PrepareAndExecuteRequest.parser(),
            new Service.PrepareAndExecuteRequest()));
    reqParsers.put(PrepareRequest.class.getName(),
        new RequestTranslator(PrepareRequest.parser(), new Service.PrepareRequest()));
    reqParsers.put(SchemasRequest.class.getName(),
        new RequestTranslator(SchemasRequest.parser(), new Service.SchemasRequest()));
    reqParsers.put(TablesRequest.class.getName(),
        new RequestTranslator(TablesRequest.parser(), new Service.TablesRequest()));
    reqParsers.put(TableTypesRequest.class.getName(),
        new RequestTranslator(TableTypesRequest.parser(), new Service.TableTypesRequest()));
    reqParsers.put(TypeInfoRequest.class.getName(),
        new RequestTranslator(TypeInfoRequest.parser(), new Service.TypeInfoRequest()));
    reqParsers.put(ExecuteRequest.class.getName(),
        new RequestTranslator(ExecuteRequest.parser(), new Service.ExecuteRequest()));
    reqParsers.put(SyncResultsRequest.class.getName(),
        new RequestTranslator(SyncResultsRequest.parser(), new Service.SyncResultsRequest()));
    reqParsers.put(CommitRequest.class.getName(),
        new RequestTranslator(CommitRequest.parser(), new Service.CommitRequest()));
    reqParsers.put(RollbackRequest.class.getName(),
        new RequestTranslator(RollbackRequest.parser(), new Service.RollbackRequest()));
    reqParsers.put(PrepareAndExecuteBatchRequest.class.getName(),
        new RequestTranslator(PrepareAndExecuteBatchRequest.parser(),
            new Service.PrepareAndExecuteBatchRequest()));
    reqParsers.put(ExecuteBatchRequest.class.getName(),
        new RequestTranslator(ExecuteBatchRequest.parser(),
            new Service.ExecuteBatchRequest()));

    REQUEST_PARSERS = Collections.unmodifiableMap(reqParsers);

    Map<String, ResponseTranslator> respParsers = new ConcurrentHashMap<>();
    respParsers.put(OpenConnectionResponse.class.getName(),
        new ResponseTranslator(OpenConnectionResponse.parser(),
            new Service.OpenConnectionResponse()));
    respParsers.put(CloseConnectionResponse.class.getName(),
        new ResponseTranslator(CloseConnectionResponse.parser(),
            new Service.CloseConnectionResponse()));
    respParsers.put(CloseStatementResponse.class.getName(),
        new ResponseTranslator(CloseStatementResponse.parser(),
            new Service.CloseStatementResponse()));
    respParsers.put(ConnectionSyncResponse.class.getName(),
        new ResponseTranslator(ConnectionSyncResponse.parser(),
            new Service.ConnectionSyncResponse()));
    respParsers.put(CreateStatementResponse.class.getName(),
        new ResponseTranslator(CreateStatementResponse.parser(),
            new Service.CreateStatementResponse()));
    respParsers.put(DatabasePropertyResponse.class.getName(),
        new ResponseTranslator(DatabasePropertyResponse.parser(),
            new Service.DatabasePropertyResponse()));
    respParsers.put(ExecuteResponse.class.getName(),
        new ResponseTranslator(ExecuteResponse.parser(), new Service.ExecuteResponse()));
    respParsers.put(FetchResponse.class.getName(),
        new ResponseTranslator(FetchResponse.parser(), new Service.FetchResponse()));
    respParsers.put(PrepareResponse.class.getName(),
        new ResponseTranslator(PrepareResponse.parser(), new Service.PrepareResponse()));
    respParsers.put(ResultSetResponse.class.getName(),
        new ResponseTranslator(ResultSetResponse.parser(), new Service.ResultSetResponse()));
    respParsers.put(ErrorResponse.class.getName(),
        new ResponseTranslator(ErrorResponse.parser(), new Service.ErrorResponse()));
    respParsers.put(SyncResultsResponse.class.getName(),
        new ResponseTranslator(SyncResultsResponse.parser(), new Service.SyncResultsResponse()));
    respParsers.put(RpcMetadata.class.getName(),
        new ResponseTranslator(RpcMetadata.parser(), new RpcMetadataResponse()));
    respParsers.put(CommitResponse.class.getName(),
        new ResponseTranslator(CommitResponse.parser(), new Service.CommitResponse()));
    respParsers.put(RollbackResponse.class.getName(),
        new ResponseTranslator(RollbackResponse.parser(), new Service.RollbackResponse()));
    respParsers.put(ExecuteBatchResponse.class.getName(),
        new ResponseTranslator(ExecuteBatchResponse.parser(), new Service.ExecuteBatchResponse()));

    RESPONSE_PARSERS = Collections.unmodifiableMap(respParsers);

    Map<Class<?>, ByteString> messageClassNames = new ConcurrentHashMap<>();
    for (Class<?> msgClz : getAllMessageClasses()) {
      messageClassNames.put(msgClz, wrapClassName(msgClz));
    }
    MESSAGE_CLASSES = Collections.unmodifiableMap(messageClassNames);
  }

  private static List<Class<?>> getAllMessageClasses() {
    List<Class<?>> messageClasses = new ArrayList<>();
    messageClasses.add(CatalogsRequest.class);
    messageClasses.add(CloseConnectionRequest.class);
    messageClasses.add(CloseStatementRequest.class);
    messageClasses.add(ColumnsRequest.class);
    messageClasses.add(CommitRequest.class);
    messageClasses.add(ConnectionSyncRequest.class);
    messageClasses.add(CreateStatementRequest.class);
    messageClasses.add(DatabasePropertyRequest.class);
    messageClasses.add(ExecuteRequest.class);
    messageClasses.add(FetchRequest.class);
    messageClasses.add(OpenConnectionRequest.class);
    messageClasses.add(PrepareAndExecuteRequest.class);
    messageClasses.add(PrepareRequest.class);
    messageClasses.add(RollbackRequest.class);
    messageClasses.add(SchemasRequest.class);
    messageClasses.add(SyncResultsRequest.class);
    messageClasses.add(TableTypesRequest.class);
    messageClasses.add(TablesRequest.class);
    messageClasses.add(TypeInfoRequest.class);
    messageClasses.add(PrepareAndExecuteBatchRequest.class);
    messageClasses.add(ExecuteBatchRequest.class);

    messageClasses.add(CloseConnectionResponse.class);
    messageClasses.add(CloseStatementResponse.class);
    messageClasses.add(CommitResponse.class);
    messageClasses.add(ConnectionSyncResponse.class);
    messageClasses.add(CreateStatementResponse.class);
    messageClasses.add(DatabasePropertyResponse.class);
    messageClasses.add(ErrorResponse.class);
    messageClasses.add(ExecuteResponse.class);
    messageClasses.add(FetchResponse.class);
    messageClasses.add(OpenConnectionResponse.class);
    messageClasses.add(PrepareResponse.class);
    messageClasses.add(ResultSetResponse.class);
    messageClasses.add(RollbackResponse.class);
    messageClasses.add(RpcMetadata.class);
    messageClasses.add(SyncResultsResponse.class);
    messageClasses.add(ExecuteBatchResponse.class);

    return messageClasses;
  }

  private static ByteString wrapClassName(Class<?> clz) {
    return UnsafeByteOperations.unsafeWrap(clz.getName().getBytes(UTF_8));
  }

  private final ThreadLocal<UnsynchronizedBuffer> threadLocalBuffer =
      new ThreadLocal<UnsynchronizedBuffer>() {
        @Override protected UnsynchronizedBuffer initialValue() {
          return new UnsynchronizedBuffer();
        }
      };

  /**
   * Fetches the concrete message's Parser implementation.
   *
   * @param className The protocol buffer class name
   * @return The Parser for the class
   * @throws IllegalArgumentException If the argument is null or if a Parser for the given
   *     class name is not found.
   */
  public static RequestTranslator getParserForRequest(String className) {
    if (null == className || className.isEmpty()) {
      throw new IllegalArgumentException("Cannot fetch parser for Request with "
          + (null == className ? "null" : "missing") + " class name");
    }

    RequestTranslator translator = REQUEST_PARSERS.get(className);
    if (null == translator) {
      throw new IllegalArgumentException("Cannot find request parser for " + className);
    }

    return translator;
  }

  /**
   * Fetches the concrete message's Parser implementation.
   *
   * @param className The protocol buffer class name
   * @return The Parser for the class
   * @throws IllegalArgumentException If the argument is null or if a Parser for the given
   *     class name is not found.
   */
  public static ResponseTranslator getParserForResponse(String className) {
    if (null == className || className.isEmpty()) {
      throw new IllegalArgumentException("Cannot fetch parser for Response with "
          + (null == className ? "null" : "missing") + " class name");
    }

    ResponseTranslator translator = RESPONSE_PARSERS.get(className);
    if (null == translator) {
      throw new IllegalArgumentException("Cannot find response parser for " + className);
    }

    return translator;
  }

  @Override public byte[] serializeResponse(Response response) throws IOException {
    // Avoid BAOS for its synchronized write methods, we don't need that concurrency control
    UnsynchronizedBuffer out = threadLocalBuffer.get();
    try {
      Message responseMsg = response.serialize();
      // Serialization of the response may be large
      if (LOG.isTraceEnabled()) {
        LOG.trace(
            "Serializing {} '{}'",
            responseMsg.getClass().getSimpleName(),
            TextFormat.shortDebugString(responseMsg)
        );
      }
      serializeMessage(out, responseMsg);
      return out.toArray();
    } finally {
      out.reset();
    }
  }

  @Override public byte[] serializeRequest(Request request) throws IOException {
    // Avoid BAOS for its synchronized write methods, we don't need that concurrency control
    UnsynchronizedBuffer out = threadLocalBuffer.get();
    try {
      Message requestMsg = request.serialize();
      // Serialization of the request may be large
      if (LOG.isTraceEnabled()) {
        LOG.trace(
            "Serializing {} '{}'",
            requestMsg.getClass().getSimpleName(),
            TextFormat.shortDebugString(requestMsg)
        );
      }
      serializeMessage(out, requestMsg);
      return out.toArray();
    } finally {
      out.reset();
    }
  }

  void serializeMessage(OutputStream out, Message msg) throws IOException {
    // Serialize the protobuf message
    UnsynchronizedBuffer buffer = threadLocalBuffer.get();
    ByteString serializedMsg;
    try {
      msg.writeTo(buffer);
      // Make a bytestring from it
      serializedMsg = UnsafeByteOperations.unsafeWrap(buffer.toArray());
    } finally {
      buffer.reset();
    }

    // Wrap the serialized message in a WireMessage
    WireMessage wireMsg = WireMessage.newBuilder().setNameBytes(getClassNameBytes(msg.getClass()))
        .setWrappedMessage(serializedMsg).build();

    // Write the WireMessage to the provided OutputStream
    wireMsg.writeTo(out);
  }

  ByteString getClassNameBytes(Class<?> clz) {
    ByteString byteString = MESSAGE_CLASSES.get(clz);
    if (null == byteString) {
      throw new IllegalArgumentException("Missing ByteString for " + clz.getName());
    }
    return byteString;
  }

  @Override public Request parseRequest(byte[] bytes) throws IOException {
    ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
    CodedInputStream inputStream = byteString.newCodedInput();
    // Enable aliasing to avoid an extra copy to get at the serialized Request inside of the
    // WireMessage.
    inputStream.enableAliasing(true);
    WireMessage wireMsg = WireMessage.parseFrom(inputStream);

    String serializedMessageClassName = wireMsg.getName();

    try {
      RequestTranslator translator = getParserForRequest(serializedMessageClassName);

      // The ByteString should be logical offsets into the original byte array
      return translator.transform(wireMsg.getWrappedMessage());
    } catch (RuntimeException e) {
      if (LOG.isDebugEnabled()) {
        LOG.debug("Failed to parse request message '{}'", TextFormat.shortDebugString(wireMsg));
      }
      throw e;
    }
  }

  @Override public Response parseResponse(byte[] bytes) throws IOException {
    ByteString byteString = UnsafeByteOperations.unsafeWrap(bytes);
    CodedInputStream inputStream = byteString.newCodedInput();
    // Enable aliasing to avoid an extra copy to get at the serialized Response inside of the
    // WireMessage.
    inputStream.enableAliasing(true);
    WireMessage wireMsg = WireMessage.parseFrom(inputStream);

    String serializedMessageClassName = wireMsg.getName();
    try {
      ResponseTranslator translator = getParserForResponse(serializedMessageClassName);

      return translator.transform(wireMsg.getWrappedMessage());
    } catch (RuntimeException e) {
      if (LOG.isDebugEnabled()) {
        LOG.debug("Failed to parse response message '{}'", TextFormat.shortDebugString(wireMsg));
      }
      throw e;
    }
  }
}

// End ProtobufTranslationImpl.java