FlightBindingService.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.arrow.flight;

import com.google.common.collect.ImmutableSet;
import io.grpc.BindableService;
import io.grpc.MethodDescriptor;
import io.grpc.MethodDescriptor.MethodType;
import io.grpc.ServerMethodDefinition;
import io.grpc.ServerServiceDefinition;
import io.grpc.ServiceDescriptor;
import io.grpc.protobuf.ProtoUtils;
import io.grpc.stub.ServerCalls;
import io.grpc.stub.StreamObserver;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import org.apache.arrow.flight.auth.ServerAuthHandler;
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.Flight.PutResult;
import org.apache.arrow.flight.impl.FlightServiceGrpc;
import org.apache.arrow.memory.BufferAllocator;

/** Extends the basic flight service to override some methods for more efficient implementations. */
class FlightBindingService implements BindableService {

  private static final String DO_GET =
      MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoGet");
  private static final String DO_PUT =
      MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoPut");
  private static final String DO_EXCHANGE =
      MethodDescriptor.generateFullMethodName(FlightConstants.SERVICE, "DoExchange");
  private static final Set<String> OVERRIDE_METHODS = ImmutableSet.of(DO_GET, DO_PUT, DO_EXCHANGE);

  private final FlightService delegate;
  private final BufferAllocator allocator;

  public FlightBindingService(
      BufferAllocator allocator,
      FlightProducer producer,
      ServerAuthHandler authHandler,
      ExecutorService executor) {
    this.allocator = allocator;
    this.delegate = new FlightService(allocator, producer, authHandler, executor);
  }

  public static MethodDescriptor<Flight.Ticket, ArrowMessage> getDoGetDescriptor(
      BufferAllocator allocator) {
    return MethodDescriptor.<Flight.Ticket, ArrowMessage>newBuilder()
        .setType(io.grpc.MethodDescriptor.MethodType.SERVER_STREAMING)
        .setFullMethodName(DO_GET)
        .setSampledToLocalTracing(false)
        .setRequestMarshaller(ProtoUtils.marshaller(Flight.Ticket.getDefaultInstance()))
        .setResponseMarshaller(ArrowMessage.createMarshaller(allocator))
        .setSchemaDescriptor(FlightServiceGrpc.getDoGetMethod().getSchemaDescriptor())
        .build();
  }

  public static MethodDescriptor<ArrowMessage, Flight.PutResult> getDoPutDescriptor(
      BufferAllocator allocator) {
    return MethodDescriptor.<ArrowMessage, Flight.PutResult>newBuilder()
        .setType(MethodType.BIDI_STREAMING)
        .setFullMethodName(DO_PUT)
        .setSampledToLocalTracing(false)
        .setRequestMarshaller(ArrowMessage.createMarshaller(allocator))
        .setResponseMarshaller(ProtoUtils.marshaller(Flight.PutResult.getDefaultInstance()))
        .setSchemaDescriptor(FlightServiceGrpc.getDoPutMethod().getSchemaDescriptor())
        .build();
  }

  public static MethodDescriptor<ArrowMessage, ArrowMessage> getDoExchangeDescriptor(
      BufferAllocator allocator) {
    return MethodDescriptor.<ArrowMessage, ArrowMessage>newBuilder()
        .setType(MethodType.BIDI_STREAMING)
        .setFullMethodName(DO_EXCHANGE)
        .setSampledToLocalTracing(false)
        .setRequestMarshaller(ArrowMessage.createMarshaller(allocator))
        .setResponseMarshaller(ArrowMessage.createMarshaller(allocator))
        .setSchemaDescriptor(FlightServiceGrpc.getDoExchangeMethod().getSchemaDescriptor())
        .build();
  }

  @Override
  public ServerServiceDefinition bindService() {
    final ServerServiceDefinition baseDefinition = delegate.bindService();

    final MethodDescriptor<Flight.Ticket, ArrowMessage> doGetDescriptor =
        getDoGetDescriptor(allocator);
    final MethodDescriptor<ArrowMessage, Flight.PutResult> doPutDescriptor =
        getDoPutDescriptor(allocator);
    final MethodDescriptor<ArrowMessage, ArrowMessage> doExchangeDescriptor =
        getDoExchangeDescriptor(allocator);

    // Make sure we preserve SchemaDescriptor fields on methods so that gRPC reflection still works.
    final ServiceDescriptor.Builder serviceDescriptorBuilder =
        ServiceDescriptor.newBuilder(FlightConstants.SERVICE)
            .setSchemaDescriptor(baseDefinition.getServiceDescriptor().getSchemaDescriptor());
    serviceDescriptorBuilder.addMethod(doGetDescriptor);
    serviceDescriptorBuilder.addMethod(doPutDescriptor);
    serviceDescriptorBuilder.addMethod(doExchangeDescriptor);
    for (MethodDescriptor<?, ?> definition : baseDefinition.getServiceDescriptor().getMethods()) {
      if (OVERRIDE_METHODS.contains(definition.getFullMethodName())) {
        continue;
      }

      serviceDescriptorBuilder.addMethod(definition);
    }

    final ServiceDescriptor serviceDescriptor = serviceDescriptorBuilder.build();
    ServerServiceDefinition.Builder serviceBuilder =
        ServerServiceDefinition.builder(serviceDescriptor);
    serviceBuilder.addMethod(
        doGetDescriptor, ServerCalls.asyncServerStreamingCall(new DoGetMethod(delegate)));
    serviceBuilder.addMethod(
        doPutDescriptor, ServerCalls.asyncBidiStreamingCall(new DoPutMethod(delegate)));
    serviceBuilder.addMethod(
        doExchangeDescriptor, ServerCalls.asyncBidiStreamingCall(new DoExchangeMethod(delegate)));

    // copy over not-overridden methods.
    for (ServerMethodDefinition<?, ?> definition : baseDefinition.getMethods()) {
      if (OVERRIDE_METHODS.contains(definition.getMethodDescriptor().getFullMethodName())) {
        continue;
      }

      serviceBuilder.addMethod(definition);
    }

    return serviceBuilder.build();
  }

  private static class DoGetMethod
      implements ServerCalls.ServerStreamingMethod<Flight.Ticket, ArrowMessage> {

    private final FlightService delegate;

    public DoGetMethod(FlightService delegate) {
      this.delegate = delegate;
    }

    @Override
    public void invoke(Flight.Ticket request, StreamObserver<ArrowMessage> responseObserver) {
      delegate.doGetCustom(request, responseObserver);
    }
  }

  private static class DoPutMethod
      implements ServerCalls.BidiStreamingMethod<ArrowMessage, PutResult> {
    private final FlightService delegate;

    public DoPutMethod(FlightService delegate) {
      this.delegate = delegate;
    }

    @Override
    public StreamObserver<ArrowMessage> invoke(StreamObserver<PutResult> responseObserver) {
      return delegate.doPutCustom(responseObserver);
    }
  }

  private static class DoExchangeMethod
      implements ServerCalls.BidiStreamingMethod<ArrowMessage, ArrowMessage> {
    private final FlightService delegate;

    public DoExchangeMethod(FlightService delegate) {
      this.delegate = delegate;
    }

    @Override
    public StreamObserver<ArrowMessage> invoke(StreamObserver<ArrowMessage> responseObserver) {
      return delegate.doExchangeCustom(responseObserver);
    }
  }
}