mirrors/arrow
0
0 Fork 0
mirror of https://github.com/apache/arrow.git synced 2026-10-28 16:47:19 +08:00
Code Activity

Use URIs for Flight locations

Browse Source
This commit is contained in:
David Li committed 2019-05-20 10:10:13 -04:00
1 parent cd34f8b19e
commit 460442cea3
38 files changed
+541 -189
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cpp/src/arrow/flight/client.cc
+24 -8
View File
@@ -37,6 +37,7 @@
#include "arrow/status.h"
#include "arrow/type.h"
#include "arrow/util/logging.h"
#include "arrow/util/uri.h"
#include "arrow/flight/client_auth.h"
#include "arrow/flight/internal.h"
@@ -230,11 +231,26 @@ class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
class FlightClient::FlightClientImpl {
public:
Status Connect(const std::string& host, int port) {
// TODO(wesm): Support other kinds of GRPC ChannelCredentials
std::stringstream ss;
ss << host << ":" << port;
std::string uri = ss.str();
Status Connect(const Location& location) {
const std::string& scheme = location.scheme();
std::stringstream grpc_uri;
std::shared_ptr<grpc::ChannelCredentials> creds;
if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) {
// TODO(wesm): Support other kinds of GRPC ChannelCredentials
grpc_uri << location.uri_->host() << ":" << location.uri_->port_text();
if (scheme == "grpc+tls") {
creds = grpc::SslCredentials(grpc::SslCredentialsOptions());
} else {
creds = grpc::InsecureChannelCredentials();
}
} else if (scheme == kSchemeGrpcUnix) {
grpc_uri << "unix://" << location.uri_->path();
creds = grpc::InsecureChannelCredentials();
} else {
return Status::NotImplemented("Flight scheme " + scheme + " is not supported.");
}
grpc::ChannelArguments args;
// Try to reconnect quickly at first, in case the server is still starting up
@@ -242,7 +258,7 @@ class FlightClient::FlightClientImpl {
// Receive messages of any size
args.SetMaxReceiveMessageSize(-1);
stub_ = pb::FlightService::NewStub(
grpc::CreateCustomChannel(ss.str(), grpc::InsecureChannelCredentials(), args));
grpc::CreateCustomChannel(grpc_uri.str(), creds, args));
return Status::OK();
}
@@ -383,10 +399,10 @@ FlightClient::FlightClient() { impl_.reset(new FlightClientImpl); }
FlightClient::~FlightClient() {}
Status FlightClient::Connect(const std::string& host, int port,
Status FlightClient::Connect(const Location& location,
std::unique_ptr<FlightClient>* client) {
client->reset(new FlightClient);
return (*client)->impl_->Connect(host, port);
return (*client)->impl_->Connect(location);
}
Status FlightClient::Authenticate(const FlightCallOptions& options,
cpp/src/arrow/flight/client.h
+2 -4
View File
@@ -64,13 +64,11 @@ class ARROW_EXPORT FlightClient {
~FlightClient();
/// \brief Connect to an unauthenticated flight service
/// \param[in] host the hostname or IP address
/// \param[in] port the port on the host
/// \param[in] location the URI
/// \param[out] client the created FlightClient
/// \return Status OK status may not indicate that the connection was
/// successful
static Status Connect(const std::string& host, int port,
std::unique_ptr<FlightClient>* client);
static Status Connect(const Location& location, std::unique_ptr<FlightClient>* client);
/// \brief Authenticate to the server using the given handler.
/// \param[in] options Per-RPC options
cpp/src/arrow/flight/flight-benchmark.cc
+6 -2
View File
@@ -79,7 +79,9 @@ Status RunPerformanceTest(const std::string& hostname, const int port) {
// Construct client and plan the query
std::unique_ptr<FlightClient> client;
RETURN_NOT_OK(FlightClient::Connect(hostname, port, &client));
Location location;
RETURN_NOT_OK(Location::ForGrpcTcp("localhost", port, &location));
RETURN_NOT_OK(FlightClient::Connect(location, &client));
FlightDescriptor descriptor;
descriptor.type = FlightDescriptor::CMD;
@@ -97,7 +99,9 @@ Status RunPerformanceTest(const std::string& hostname, const int port) {
auto ConsumeStream = [&stats, &hostname, &port](const FlightEndpoint& endpoint) {
// TODO(wesm): Use location from endpoint, same host/port for now
std::unique_ptr<FlightClient> client;
RETURN_NOT_OK(FlightClient::Connect(hostname, port, &client));
Location location;
RETURN_NOT_OK(Location::ForGrpcTcp("localhost", port, &location));
RETURN_NOT_OK(FlightClient::Connect(location, &client));
perf::Token token;
token.ParseFromString(endpoint.ticket.ticket);
cpp/src/arrow/flight/flight-test.cc
+28 -5
View File
@@ -61,8 +61,7 @@ void AssertEqual(const Ticket& expected, const Ticket& actual) {
}
void AssertEqual(const Location& expected, const Location& actual) {
ASSERT_EQ(expected.host, actual.host);
ASSERT_EQ(expected.port, actual.port);
ASSERT_EQ(expected, actual);
}
void AssertEqual(const std::vector<FlightEndpoint>& expected,
@@ -146,6 +145,20 @@ TEST(TestFlight, StartStopTestServer) {
ASSERT_FALSE(server.IsRunning());
}
TEST(TestFlight, ConnectUri) {
TestServer server("flight-test-server", 30000);
server.Start();
ASSERT_TRUE(server.IsRunning());
std::unique_ptr<FlightClient> client;
Location location1;
Location location2;
ASSERT_OK(Location::Parse("grpc://localhost:30000", &location1));
ASSERT_OK(Location::Parse("grpc://localhost:30000", &location2));
ASSERT_OK(FlightClient::Connect(location1, &client));
ASSERT_OK(FlightClient::Connect(location2, &client));
}
// ----------------------------------------------------------------------
// Client tests
@@ -169,7 +182,11 @@ class TestFlightClient : public ::testing::Test {
void TearDown() { server_->Stop(); }
Status ConnectClient() { return FlightClient::Connect("localhost", port_, &client_); }
Status ConnectClient() {
Location location;
RETURN_NOT_OK(Location::ForGrpcTcp("localhost", port_, &location));
return FlightClient::Connect(location, &client_);
}
template <typename EndpointCheckFunc>
void CheckDoGet(const FlightDescriptor& descr, const BatchVector& expected_batches,
@@ -248,7 +265,11 @@ class TestAuthHandler : public ::testing::Test {
void TearDown() { server_->Stop(); }
Status ConnectClient() { return FlightClient::Connect("localhost", port_, &client_); }
Status ConnectClient() {
Location location;
RETURN_NOT_OK(Location::ForGrpcTcp("localhost", port_, &location));
return FlightClient::Connect(location, &client_);
}
protected:
int port_;
@@ -423,7 +444,9 @@ TEST_F(TestFlightClient, Issue5095) {
TEST_F(TestFlightClient, TimeoutFires) {
// Server does not exist on this port, so call should fail
std::unique_ptr<FlightClient> client;
ASSERT_OK(FlightClient::Connect("localhost", 30001, &client));
Location location;
ASSERT_OK(Location::ForGrpcTcp("localhost", 30001, &location));
ASSERT_OK(FlightClient::Connect(location, &client));
FlightCallOptions options;
options.timeout = TimeoutDuration{0.2};
std::unique_ptr<FlightInfo> info;
cpp/src/arrow/flight/internal.cc
+2 -5
View File
@@ -119,14 +119,11 @@ Status FromProto(const pb::Criteria& pb_criteria, Criteria* criteria) {
// Location
Status FromProto(const pb::Location& pb_location, Location* location) {
location->host = pb_location.host();
location->port = pb_location.port();
return Status::OK();
return Location::Parse(pb_location.uri(), location);
}
void ToProto(const Location& location, pb::Location* pb_location) {
pb_location->set_host(location.host);
pb_location->set_port(location.port);
pb_location->set_uri(location.ToString());
}
// Ticket
cpp/src/arrow/flight/perf-server.cc
+5 -2
View File
@@ -141,7 +141,8 @@ Status GetPerfBatches(const perf::Token& token, const std::shared_ptr<Schema>& s
class FlightPerfServer : public FlightServerBase {
public:
FlightPerfServer() : location_(Location{"localhost", FLAGS_port}) {
FlightPerfServer() : location_() {
DCHECK_OK(Location::ForGrpcTcp("localhost", FLAGS_port, &location_));
perf_schema_ = schema({field("a", int64()), field("b", int64()), field("c", int64()),
field("d", int64())});
}
@@ -204,8 +205,10 @@ int main(int argc, char** argv) {
g_server.reset(new arrow::flight::FlightPerfServer);
arrow::flight::Location location;
ARROW_CHECK_OK(arrow::flight::Location::ForGrpcTcp("0.0.0.0", FLAGS_port, &location));
ARROW_CHECK_OK(
g_server->Init(std::unique_ptr<arrow::flight::NoOpAuthHandler>(), FLAGS_port));
g_server->Init(std::unique_ptr<arrow::flight::NoOpAuthHandler>(), location));
// Exit with a clean error code (0) on SIGTERM
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
std::cout << "Server port: " << FLAGS_port << std::endl;
cpp/src/arrow/flight/server.cc
+21 -4
View File
@@ -40,6 +40,7 @@
#include "arrow/status.h"
#include "arrow/util/logging.h"
#include "arrow/util/stl.h"
#include "arrow/util/uri.h"
#include "arrow/flight/internal.h"
#include "arrow/flight/serialization-internal.h"
@@ -410,7 +411,6 @@ class FlightServiceImpl : public FlightService::Service {
#endif
struct FlightServerBase::Impl {
std::string address_;
std::unique_ptr<FlightServiceImpl> service_;
std::unique_ptr<grpc::Server> server_;
@@ -442,17 +442,34 @@ FlightServerBase::FlightServerBase() { impl_.reset(new Impl); }
FlightServerBase::~FlightServerBase() {}
Status FlightServerBase::Init(std::unique_ptr<ServerAuthHandler> auth_handler, int port) {
impl_->address_ = "localhost:" + std::to_string(port);
Status FlightServerBase::Init(std::unique_ptr<ServerAuthHandler> auth_handler,
const Location& location) {
std::shared_ptr<ServerAuthHandler> handler = std::move(auth_handler);
impl_->service_.reset(new FlightServiceImpl(handler, this));
grpc::ServerBuilder builder;
// Allow uploading messages of any length
builder.SetMaxReceiveMessageSize(-1);
builder.AddListeningPort(impl_->address_, grpc::InsecureServerCredentials());
const std::string scheme = location.scheme();
if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp) {
std::stringstream address;
address << location.uri_->host() << ':' << location.uri_->port_text();
builder.AddListeningPort(address.str(), grpc::InsecureServerCredentials());
} else if (scheme == kSchemeGrpcUnix) {
std::stringstream address;
address << "unix:" << location.uri_->path();
builder.AddListeningPort(address.str(), grpc::InsecureServerCredentials());
} else {
return Status::NotImplemented("Scheme is not supported: " + scheme);
}
builder.RegisterService(impl_->service_.get());
// Disable SO_REUSEPORT - it makes debugging/testing a pain as
// leftover processes can handle requests on accident
builder.AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0);
impl_->server_ = builder.BuildAndStart();
if (!impl_->server_) {
return Status::UnknownError("Server did not start properly");
cpp/src/arrow/flight/server.h
+3 -4
View File
@@ -98,13 +98,12 @@ class ARROW_EXPORT FlightServerBase {
// Lifecycle methods.
/// \brief Initialize an insecure TCP server listening on localhost
/// at the given port.
/// \brief Initialize a Flight server listening at the given location.
/// This method must be called before any other method.
/// \param[in] port The port to serve on.
/// \param[in] auth_handler The authentication handler. May be
/// nullptr if no authentication is desired.
Status Init(std::unique_ptr<ServerAuthHandler> auth_handler, int port);
/// \param[in] location The location to serve on.
Status Init(std::unique_ptr<ServerAuthHandler> auth_handler, const Location& location);
/// \brief Set the server to stop when receiving any of the given signal
/// numbers.
cpp/src/arrow/flight/test-integration-client.cc
+7 -6
View File
@@ -89,8 +89,7 @@ arrow::Status ConsumeFlightLocation(const arrow::flight::Location& location,
const std::shared_ptr<arrow::Schema>& schema,
std::shared_ptr<arrow::Table>* retrieved_data) {
std::unique_ptr<arrow::flight::FlightClient> read_client;
RETURN_NOT_OK(
arrow::flight::FlightClient::Connect(location.host, location.port, &read_client));
RETURN_NOT_OK(arrow::flight::FlightClient::Connect(location, &read_client));
std::unique_ptr<arrow::RecordBatchReader> stream;
RETURN_NOT_OK(read_client->DoGet(ticket, &stream));
@@ -103,7 +102,10 @@ int main(int argc, char** argv) {
gflags::ParseCommandLineFlags(&argc, &argv, true);
std::unique_ptr<arrow::flight::FlightClient> client;
ABORT_NOT_OK(arrow::flight::FlightClient::Connect(FLAGS_host, FLAGS_port, &client));
std::stringstream uri;
arrow::flight::Location location;
ABORT_NOT_OK(arrow::flight::Location::ForGrpcTcp(FLAGS_host, FLAGS_port, &location));
ABORT_NOT_OK(arrow::flight::FlightClient::Connect(location, &client));
arrow::flight::FlightDescriptor descr{
arrow::flight::FlightDescriptor::PATH, "", {FLAGS_path}};
@@ -143,12 +145,11 @@ int main(int argc, char** argv) {
auto locations = endpoint.locations;
if (locations.size() == 0) {
locations = {arrow::flight::Location{FLAGS_host, FLAGS_port}};
locations = {location};
}
for (const auto location : locations) {
std::cout << "Verifying location " << location.host << ':' << location.port
<< std::endl;
std::cout << "Verifying location " << location.ToString() << std::endl;
// 3. Download the data from the server.
std::shared_ptr<arrow::Table> retrieved_data;
ABORT_NOT_OK(ConsumeFlightLocation(location, ticket, schema, &retrieved_data));
cpp/src/arrow/flight/test-integration-server.cc
+3 -1
View File
@@ -124,7 +124,9 @@ int main(int argc, char** argv) {
gflags::ParseCommandLineFlags(&argc, &argv, true);
g_server.reset(new arrow::flight::FlightIntegrationTestServer);
ARROW_CHECK_OK(g_server->Init(nullptr, FLAGS_port));
arrow::flight::Location location;
ARROW_CHECK_OK(arrow::flight::Location::ForGrpcTcp("0.0.0.0", FLAGS_port, &location));
ARROW_CHECK_OK(g_server->Init(nullptr, location));
// Exit with a clean error code (0) on SIGTERM
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
cpp/src/arrow/flight/test-server.cc
+3 -1
View File
@@ -139,8 +139,10 @@ int main(int argc, char** argv) {
gflags::ParseCommandLineFlags(&argc, &argv, true);
g_server.reset(new arrow::flight::FlightTestServer);
arrow::flight::Location location;
ARROW_CHECK_OK(arrow::flight::Location::ForGrpcTcp("0.0.0.0", FLAGS_port, &location));
ARROW_CHECK_OK(
g_server->Init(std::unique_ptr<arrow::flight::NoOpAuthHandler>(), FLAGS_port));
g_server->Init(std::unique_ptr<arrow::flight::NoOpAuthHandler>(), location));
// Exit with a clean error code (0) on SIGTERM
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
cpp/src/arrow/flight/test-util.cc
+16 -5
View File
@@ -119,7 +119,9 @@ bool TestServer::IsRunning() { return server_process_->running(); }
int TestServer::port() const { return port_; }
Status InProcessTestServer::Start(std::unique_ptr<ServerAuthHandler> auth_handler) {
RETURN_NOT_OK(server_->Init(std::move(auth_handler), port_));
Location location;
RETURN_NOT_OK(Location::ForGrpcTcp("localhost", port_, &location));
RETURN_NOT_OK(server_->Init(std::move(auth_handler), location));
thread_ = std::thread([this]() { ARROW_EXPECT_OK(server_->Serve()); });
return Status::OK();
}
@@ -169,12 +171,21 @@ std::shared_ptr<Schema> ExampleDictSchema() {
}
std::vector<FlightInfo> ExampleFlightInfo() {
Location location1;
Location location2;
Location location3;
Location location4;
ARROW_EXPECT_OK(Location::ForGrpcTcp("foo1.bar.com", 12345, &location1));
ARROW_EXPECT_OK(Location::ForGrpcTcp("foo2.bar.com", 12345, &location2));
ARROW_EXPECT_OK(Location::ForGrpcTcp("foo3.bar.com", 12345, &location3));
ARROW_EXPECT_OK(Location::ForGrpcTcp("foo4.bar.com", 12345, &location4));
FlightInfo::Data flight1, flight2, flight3;
FlightEndpoint endpoint1({{"ticket-ints-1"}, {{"foo1.bar.com", 92385}}});
FlightEndpoint endpoint2({{"ticket-ints-2"}, {{"foo2.bar.com", 92385}}});
FlightEndpoint endpoint3({{"ticket-cmd"}, {{"foo3.bar.com", 92385}}});
FlightEndpoint endpoint4({{"ticket-dicts-1"}, {{"foo4.bar.com", 92385}}});
FlightEndpoint endpoint1({{"ticket-ints-1"}, {location1}});
FlightEndpoint endpoint2({{"ticket-ints-2"}, {location2}});
FlightEndpoint endpoint3({{"ticket-cmd"}, {location3}});
FlightEndpoint endpoint4({{"ticket-dicts-1"}, {location4}});
FlightDescriptor descr1{FlightDescriptor::PATH, "", {"examples", "ints"}};
FlightDescriptor descr2{FlightDescriptor::CMD, "my_command", {}};
cpp/src/arrow/flight/types.cc
+29
View File
@@ -25,6 +25,7 @@
#include "arrow/ipc/dictionary.h"
#include "arrow/ipc/reader.h"
#include "arrow/status.h"
#include "arrow/util/uri.h"
namespace arrow {
namespace flight {
@@ -83,6 +84,34 @@ Status FlightInfo::GetSchema(ipc::DictionaryMemo* dictionary_memo,
return Status::OK();
}
Location::Location() { uri_ = std::make_shared<arrow::internal::Uri>(); }
Status Location::Parse(const std::string& uri_string, Location* location) {
return location->uri_->Parse(uri_string);
}
Status Location::ForGrpcTcp(const std::string& host, const int port, Location* location) {
std::stringstream uri_string;
uri_string << "grpc+tcp://" << host << ':' << port;
return Location::Parse(uri_string.str(), location);
}
Status Location::ForGrpcUnix(const std::string& path, Location* location) {
std::stringstream uri_string;
uri_string << "grpc+unix://" << path;
return Location::Parse(uri_string.str(), location);
}
std::string Location::ToString() const { return uri_->to_string(); }
std::string Location::scheme() const {
std::string scheme = uri_->scheme();
if (scheme.empty()) {
// Default to grpc+tcp
return "grpc+tcp";
}
return scheme;
}
SimpleFlightListing::SimpleFlightListing(const std::vector<FlightInfo>& flights)
: position_(0), flights_(flights) {}
cpp/src/arrow/flight/types.h
+53 -4
View File
@@ -39,7 +39,13 @@ namespace ipc {
class DictionaryMemo;
}
} // namespace ipc
namespace internal {
class Uri;
} // namespace internal
namespace flight {
@@ -115,10 +121,53 @@ struct Ticket {
std::string ticket;
};
/// \brief A host location (hostname and port)
class FlightClient;
class FlightServerBase;
static const char* kSchemeGrpc = "grpc";
static const char* kSchemeGrpcTcp = "grpc+tcp";
static const char* kSchemeGrpcUnix = "grpc+unix";
static const char* kSchemeGrpcTls = "grpc+tls";
/// \brief A host location (a URI)
struct Location {
std::string host;
int32_t port;
public:
/// \brief Initialize a blank location.
Location();
/// \brief Initialize a location by parsing a URI string
static Status Parse(const std::string& uri_string, Location* location);
/// \brief Initialize a location for a non-TLS, gRPC-based Flight
/// service from a host and port
/// \param[in] host The hostname to connect to
/// \param[in] port The port
/// \param[out] location The resulting location
static Status ForGrpcTcp(const std::string& host, const int port, Location* location);
/// \brief Initialize a location for a domain socket-based Flight
/// service
/// \param[in] path The path to the domain socket
/// \param[out] location The resulting location
static Status ForGrpcUnix(const std::string& path, Location* location);
/// \brief Get a representation of this URI as a string.
std::string ToString() const;
/// \brief Get the scheme of this URI.
std::string scheme() const;
friend bool operator==(const Location& left, const Location& right) {
return left.ToString() == right.ToString();
}
friend bool operator!=(const Location& left, const Location& right) {
return !(left == right);
}
private:
friend class FlightClient;
friend class FlightServerBase;
std::shared_ptr<arrow::internal::Uri> uri_;
};
/// \brief A flight ticket and list of locations where the ticket can be
cpp/src/arrow/util/uri.cc
+5 -1
View File
@@ -52,7 +52,7 @@ bool IsTextRangeSet(const UriTextRangeStructA& range) { return range.first != nu
} // namespace
struct Uri::Impl {
Impl() : port_(-1) { memset(&uri_, 0, sizeof(uri_)); }
Impl() : string_rep_(""), port_(-1) { memset(&uri_, 0, sizeof(uri_)); }
~Impl() { uriFreeUriMembersA(&uri_); }
@@ -71,6 +71,7 @@ struct Uri::Impl {
UriUriA uri_;
// Keep alive strings that uriparser stores pointers to
std::vector<std::string> data_;
std::string string_rep_;
int32_t port_;
};
@@ -119,10 +120,13 @@ std::string Uri::path() const {
return ss.str();
}
const std::string& Uri::to_string() const { return impl_->string_rep_; }
Status Uri::Parse(const std::string& uri_string) {
impl_->Reset();
const auto& s = impl_->KeepString(uri_string);
impl_->string_rep_ = s;
const char* error_pos;
if (uriParseSingleUriExA(&impl_->uri_, s.data(), s.data() + s.size(), &error_pos) !=
URI_SUCCESS) {
cpp/src/arrow/util/uri.h
+2
View File
@@ -54,6 +54,8 @@ class ARROW_EXPORT Uri {
int32_t port() const;
/// The URI path component.
std::string path() const;
/// Get the string representation of this URI.
const std::string& to_string() const;
/// Factory function to parse a URI from its string representation.
Status Parse(const std::string& uri_string);
format/Flight.proto
+4 -5
View File
@@ -246,7 +246,7 @@ message FlightEndpoint {
Ticket ticket = 1;
/*
* A list of locations where this ticket can be redeemed. If the list is
* A list of URIs where this ticket can be redeemed. If the list is
* empty, the expectation is that the ticket can only be redeemed on the
* current service where the ticket was generated.
*/
@@ -254,12 +254,11 @@ message FlightEndpoint {
}
/*
* A location where a flight service will accept retrieval of a particular
* stream given a ticket.
* A location where a Flight service will accept retrieval of a particular
* stream given a ticket.
*/
message Location {
string host = 1;
int32 port = 2;
string uri = 1;
}
/*
java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java
+19 -4
View File
@@ -20,6 +20,8 @@ package org.apache.arrow.flight;
import static io.grpc.stub.ClientCalls.asyncClientStreamingCall;
import static io.grpc.stub.ClientCalls.asyncServerStreamingCall;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Iterator;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
@@ -74,8 +76,7 @@ public class FlightClient implements AutoCloseable {
*/
public FlightClient(BufferAllocator incomingAllocator, Location location) {
final ManagedChannelBuilder<?> channelBuilder =
ManagedChannelBuilder.forAddress(location.getHost(),
location.getPort())
ManagedChannelBuilder.forAddress(location.getUri().getHost(), location.getUri().getPort())
.maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS)
.maxInboundMessageSize(FlightServer.MAX_GRPC_MESSAGE_SIZE)
.usePlaintext();
@@ -97,7 +98,15 @@ public class FlightClient implements AutoCloseable {
public Iterable<FlightInfo> listFlights(Criteria criteria, CallOption... options) {
return ImmutableList.copyOf(CallOptions.wrapStub(blockingStub, options).listFlights(criteria.asCriteria()))
.stream()
.map(FlightInfo::new)
.map(t -> {
try {
return new FlightInfo(t);
} catch (URISyntaxException e) {
// We don't expect this will happen for conforming Flight implementations. For instance, a Java server
// itself wouldn't be able to construct an invalid Location.
throw new RuntimeException(e);
}
})
.collect(Collectors.toList());
}
@@ -177,7 +186,13 @@ public class FlightClient implements AutoCloseable {
* @param options RPC-layer hints for this call.
*/
public FlightInfo getInfo(FlightDescriptor descriptor, CallOption... options) {
return new FlightInfo(CallOptions.wrapStub(blockingStub, options).getFlightInfo(descriptor.toProtocol()));
try {
return new FlightInfo(CallOptions.wrapStub(blockingStub, options).getFlightInfo(descriptor.toProtocol()));
} catch (URISyntaxException e) {
// We don't expect this will happen for conforming Flight implementations. For instance, a Java server
// itself wouldn't be able to construct an invalid Location.
throw new RuntimeException(e);
}
}
/**
java/flight/src/main/java/org/apache/arrow/flight/FlightEndpoint.java
+8 -7
View File
@@ -17,6 +17,8 @@
package org.apache.arrow.flight;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
@@ -46,9 +48,11 @@ public class FlightEndpoint {
/**
* Constructs from the protocol buffer representation.
*/
public FlightEndpoint(Flight.FlightEndpoint flt) {
locations = flt.getLocationList().stream()
.map(t -> new Location(t)).collect(Collectors.toList());
public FlightEndpoint(Flight.FlightEndpoint flt) throws URISyntaxException {
locations = new ArrayList<>();
for (final Flight.Location location : flt.getLocationList()) {
locations.add(new Location(location.getUri()));
}
ticket = new Ticket(flt.getTicket());
}
@@ -68,10 +72,7 @@ public class FlightEndpoint {
.setTicket(ticket.toProtocol());
for (Location l : locations) {
b.addLocation(Flight.Location.newBuilder()
.setHost(l.getHost())
.setPort(l.getPort())
.build());
b.addLocation(l.toProtocol());
}
return b.build();
}
java/flight/src/main/java/org/apache/arrow/flight/FlightInfo.java
+7 -2
View File
@@ -19,8 +19,10 @@ package org.apache.arrow.flight;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
@@ -67,7 +69,7 @@ public class FlightInfo {
/**
* Constructs from the protocol buffer representation.
*/
FlightInfo(Flight.FlightInfo pbFlightInfo) {
FlightInfo(Flight.FlightInfo pbFlightInfo) throws URISyntaxException {
try {
final ByteBuffer schemaBuf = pbFlightInfo.getSchema().asReadOnlyByteBuffer();
schema = pbFlightInfo.getSchema().size() > 0 ?
@@ -78,7 +80,10 @@ public class FlightInfo {
throw new RuntimeException(e);
}
descriptor = new FlightDescriptor(pbFlightInfo.getFlightDescriptor());
endpoints = pbFlightInfo.getEndpointList().stream().map(t -> new FlightEndpoint(t)).collect(Collectors.toList());
endpoints = new ArrayList<>();
for (final Flight.FlightEndpoint endpoint : pbFlightInfo.getEndpointList()) {
endpoints.add(new FlightEndpoint(endpoint));
}
bytes = pbFlightInfo.getTotalBytes();
records = pbFlightInfo.getTotalRecords();
}
java/flight/src/main/java/org/apache/arrow/flight/FlightServer.java
+20 -4
View File
@@ -18,6 +18,7 @@
package org.apache.arrow.flight;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.concurrent.TimeUnit;
import org.apache.arrow.flight.auth.ServerAuthHandler;
@@ -26,8 +27,8 @@ import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.VectorSchemaRoot;
import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.ServerInterceptors;
import io.grpc.netty.NettyServerBuilder;
public class FlightServer implements AutoCloseable {
@@ -42,16 +43,31 @@ public class FlightServer implements AutoCloseable {
* Constructs a new instance.
*
* @param allocator The allocator to use for storing/copying arrow data.
* @param port The port to bind to.
* @param location The location to serve on.
* @param producer The underlying business logic for the server.
* @param authHandler The authorization handler for the server.
*/
public FlightServer(
BufferAllocator allocator,
int port,
Location location,
FlightProducer producer,
ServerAuthHandler authHandler) {
this.server = ServerBuilder.forPort(port)
final NettyServerBuilder builder;
switch (location.getUri().getScheme()) {
case LocationSchemes.GRPC_DOMAIN_SOCKET: {
// TODO: need reflection to check if domain sockets are available
throw new UnsupportedOperationException("Domain sockets are not available.");
}
case LocationSchemes.GRPC:
case LocationSchemes.GRPC_INSECURE: {
builder = NettyServerBuilder
.forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
break;
}
default:
throw new IllegalArgumentException("Scheme is not supported: " + location.getUri().getScheme());
}
this.server = builder
.maxInboundMessageSize(MAX_GRPC_MESSAGE_SIZE)
.addService(
ServerInterceptors.intercept(
java/flight/src/main/java/org/apache/arrow/flight/Location.java
+33 -17
View File
@@ -17,40 +17,56 @@
package org.apache.arrow.flight;
import java.net.URI;
import java.net.URISyntaxException;
import org.apache.arrow.flight.impl.Flight;
/** A URI where a Flight stream is available. */
public class Location {
private final String host;
private final int port;
private final URI uri;
/**
* Constructs a new instance.
*
* @param host the host containing a flight server
* @param port the port the server is listening on.
* @param uri the URI of the Flight service
*/
public Location(String host, int port) {
public Location(String uri) throws URISyntaxException {
super();
this.host = host;
this.port = port;
this.uri = new URI(uri);
}
Location(Flight.Location location) {
this.host = location.getHost();
this.port = location.getPort();
public Location(URI uri) {
super();
this.uri = uri;
}
public String getHost() {
return host;
public URI getUri() {
return uri;
}
public int getPort() {
return port;
/**
* Convert this Location into its protocol-level representation.
*/
public Flight.Location toProtocol() {
return Flight.Location.newBuilder().setUri(uri.toString()).build();
}
Flight.Location toProtocol() {
return Flight.Location.newBuilder().setHost(host).setPort(port).build();
/** Construct a URI for a Flight+gRPC server without transport security. */
public static Location forGrpcInsecure(String host, int port) {
try {
return new Location(new URI(LocationSchemes.GRPC_INSECURE, null, host, port, null, null, null));
} catch (URISyntaxException e) {
throw new IllegalArgumentException(e);
}
}
/** Construct a URI for a Flight+gRPC server with transport security. */
public static Location forGrpcTls(String host, int port) {
try {
return new Location(new URI(LocationSchemes.GRPC_TLS, null, host, port, null, null, null));
} catch (URISyntaxException e) {
throw new IllegalArgumentException(e);
}
}
}
java/flight/src/main/java/org/apache/arrow/flight/LocationSchemes.java
+28
View File
@@ -0,0 +1,28 @@
/*
* 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;
/**
* Constants representing well-known URI schemes for Flight services.
*/
public class LocationSchemes {
public static final String GRPC = "grpc";
public static final String GRPC_INSECURE = "grpc+tcp";
public static final String GRPC_DOMAIN_SOCKET = "grpc+unix";
public static final String GRPC_TLS = "grpc+tls";
}
java/flight/src/main/java/org/apache/arrow/flight/example/ExampleFlightServer.java
+2 -2
View File
@@ -46,7 +46,7 @@ public class ExampleFlightServer implements AutoCloseable {
this.allocator = allocator.newChildAllocator("flight-server", 0, Long.MAX_VALUE);
this.location = location;
this.mem = new InMemoryStore(this.allocator, location);
this.flightServer = new FlightServer(allocator, location.getPort(), mem, ServerAuthHandler.NO_OP);
this.flightServer = new FlightServer(allocator, location, mem, ServerAuthHandler.NO_OP);
}
public Location getLocation() {
@@ -75,7 +75,7 @@ public class ExampleFlightServer implements AutoCloseable {
*/
public static void main(String[] args) throws Exception {
final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
final ExampleFlightServer efs = new ExampleFlightServer(a, new Location("localhost", 12233));
final ExampleFlightServer efs = new ExampleFlightServer(a, Location.forGrpcInsecure("localhost", 12233));
efs.start();
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
try {
java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java
+5 -3
View File
@@ -19,6 +19,7 @@ package org.apache.arrow.flight.example.integration;
import java.io.File;
import java.io.IOException;
import java.net.URISyntaxException;
import java.util.Collections;
import java.util.List;
@@ -80,7 +81,8 @@ class IntegrationTestClient {
final int port = Integer.parseInt(cmd.getOptionValue("port", "31337"));
final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
final FlightClient client = new FlightClient(allocator, new Location(host, port));
final Location defaultLocation = Location.forGrpcInsecure(host, port);
final FlightClient client = new FlightClient(allocator, defaultLocation);
final String inputPath = cmd.getOptionValue("j");
@@ -114,10 +116,10 @@ class IntegrationTestClient {
// 3. Download the data from the server.
List<Location> locations = endpoint.getLocations();
if (locations.size() == 0) {
locations = Collections.singletonList(new Location(host, port));
locations = Collections.singletonList(defaultLocation);
}
for (Location location : locations) {
System.out.println("Verifying location " + location.getHost() + ":" + location.getPort());
System.out.println("Verifying location " + location.getUri());
FlightClient readClient = new FlightClient(allocator, location);
FlightStream stream = readClient.getStream(endpoint.getTicket());
VectorSchemaRoot downloadedRoot;
java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestServer.java
+1 -1
View File
@@ -43,7 +43,7 @@ class IntegrationTestServer {
final int port = Integer.parseInt(cmd.getOptionValue("port", "31337"));
final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
final ExampleFlightServer efs = new ExampleFlightServer(allocator, new Location("localhost", port));
final ExampleFlightServer efs = new ExampleFlightServer(allocator, Location.forGrpcInsecure("localhost", port));
efs.start();
// Print out message for integration test script
System.out.println("Server listening on localhost:" + port);
java/flight/src/test/java/org/apache/arrow/flight/TestBackPressure.java
+4 -3
View File
@@ -47,7 +47,7 @@ public class TestBackPressure {
try (
final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
final PerformanceTestServer server = FlightTestUtil.getStartedServer(
(port) -> (new PerformanceTestServer(a, new Location(FlightTestUtil.LOCALHOST, port))));
(port) -> (new PerformanceTestServer(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port))));
final FlightClient client = new FlightClient(a, server.getLocation())
) {
FlightStream fs1 = client.getStream(client.getInfo(
@@ -123,10 +123,11 @@ public class TestBackPressure {
BufferAllocator serverAllocator = allocator.newChildAllocator("server", 0, Long.MAX_VALUE);
FlightServer server =
FlightTestUtil.getStartedServer(
(port) -> new FlightServer(serverAllocator, port, producer, ServerAuthHandler.NO_OP));
(port) -> new FlightServer(serverAllocator, Location.forGrpcInsecure("localhost", port), producer,
ServerAuthHandler.NO_OP));
BufferAllocator clientAllocator = allocator.newChildAllocator("client", 0, Long.MAX_VALUE);
FlightClient client =
new FlightClient(clientAllocator, new Location(FlightTestUtil.LOCALHOST, server.getPort()))
new FlightClient(clientAllocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
) {
FlightStream stream = client.getStream(new Ticket(new byte[1]));
VectorSchemaRoot root = stream.getRoot();
java/flight/src/test/java/org/apache/arrow/flight/TestBasicOperation.java
+16 -4
View File
@@ -17,6 +17,7 @@
package org.apache.arrow.flight;
import java.net.URISyntaxException;
import java.util.concurrent.Callable;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
@@ -138,10 +139,12 @@ public class TestBasicOperation {
BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
Producer producer = new Producer(a);
FlightServer s =
FlightTestUtil.getStartedServer((port) -> new FlightServer(a, port, producer, ServerAuthHandler.NO_OP))) {
FlightTestUtil.getStartedServer(
(port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
ServerAuthHandler.NO_OP))) {
try (
FlightClient c = new FlightClient(a, new Location(FlightTestUtil.LOCALHOST, s.getPort()));
FlightClient c = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()));
) {
try (BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE)) {
consumer.accept(c, testAllocator);
@@ -170,7 +173,12 @@ public class TestBasicOperation {
.setType(DescriptorType.CMD)
.setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8)))
.build();
listener.onNext(new FlightInfo(getInfo));
try {
listener.onNext(new FlightInfo(getInfo));
} catch (URISyntaxException e) {
listener.onError(e);
return;
}
listener.onCompleted();
}
@@ -231,7 +239,11 @@ public class TestBasicOperation {
.setType(DescriptorType.CMD)
.setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8)))
.build();
return new FlightInfo(getInfo);
try {
return new FlightInfo(getInfo);
} catch (URISyntaxException e) {
throw new RuntimeException(e);
}
}
@Override
java/flight/src/test/java/org/apache/arrow/flight/TestCallOptions.java
+4 -2
View File
@@ -70,8 +70,10 @@ public class TestCallOptions {
BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
Producer producer = new Producer(a);
FlightServer s =
FlightTestUtil.getStartedServer((port) -> new FlightServer(a, port, producer, ServerAuthHandler.NO_OP));
FlightClient client = new FlightClient(a, new Location(FlightTestUtil.LOCALHOST, s.getPort()))) {
FlightTestUtil.getStartedServer(
(port) -> new FlightServer(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port), producer,
ServerAuthHandler.NO_OP));
FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))) {
testFn.accept(client);
} catch (InterruptedException | IOException e) {
throw new RuntimeException(e);
java/flight/src/test/java/org/apache/arrow/flight/TestLargeMessage.java
+8 -4
View File
@@ -44,9 +44,11 @@ public class TestLargeMessage {
try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
final Producer producer = new Producer(a);
final FlightServer s =
FlightTestUtil.getStartedServer((port) -> new FlightServer(a, port, producer, ServerAuthHandler.NO_OP))) {
FlightTestUtil.getStartedServer(
(port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
ServerAuthHandler.NO_OP))) {
try (FlightClient client = new FlightClient(a, new Location(FlightTestUtil.LOCALHOST, s.getPort()))) {
try (FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))) {
FlightStream stream = client.getStream(new Ticket(new byte[]{}));
try (VectorSchemaRoot root = stream.getRoot()) {
while (stream.next()) {
@@ -73,9 +75,11 @@ public class TestLargeMessage {
try (final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
final Producer producer = new Producer(a);
final FlightServer s =
FlightTestUtil.getStartedServer((port) -> new FlightServer(a, port, producer, ServerAuthHandler.NO_OP))) {
FlightTestUtil.getStartedServer(
(port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
ServerAuthHandler.NO_OP))) {
try (FlightClient client = new FlightClient(a, new Location(FlightTestUtil.LOCALHOST, s.getPort()));
try (FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()));
BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE);
VectorSchemaRoot root = generateData(testAllocator)) {
final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root);
java/flight/src/test/java/org/apache/arrow/flight/auth/TestAuth.java
+2 -2
View File
@@ -121,7 +121,7 @@ public class TestAuth {
server = FlightTestUtil.getStartedServer((port) -> new FlightServer(
allocator,
port,
Location.forGrpcInsecure("localhost", port),
new NoOpFlightProducer() {
@Override
public void listFlights(CallContext context, Criteria criteria,
@@ -151,7 +151,7 @@ public class TestAuth {
}
},
new BasicServerAuthHandler(validator)));
client = new FlightClient(allocator, new Location(FlightTestUtil.LOCALHOST, server.getPort()));
client = new FlightClient(allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()));
}
@After
java/flight/src/test/java/org/apache/arrow/flight/example/TestExampleServer.java
+3 -1
View File
@@ -18,12 +18,14 @@
package org.apache.arrow.flight.example;
import java.io.IOException;
import java.net.URISyntaxException;
import org.apache.arrow.flight.FlightClient;
import org.apache.arrow.flight.FlightClient.ClientStreamListener;
import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.FlightTestUtil;
import org.apache.arrow.flight.Location;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
@@ -49,7 +51,7 @@ public class TestExampleServer {
public void start() throws IOException {
allocator = new RootAllocator(Long.MAX_VALUE);
Location l = new Location("localhost", 12233);
Location l = Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, 12233);
if (!Boolean.getBoolean("disableServer")) {
System.out.println("Starting server.");
server = new ExampleFlightServer(allocator, l);
java/flight/src/test/java/org/apache/arrow/flight/perf/PerformanceTestServer.java
+1 -1
View File
@@ -64,7 +64,7 @@ public class PerformanceTestServer implements AutoCloseable {
this.allocator = incomingAllocator.newChildAllocator("perf-server", 0, Long.MAX_VALUE);
this.location = location;
this.producer = new PerfProducer();
this.flightServer = new FlightServer(this.allocator, location.getPort(), producer, ServerAuthHandler.NO_OP);
this.flightServer = new FlightServer(this.allocator, location, producer, ServerAuthHandler.NO_OP);
}
public Location getLocation() {
java/flight/src/test/java/org/apache/arrow/flight/perf/TestPerf.java
+1 -1
View File
@@ -82,7 +82,7 @@ public class TestPerf {
final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
final PerformanceTestServer server =
FlightTestUtil.getStartedServer((port) -> new PerformanceTestServer(a,
new Location(FlightTestUtil.LOCALHOST, port)));
Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port)));
final FlightClient client = new FlightClient(a, server.getLocation());
) {
final FlightInfo info = client.getInfo(getPerfFlightDescriptor(50_000_000L, 4095, 2));
python/examples/flight/client.py
+1 -1
View File
@@ -131,7 +131,7 @@ def main():
}
host, port = args.host.split(':')
port = int(port)
client = pyarrow.flight.FlightClient.connect(host, port)
client = pyarrow.flight.FlightClient.connect(f"grpc://{host}:{port}")
while True:
try:
action = pyarrow.flight.Action("healthcheck", b"")
python/pyarrow/_flight.pyx
+100 -36
View File
@@ -20,6 +20,8 @@
import collections
import enum
import six
from cython.operator cimport dereference as deref
from pyarrow.compat import frombytes, tobytes
@@ -69,6 +71,13 @@ cdef class Action:
def body(self):
return pyarrow_wrap_buffer(self.action.body)
@staticmethod
cdef CAction unwrap(action):
if not isinstance(action, Action):
raise TypeError("Must provide Action, not '{}'".format(
type(action)))
return (<Action> action).action
_ActionType = collections.namedtuple('_ActionType', ['type', 'description'])
@@ -158,18 +167,73 @@ cdef class FlightDescriptor:
def __repr__(self):
return "<FlightDescriptor type: {!r}>".format(self.descriptor_type)
@staticmethod
cdef FlightDescriptor unwrap(descriptor):
if not isinstance(descriptor, FlightDescriptor):
raise TypeError("Must provide a FlightDescriptor, not '{}'".format(
type(descriptor)))
return <FlightDescriptor> descriptor
class Ticket:
cdef class Ticket:
"""A ticket for requesting a Flight stream."""
cdef:
CTicket ticket
def __init__(self, ticket):
self.ticket = ticket
self.ticket.ticket = tobytes(ticket)
def __repr__(self):
return '<Ticket {}>'.format(self.ticket)
return '<Ticket {}>'.format(self.ticket.ticket)
class Location(collections.namedtuple('Location', ['host', 'port'])):
"""A location where a Flight stream is available."""
cdef class Location:
"""The location of a Flight service."""
cdef:
CLocation location
def __init__(self, uri):
check_status(CLocation.Parse(tobytes(uri), &self.location))
def __repr__(self):
return '<Location {}>'.format(self.location.ToString())
@staticmethod
def for_grpc_tcp(host, port):
"""Create a Location for a TCP-based gRPC service."""
cdef:
c_string c_host = tobytes(host)
int c_port = port
Location result = Location.__new__(Location)
check_status(CLocation.ForGrpcTcp(c_host, c_port, &result.location))
return result
@staticmethod
def for_grpc_unix(path):
"""Create a Location for a domain socket-based gRPC service."""
cdef:
c_string c_path = tobytes(path)
Location result = Location.__new__(Location)
check_status(CLocation.ForGrpcUnix(c_path, &result.location))
return result
@staticmethod
cdef Location wrap(CLocation location):
cdef Location result = Location.__new__(Location)
result.location = location
return result
@staticmethod
cdef CLocation unwrap(object location):
cdef CLocation c_location
if isinstance(location, (six.text_type, six.binary_type)):
check_status(CLocation.Parse(tobytes(location), &c_location))
return c_location
elif not isinstance(location, Location):
raise TypeError("Must provide a Location, not '{}'".format(
type(location)))
return (<Location> location).location
cdef class FlightEndpoint:
@@ -184,11 +248,16 @@ cdef class FlightEndpoint:
----------
ticket : Ticket or bytes
the ticket needed to access this flight
locations : list of Location or tuples of (host, port)
locations : list of string URIs
locations where this flight is available
Raises
------
ArrowException
If one of the location URIs is not a valid URI.
"""
cdef:
CLocation c_location = CLocation()
CLocation c_location
if isinstance(ticket, Ticket):
self.endpoint.ticket.ticket = tobytes(ticket.ticket)
@@ -196,9 +265,8 @@ cdef class FlightEndpoint:
self.endpoint.ticket.ticket = tobytes(ticket)
for location in locations:
# Accepts Location namedtuple or tuple
c_location.host = tobytes(location[0])
c_location.port = location[1]
c_location = CLocation()
check_status(CLocation.Parse(tobytes(location), &c_location))
self.endpoint.locations.push_back(c_location)
@property
@@ -207,7 +275,7 @@ cdef class FlightEndpoint:
@property
def locations(self):
return [Location(frombytes(location.host), location.port)
return [Location.wrap(location)
for location in self.endpoint.locations]
@@ -313,27 +381,21 @@ cdef class FlightClient:
.format(self.__class__.__name__))
@staticmethod
def connect(*args):
def connect(location):
"""Connect to a Flight service on the given host and port."""
cdef:
FlightClient result = FlightClient.__new__(FlightClient)
int c_port = 0
c_string c_host
CLocation c_location
if len(args) == 1:
# Accept namedtuple or plain tuple
c_host = tobytes(args[0][0])
c_port = args[0][1]
elif len(args) == 2:
# Accept separate host, port
c_host = tobytes(args[0])
c_port = args[1]
if isinstance(location, Location):
c_location = (<Location> location).location
else:
raise TypeError("FlightClient.connect() takes 1 "
"or 2 arguments ({} given)".format(len(args)))
c_location = CLocation()
check_status(CLocation.Parse(tobytes(location), &c_location))
with nogil:
check_status(CFlightClient.Connect(c_host, c_port, &result.client))
check_status(CFlightClient.Connect(c_location, &result.client))
return result
@@ -375,10 +437,12 @@ cdef class FlightClient:
cdef:
unique_ptr[CResultStream] results
Result result
CAction c_action = Action.unwrap(action)
CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
with nogil:
check_status(self.client.get().DoAction(deref(c_options),
action.action, &results))
check_status(
self.client.get().DoAction(deref(c_options), c_action,
&results))
while True:
result = Result.__new__(Result)
@@ -414,25 +478,23 @@ cdef class FlightClient:
cdef:
FlightInfo result = FlightInfo.__new__(FlightInfo)
CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
FlightDescriptor c_descriptor = \
FlightDescriptor.unwrap(descriptor)
with nogil:
check_status(self.client.get().GetFlightInfo(
deref(c_options), descriptor.descriptor, &result.info))
deref(c_options), c_descriptor.descriptor, &result.info))
return result
def do_get(self, ticket: Ticket, options: FlightCallOptions = None):
"""Request the data for a flight."""
cdef:
# TODO: introduce unwrap
CTicket c_ticket
unique_ptr[CRecordBatchReader] reader
CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
c_ticket.ticket = ticket.ticket
with nogil:
check_status(
self.client.get().DoGet(deref(c_options), c_ticket, &reader))
check_status(self.client.get().DoGet(deref(c_options), ticket.ticket, &reader))
result = FlightRecordBatchReader()
result.reader.reset(reader.release())
return result
@@ -444,10 +506,12 @@ cdef class FlightClient:
shared_ptr[CSchema] c_schema = pyarrow_unwrap_schema(schema)
unique_ptr[CRecordBatchWriter] writer
CFlightCallOptions* c_options = FlightCallOptions.unwrap(options)
FlightDescriptor c_descriptor = \
FlightDescriptor.unwrap(descriptor)
with nogil:
check_status(self.client.get().DoPut(
deref(c_options), descriptor.descriptor, c_schema, &writer))
deref(c_options), c_descriptor.descriptor, c_schema, &writer))
result = FlightRecordBatchWriter()
result.writer.reset(writer.release())
return result
@@ -921,11 +985,11 @@ cdef class FlightServerBase:
cdef:
unique_ptr[PyFlightServer] server
def run(self, port, auth_handler=None):
def run(self, location, auth_handler=None):
cdef:
PyFlightServerVtable vtable = PyFlightServerVtable()
int c_port = port
PyFlightServer* c_server
CLocation c_location = Location.unwrap(location)
unique_ptr[CServerAuthHandler] c_auth_handler
if auth_handler:
@@ -945,7 +1009,7 @@ cdef class FlightServerBase:
c_server = new PyFlightServer(self, vtable)
self.server.reset(c_server)
with nogil:
check_status(c_server.Init(move(c_auth_handler), c_port))
check_status(c_server.Init(move(c_auth_handler), c_location))
check_status(c_server.ServeWithSignals())
def list_flights(self, context, criteria):
python/pyarrow/includes/libarrow_flight.pxd
+10 -4
View File
@@ -80,9 +80,14 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
cdef cppclass CLocation" arrow::flight::Location":
CLocation()
c_string ToString()
c_string host
int32_t port
@staticmethod
CStatus Parse(c_string& uri_string, CLocation* location)
@staticmethod
CStatus ForGrpcTcp(c_string& host, int port, CLocation* location)
@staticmethod
CStatus ForGrpcUnix(c_string& path, CLocation* location)
cdef cppclass CFlightEndpoint" arrow::flight::FlightEndpoint":
CFlightEndpoint()
@@ -150,7 +155,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
cdef cppclass CFlightClient" arrow::flight::FlightClient":
@staticmethod
CStatus Connect(const c_string& host, int port,
CStatus Connect(const CLocation& location,
unique_ptr[CFlightClient]* client)
CStatus Authenticate(CFlightCallOptions& options,
@@ -224,7 +229,8 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
cdef cppclass PyFlightServer:
PyFlightServer(object server, PyFlightServerVtable vtable)
CStatus Init(unique_ptr[CServerAuthHandler] auth_handler, int port)
CStatus Init(unique_ptr[CServerAuthHandler] auth_handler,
const CLocation& location)
CStatus ServeWithSignals() except *
void Shutdown()
python/pyarrow/tests/test_flight.py
+55 -33
View File
@@ -19,6 +19,7 @@
import base64
import contextlib
import socket
import tempfile
import threading
import time
@@ -206,26 +207,30 @@ class TokenClientAuthHandler(flight.ClientAuthHandler):
@contextlib.contextmanager
def flight_server(server_base, *args, **kwargs):
"""Spawn a Flight server on a free port, shutting it down when done."""
# Find a free port
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
with contextlib.closing(sock) as sock:
sock.bind(('', 0))
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
port = sock.getsockname()[1]
auth_handler = kwargs.pop('auth_handler', None)
location = kwargs.pop('location', None)
if location is None:
# Find a free port
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
with contextlib.closing(sock) as sock:
sock.bind(('', 0))
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
port = sock.getsockname()[1]
location = flight.Location.for_grpc_tcp("localhost", port)
else:
port = None
auth_handler = kwargs.get('auth_handler')
ctor_kwargs = kwargs
if auth_handler:
del ctor_kwargs['auth_handler']
server_instance = server_base(*args, **ctor_kwargs)
def _server_thread():
server_instance.run(port, auth_handler=auth_handler)
server_instance.run(location, auth_handler=auth_handler)
thread = threading.Thread(target=_server_thread, daemon=True)
thread.start()
yield port
yield location
server_instance.shutdown()
thread.join()
@@ -244,12 +249,29 @@ def test_flight_do_get_ints():
def test_flight_do_get_dicts():
table = simple_dicts_table()
with flight_server(ConstantFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(ConstantFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
data = client.do_get(flight.Ticket(b'dicts')).read_all()
assert data.equals(table)
def test_flight_domain_socket():
"""Try a simple do_get call over a domain socket."""
data = [
pa.array([-10, -5, 0, 5, 10])
]
table = pa.Table.from_arrays(data, names=['a'])
with tempfile.NamedTemporaryFile() as sock:
sock.close()
location = flight.Location.for_grpc_unix(sock.name)
with flight_server(ConstantFlightServer,
location=location) as server_location:
client = flight.FlightClient.connect(server_location)
data = client.do_get(flight.Ticket(b'')).read_all()
assert data.equals(table)
@pytest.mark.slow
def test_flight_large_message():
"""Try sending/receiving a large message via Flight.
@@ -261,8 +283,8 @@ def test_flight_large_message():
pa.array(range(0, 10 * 1024 * 1024))
], names=['a'])
with flight_server(EchoFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(EchoFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
writer = client.do_put(flight.FlightDescriptor.for_path('test'),
data.schema)
# Write a single giant chunk
@@ -278,8 +300,8 @@ def test_flight_generator_stream():
pa.array(range(0, 10 * 1024))
], names=['a'])
with flight_server(EchoStreamFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(EchoStreamFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
writer = client.do_put(flight.FlightDescriptor.for_path('test'),
data.schema)
writer.write_table(data)
@@ -290,8 +312,8 @@ def test_flight_generator_stream():
def test_flight_invalid_generator_stream():
"""Try streaming data with mismatched schemas."""
with flight_server(InvalidStreamFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(InvalidStreamFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
with pytest.raises(pa.ArrowException):
client.do_get(flight.Ticket(b'')).read_all()
@@ -300,8 +322,8 @@ def test_timeout_fires():
"""Make sure timeouts fire on slow requests."""
# Do this in a separate thread so that if it fails, we don't hang
# the entire test process
with flight_server(SlowFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(SlowFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
action = flight.Action("", b"")
options = flight.FlightCallOptions(timeout=0.2)
with pytest.raises(pa.ArrowIOError, match="Deadline Exceeded"):
@@ -310,8 +332,8 @@ def test_timeout_fires():
def test_timeout_passes():
"""Make sure timeouts do not fire on fast requests."""
with flight_server(ConstantFlightServer) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
with flight_server(ConstantFlightServer) as server_location:
client = flight.FlightClient.connect(server_location)
options = flight.FlightCallOptions(timeout=0.2)
client.do_get(flight.Ticket(b'ints'), options=options).read_all()
@@ -328,8 +350,8 @@ token_auth_handler = TokenServerAuthHandler(creds={
def test_http_basic_unauth():
"""Test that auth fails when not authenticated."""
with flight_server(EchoStreamFlightServer,
auth_handler=basic_auth_handler) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
auth_handler=basic_auth_handler) as server_location:
client = flight.FlightClient.connect(server_location)
action = flight.Action("who-am-i", b"")
with pytest.raises(pa.ArrowException, match=".*unauthenticated.*"):
list(client.do_action(action))
@@ -338,8 +360,8 @@ def test_http_basic_unauth():
def test_http_basic_auth():
"""Test a Python implementation of HTTP basic authentication."""
with flight_server(EchoStreamFlightServer,
auth_handler=basic_auth_handler) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
auth_handler=basic_auth_handler) as server_location:
client = flight.FlightClient.connect(server_location)
action = flight.Action("who-am-i", b"")
client.authenticate(HttpBasicClientAuthHandler('test', 'p4ssw0rd'))
identity = next(client.do_action(action))
@@ -349,8 +371,8 @@ def test_http_basic_auth():
def test_http_basic_auth_invalid_password():
"""Test that auth fails with the wrong password."""
with flight_server(EchoStreamFlightServer,
auth_handler=basic_auth_handler) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
auth_handler=basic_auth_handler) as server_location:
client = flight.FlightClient.connect(server_location)
action = flight.Action("who-am-i", b"")
client.authenticate(HttpBasicClientAuthHandler('test', 'wrong'))
with pytest.raises(pa.ArrowException, match=".*wrong password.*"):
@@ -360,8 +382,8 @@ def test_http_basic_auth_invalid_password():
def test_token_auth():
"""Test an auth mechanism that uses a handshake."""
with flight_server(EchoStreamFlightServer,
auth_handler=token_auth_handler) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
auth_handler=token_auth_handler) as server_location:
client = flight.FlightClient.connect(server_location)
action = flight.Action("who-am-i", b"")
client.authenticate(TokenClientAuthHandler('test', 'p4ssw0rd'))
identity = next(client.do_action(action))
@@ -371,7 +393,7 @@ def test_token_auth():
def test_token_auth_invalid():
"""Test an auth mechanism that uses a handshake."""
with flight_server(EchoStreamFlightServer,
auth_handler=token_auth_handler) as server_port:
client = flight.FlightClient.connect('localhost', server_port)
auth_handler=token_auth_handler) as server_location:
client = flight.FlightClient.connect(server_location)
with pytest.raises(pa.ArrowException, match=".*unauthenticated.*"):
client.authenticate(TokenClientAuthHandler('test', 'wrong'))