org.apache.arrow.flight.FlightClient Java Examples

@Test
public void connect() throws Exception {
  InetAddress ip = InetAddress.getLocalHost();
  Location location = Location.forGrpcInsecure(ip.getHostName(), 47470);
  try (FlightClient c = flightClient(getAllocator(), location)) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    String sql = "select * from sys.options";
    FlightInfo info = c.getInfo(FlightDescriptor.command(sql.getBytes()));
    long total = info.getEndpoints().stream()
      .map(this::submit)
      .map(TestFlightEndpoint::get)
      .mapToLong(Long::longValue)
      .sum();

    Assert.assertTrue(total > 1);
    System.out.println(total);
  }
}
 

@Override
public Long call() {
  long count = 0;
  int readIndex = 0;
  logger.debug("starting work on flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
  try (FlightClient c = flightClient(allocator, endpoint.getLocations().get(0))) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    logger.debug("trying to get stream for flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    FlightStream fs = c.getStream(endpoint.getTicket());
    logger.debug("got stream for flight endpoint with ticket {} to {}. Will now try and read", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    while (fs.next()) {
      long thisCount = fs.getRoot().getRowCount();
      count += thisCount;
      logger.debug("got results from stream for flight endpoint with ticket {} to {}. This is read {} and we got {} rows back for a total of {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, thisCount, count);
      fs.getRoot().clear();
    }
  } catch (InterruptedException e) {

  } catch (Throwable t) {
    logger.error("Error in stream fetch", t);
  }
  logger.debug("got all results from stream for flight endpoint with ticket {} to {}. We read {} batches and we got {} rows back", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, count);
  return count;
}
 

@Test
public void connect() throws Exception {
  certs();
  InetAddress ip = InetAddress.getLocalHost();
  Location location = Location.forGrpcTls(ip.getHostName(), 47470);
  try (FlightClient c = flightClient(getAllocator(), location)) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    String sql = "select * from sys.options";
    FlightInfo info = c.getInfo(FlightDescriptor.command(sql.getBytes()));
    long total = info.getEndpoints().stream()
      .map(this::submit)
      .map(TestSslFlightEndpoint::get)
      .mapToLong(Long::longValue)
      .sum();

    Assert.assertTrue(total > 1);
    System.out.println(total);
  }
}
 

@Override
public Long call() {
  long count = 0;
  int readIndex = 0;
  logger.debug("starting work on flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
  try (FlightClient c = flightClient(allocator, endpoint.getLocations().get(0))) {
    c.authenticate(new BasicClientAuthHandler(SystemUser.SYSTEM_USERNAME, null));
    logger.debug("trying to get stream for flight endpoint with ticket {} to {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    FlightStream fs = c.getStream(endpoint.getTicket());
    logger.debug("got stream for flight endpoint with ticket {} to {}. Will now try and read", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri());
    while (fs.next()) {
      long thisCount = fs.getRoot().getRowCount();
      count += thisCount;
      logger.debug("got results from stream for flight endpoint with ticket {} to {}. This is read {} and we got {} rows back for a total of {}", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, thisCount, count);
      fs.getRoot().clear();
    }
  } catch (InterruptedException e) {

  } catch (Throwable t) {
    logger.error("Error in stream fetch", t);
  }
  logger.debug("got all results from stream for flight endpoint with ticket {} to {}. We read {} batches and we got {} rows back", new String(endpoint.getTicket().getBytes()), endpoint.getLocations().get(0).getUri(), ++readIndex, count);
  return count;
}
 

public FlightDataSourceReader(Broadcast<FactoryOptions> dataSourceOptions) {
  clientFactory = new FlightClientFactory(
    dataSourceOptions.value().getLocation(),
    dataSourceOptions.value().getUsername(),
    dataSourceOptions.value().getPassword(),
    dataSourceOptions.value().isParallel()
  );
  defaultLocation = dataSourceOptions.value().getLocation();
  sql = dataSourceOptions.value().getSql();
  this.dataSourceOptions = dataSourceOptions;
  descriptor = getDescriptor(sql);
  try (FlightClient client = clientFactory.apply()) {
    info = client.getSchema(descriptor);
  } catch (InterruptedException e) {
    throw new RuntimeException(e);
  }
}
 

@Override
public Filter[] pushFilters(Filter[] filters) {
  List<Filter> notPushed = Lists.newArrayList();
  List<Filter> pushed = Lists.newArrayList();
  for (Filter filter : filters) {
    boolean isPushed = canBePushed(filter);
    if (isPushed) {
      pushed.add(filter);
    } else {
      notPushed.add(filter);
    }
  }
  this.pushed = pushed.toArray(new Filter[0]);
  if (!pushed.isEmpty()) {
    String whereClause = generateWhereClause(pushed);
    mergeWhereDescriptors(whereClause);
    try (FlightClient client = clientFactory.apply()) {
      info = client.getSchema(descriptor);
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }
  return notPushed.toArray(new Filter[0]);
}
 

private synchronized void refreshClients() {
  List<FlightClient> oldClients = clients;
  clients = context.getExecutors().stream()
    .map(e -> FlightClient.builder().allocator(allocator).location(Location.forGrpcInsecure(e.getAddress(), FLIGHT_PORT)).build()).collect(Collectors.toList());
  try {
    AutoCloseables.close(oldClients);
  } catch (Exception ex) {
    logger.error("Failure while refreshing clients.", ex);
  }
}
 

private static FlightClient flightClient(BufferAllocator allocator, Location location) {
  try {
    InputStream certStream = certs();
    return FlightClient.builder()
      .allocator(allocator)
      .location(location)
      .useTls()
      .trustedCertificates(certStream)
      .build();
  } catch (GeneralSecurityException | IOException e) {
    throw new RuntimeException(e);
  }
}
 

private List<InputPartition<ColumnarBatch>> planBatchInputPartitionsParallel() {

    try (FlightClient client = clientFactory.apply()) {
      FlightInfo info = client.getInfo(FlightDescriptor.command(sql.getBytes()));
      return planBatchInputPartitionsSerial(info);
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }
 

public FlightClient apply() {
  FlightClient client = FlightClient.builder(allocator, defaultLocation).build();
  client.authenticateBasic(username, password);
  if (parallel) {
    Iterator<Result> res = client.doAction(new Action("PARALLEL"));
    res.forEachRemaining(Object::toString);
  }
  return client;

}
 

/**
 * Factory for returning JobData over gRPC
 * @param jobsService reference to job service
 * @param bufferAllocator allocator for vectors
 * @param jobId jobid for corresponding jobresults
 * @param offset start index of results
 * @param limit max number of results to fetch
 * @return JobDataFragment
 */
public static JobDataFragment getJobData(JobsService jobsService, BufferAllocator bufferAllocator, JobId jobId, int offset, int limit) {
  final FlightClient flightClient = jobsService.getJobsClient().getFlightClient();
  final Ticket ticket = new JobsFlightTicket(jobId.getId(), offset, limit).toTicket();
  try (FlightStream flightStream = flightClient.getStream(ticket)) {
    return new JobDataFragmentImpl(new RecordBatches(JobDataClientUtils.getData(
      flightStream, bufferAllocator, limit)), offset, jobId);
  } catch (FlightRuntimeException fre) {
    Optional<UserException> ue = JobsRpcUtils.fromFlightRuntimeException(fre);
    throw ue.isPresent() ? ue.get() : fre;
  } catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

public DatasetBuilder(List<FlightClient> clients, EntityPath key) {
  super();
  this.clients = clients;
  this.key = key;
  buildIfNecessary();
}
 

public DatasetBuilder(List<FlightClient> clients, EntityPath key, List<FlightInfo> infos) {
  super();
  this.clients = clients;
  this.key = key;
  this.infos = infos;
}
 

private static FlightClient flightClient(BufferAllocator allocator, Location location) {
  return FlightClient.builder().allocator(allocator).location(location).build();
}
 

@Before
public void setupClient() {
  client = FlightClient.builder().allocator(allocator)
    .location(location)
    .build();
}
 

/**
 * Get the Arrow Flight Client for data requests to jobs service.
 *
 * @return Flight Client
 */
public FlightClient getFlightClient() {
  return flightClient;
}