org.apache.arrow.flight.FlightServer Java Examples

public static Pair<Location, FlightServer.Builder> sslHelper(FlightServer.Builder serverBuilder, DremioConfig config, boolean useSsl, String hostname, int port, String sslHostname) {
  Location location;
  Location exLocation;
  try {
    if (!useSsl) {
      throw new UnsupportedOperationException("Don't use ssl");
    }

    Pair<InputStream, InputStream> pair = ssl(config, sslHostname);
    location = Location.forGrpcTls(hostname, port);
    exLocation = Location.forGrpcTls(sslHostname, port);
    serverBuilder.useTls(pair.getRight(), pair.getLeft()).location(location);
  } catch (Exception e) {
    location = Location.forGrpcInsecure(hostname, port);
    exLocation = Location.forGrpcInsecure(sslHostname, port);
    serverBuilder.location(location);
  }
  return Pair.of(exLocation, serverBuilder);
}
 

@Override
public Void initialize(BindingProvider provider) throws Exception {
  if (!Boolean.parseBoolean(PropertyHelper.getFromEnvProperty("dremio.flight.enabled", Boolean.toString(false)))) {
    logger.info("Flight plugin is not enabled, skipping initialization");
    return null;
  }
  this.allocator = provider.provider(BootStrapContext.class).get().getAllocator().newChildAllocator("arrow-flight", 0, Long.MAX_VALUE);


  AuthValidator validator = new AuthValidator(provider.provider(UserService.class), provider.provider(SabotContext.class));
  FlightServer.Builder serverBuilder = FlightServer.builder().allocator(allocator).authHandler(new BasicServerAuthHandler(validator));
  DremioConfig config = null;
  try {
    config = provider.lookup(DremioConfig.class);
  } catch (Throwable t) {
  }
  Pair<Location, FlightServer.Builder> pair = SslHelper.sslHelper(
    serverBuilder,
    config,
    useSsl,
    InetAddress.getLocalHost().getHostName(),
    port,
    host);
  producer = new Producer(
    pair.getKey(),
    provider.provider(UserWorker.class),
    provider.provider(SabotContext.class),
    allocator,
    validator);
  pair.getRight().producer(producer);
  server = pair.getRight().build();
  server.start();
  logger.info("set up flight plugin on port {} and host {}", pair.getKey().getUri().getPort(), pair.getKey().getUri().getHost());
  return null;
}
 

public FormationPlugin(SabotContext context, String name, Provider<StoragePluginId> pluginIdProvider) {
  this.context = context;
  this.pluginIdProvider = pluginIdProvider;
  if (!Boolean.parseBoolean(PropertyHelper.getFromEnvProperty("dremio.flight.parallel.enabled", Boolean.toString(false)))) { //todo add this and others to DremioConfig
    allocator = null;
    thisLocation = null;
    server = null;
    producer = null;
    validator = null;
    logger.info("Parallel flight plugin is not enabled, skipping initialization");
    return;
  }

  this.allocator = context.getAllocator().newChildAllocator("formation-" + name, 0, Long.MAX_VALUE);
  String hostname = PropertyHelper.getFromEnvProperty("dremio.flight.host", context.getEndpoint().getAddress());
  int port = Integer.parseInt(PropertyHelper.getFromEnvProperty("dremio.flight.port", Integer.toString(FLIGHT_PORT)));
  FlightServer.Builder serverBuilder = FlightServer.builder().allocator(this.allocator);
  Pair<Location, FlightServer.Builder> pair = SslHelper.sslHelper(
    serverBuilder,
    context.getDremioConfig(),
    Boolean.parseBoolean(PropertyHelper.getFromEnvProperty("dremio.formation.use-ssl", "false")),
    context.getEndpoint().getAddress(),
    port,
    hostname);
  thisLocation = pair.getKey();
  this.producer = new FormationFlightProducer(thisLocation, allocator);
  this.validator = new AuthValidator(context.isUserAuthenticationEnabled() ? context.getUserService() : null, context);
  this.server = pair.getRight().producer(producer).authHandler(new BasicServerAuthHandler(validator)).build();
  logger.info("set up formation plugin on port {} and host {}", thisLocation.getUri().getPort(), thisLocation.getUri().getHost());
}
 

@BeforeClass
public static void setUp() throws Exception {
  server = FlightTestUtil.getStartedServer(location -> FlightServer.builder(allocator, location, new TestProducer()).authHandler(
    new ServerAuthHandler() {
      @Override
      public Optional<String> isValid(byte[] token) {
        return Optional.of("xxx");
      }

      @Override
      public boolean authenticate(ServerAuthSender outgoing, Iterator<byte[]> incoming) {
        incoming.next();
        outgoing.send(new byte[0]);
        return true;
      }
    }).build()
  );
  location = server.getLocation();
  conf = new SparkConf()
    .setAppName("flightTest")
    .setMaster("local[*]")
    .set("spark.driver.allowMultipleContexts", "true")
    .set("spark.flight.endpoint.host", location.getUri().getHost())
    .set("spark.flight.endpoint.port", Integer.toString(location.getUri().getPort()))
    .set("spark.flight.auth.username", "xxx")
    .set("spark.flight.auth.password", "yyy")
  ;
  sc = new JavaSparkContext(conf);
  csc = FlightSparkContext.flightContext(sc);
}
 

@BeforeClass
public static void setup() throws IOException {
  server = FlightServer.builder()
    .allocator(allocator)
    .producer(new SimpleErrorProducer())
    .location(location)
    .build();
  server.start();
}