first commit

This commit is contained in:
Ayxan
2022-05-23 00:16:32 +04:00
commit d660f2a4ca
24786 changed files with 4428337 additions and 0 deletions

View File

@@ -0,0 +1,27 @@
// 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.
#pragma once
#include "arrow/flight/client.h"
#include "arrow/flight/client_auth.h"
#include "arrow/flight/client_middleware.h"
#include "arrow/flight/middleware.h"
#include "arrow/flight/server.h"
#include "arrow/flight/server_auth.h"
#include "arrow/flight/server_middleware.h"
#include "arrow/flight/types.h"

View File

@@ -0,0 +1,434 @@
// 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.
/// \brief Implementation of Flight RPC client. API should be
/// considered experimental for now
#pragma once
#include <chrono>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "arrow/ipc/options.h"
#include "arrow/ipc/reader.h"
#include "arrow/ipc/writer.h"
#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/util/cancel.h"
#include "arrow/util/variant.h"
#include "arrow/flight/type_fwd.h"
#include "arrow/flight/types.h" // IWYU pragma: keep
#include "arrow/flight/visibility.h"
namespace arrow {
class RecordBatch;
class Schema;
namespace flight {
/// \brief A duration type for Flight call timeouts.
typedef std::chrono::duration<double, std::chrono::seconds::period> TimeoutDuration;
/// \brief Hints to the underlying RPC layer for Arrow Flight calls.
class ARROW_FLIGHT_EXPORT FlightCallOptions {
public:
/// Create a default set of call options.
FlightCallOptions();
/// \brief An optional timeout for this call. Negative durations
/// mean an implementation-defined default behavior will be used
/// instead. This is the default value.
TimeoutDuration timeout;
/// \brief IPC reader options, if applicable for the call.
ipc::IpcReadOptions read_options;
/// \brief IPC writer options, if applicable for the call.
ipc::IpcWriteOptions write_options;
/// \brief Headers for client to add to context.
std::vector<std::pair<std::string, std::string>> headers;
/// \brief A token to enable interactive user cancellation of long-running requests.
StopToken stop_token;
/// \brief An optional memory manager to control where to allocate incoming data.
std::shared_ptr<MemoryManager> memory_manager;
};
/// \brief Indicate that the client attempted to write a message
/// larger than the soft limit set via write_size_limit_bytes.
class ARROW_FLIGHT_EXPORT FlightWriteSizeStatusDetail : public arrow::StatusDetail {
public:
explicit FlightWriteSizeStatusDetail(int64_t limit, int64_t actual)
: limit_(limit), actual_(actual) {}
const char* type_id() const override;
std::string ToString() const override;
int64_t limit() const { return limit_; }
int64_t actual() const { return actual_; }
/// \brief Extract this status detail from a status, or return
/// nullptr if the status doesn't contain this status detail.
static std::shared_ptr<FlightWriteSizeStatusDetail> UnwrapStatus(
const arrow::Status& status);
private:
int64_t limit_;
int64_t actual_;
};
struct ARROW_FLIGHT_EXPORT FlightClientOptions {
/// \brief Root certificates to use for validating server
/// certificates.
std::string tls_root_certs;
/// \brief Override the hostname checked by TLS. Use with caution.
std::string override_hostname;
/// \brief The client certificate to use if using Mutual TLS
std::string cert_chain;
/// \brief The private key associated with the client certificate for Mutual TLS
std::string private_key;
/// \brief A list of client middleware to apply.
std::vector<std::shared_ptr<ClientMiddlewareFactory>> middleware;
/// \brief A soft limit on the number of bytes to write in a single
/// batch when sending Arrow data to a server.
///
/// Used to help limit server memory consumption. Only enabled if
/// positive. When enabled, FlightStreamWriter.Write* may yield a
/// IOError with error detail FlightWriteSizeStatusDetail.
int64_t write_size_limit_bytes = 0;
/// \brief Generic connection options, passed to the underlying
/// transport; interpretation is implementation-dependent.
std::vector<std::pair<std::string, util::Variant<int, std::string>>> generic_options;
/// \brief Use TLS without validating the server certificate. Use with caution.
bool disable_server_verification = false;
/// \brief Get default options.
static FlightClientOptions Defaults();
};
/// \brief A RecordBatchReader exposing Flight metadata and cancel
/// operations.
class ARROW_FLIGHT_EXPORT FlightStreamReader : public MetadataRecordBatchReader {
public:
/// \brief Try to cancel the call.
virtual void Cancel() = 0;
using MetadataRecordBatchReader::ToRecordBatches;
/// \brief Consume entire stream as a vector of record batches
virtual arrow::Result<std::vector<std::shared_ptr<RecordBatch>>> ToRecordBatches(
const StopToken& stop_token) = 0;
using MetadataRecordBatchReader::ReadAll;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToRecordBatches instead.")
Status ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches,
const StopToken& stop_token);
using MetadataRecordBatchReader::ToTable;
/// \brief Consume entire stream as a Table
arrow::Result<std::shared_ptr<Table>> ToTable(const StopToken& stop_token);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToTable instead.")
Status ReadAll(std::shared_ptr<Table>* table, const StopToken& stop_token);
};
// Silence warning
// "non dll-interface class RecordBatchReader used as base for dll-interface class"
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4275)
#endif
/// \brief A RecordBatchWriter that also allows sending
/// application-defined metadata via the Flight protocol.
class ARROW_FLIGHT_EXPORT FlightStreamWriter : public MetadataRecordBatchWriter {
public:
/// \brief Indicate that the application is done writing to this stream.
///
/// The application may not write to this stream after calling
/// this. This differs from closing the stream because this writer
/// may represent only one half of a readable and writable stream.
virtual Status DoneWriting() = 0;
};
#ifdef _MSC_VER
#pragma warning(pop)
#endif
/// \brief A reader for application-specific metadata sent back to the
/// client during an upload.
class ARROW_FLIGHT_EXPORT FlightMetadataReader {
public:
virtual ~FlightMetadataReader();
/// \brief Read a message from the server.
virtual Status ReadMetadata(std::shared_ptr<Buffer>* out) = 0;
};
/// \brief Client class for Arrow Flight RPC services.
/// API experimental for now
class ARROW_FLIGHT_EXPORT FlightClient {
public:
~FlightClient();
/// \brief Connect to an unauthenticated flight service
/// \param[in] location the URI
/// \return Arrow result with the created FlightClient, OK status may not indicate that
/// the connection was successful
static arrow::Result<std::unique_ptr<FlightClient>> Connect(const Location& location);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Connect(const Location& location, std::unique_ptr<FlightClient>* client);
/// \brief Connect to an unauthenticated flight service
/// \param[in] location the URI
/// \param[in] options Other options for setting up the client
/// \return Arrow result with the created FlightClient, OK status may not indicate that
/// the connection was successful
static arrow::Result<std::unique_ptr<FlightClient>> Connect(
const Location& location, const FlightClientOptions& options);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Connect(const Location& location, const FlightClientOptions& options,
std::unique_ptr<FlightClient>* client);
/// \brief Authenticate to the server using the given handler.
/// \param[in] options Per-RPC options
/// \param[in] auth_handler The authentication mechanism to use
/// \return Status OK if the client authenticated successfully
Status Authenticate(const FlightCallOptions& options,
std::unique_ptr<ClientAuthHandler> auth_handler);
/// \brief Authenticate to the server using basic HTTP style authentication.
/// \param[in] options Per-RPC options
/// \param[in] username Username to use
/// \param[in] password Password to use
/// \return Arrow result with bearer token and status OK if client authenticated
/// sucessfully
arrow::Result<std::pair<std::string, std::string>> AuthenticateBasicToken(
const FlightCallOptions& options, const std::string& username,
const std::string& password);
/// \brief Perform the indicated action, returning an iterator to the stream
/// of results, if any
/// \param[in] options Per-RPC options
/// \param[in] action the action to be performed
/// \return Arrow result with an iterator object for reading the returned results
arrow::Result<std::unique_ptr<ResultStream>> DoAction(const FlightCallOptions& options,
const Action& action);
arrow::Result<std::unique_ptr<ResultStream>> DoAction(const Action& action) {
return DoAction({}, action);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoAction(const FlightCallOptions& options, const Action& action,
std::unique_ptr<ResultStream>* results);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoAction(const Action& action, std::unique_ptr<ResultStream>* results) {
return DoAction({}, action).Value(results);
}
/// \brief Retrieve a list of available Action types
/// \param[in] options Per-RPC options
/// \return Arrow result with the available actions
arrow::Result<std::vector<ActionType>> ListActions(const FlightCallOptions& options);
arrow::Result<std::vector<ActionType>> ListActions() {
return ListActions(FlightCallOptions());
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status ListActions(const FlightCallOptions& options, std::vector<ActionType>* actions);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status ListActions(std::vector<ActionType>* actions) {
return ListActions().Value(actions);
}
/// \brief Request access plan for a single flight, which may be an existing
/// dataset or a command to be executed
/// \param[in] options Per-RPC options
/// \param[in] descriptor the dataset request, whether a named dataset or
/// command
/// \return Arrow result with the FlightInfo describing where to access the dataset
arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfo(
const FlightCallOptions& options, const FlightDescriptor& descriptor);
arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfo(
const FlightDescriptor& descriptor) {
return GetFlightInfo({}, descriptor);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetFlightInfo(const FlightCallOptions& options,
const FlightDescriptor& descriptor,
std::unique_ptr<FlightInfo>* info);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetFlightInfo(const FlightDescriptor& descriptor,
std::unique_ptr<FlightInfo>* info) {
return GetFlightInfo({}, descriptor).Value(info);
}
/// \brief Request schema for a single flight, which may be an existing
/// dataset or a command to be executed
/// \param[in] options Per-RPC options
/// \param[in] descriptor the dataset request, whether a named dataset or
/// command
/// \return Arrow result with the SchemaResult describing the dataset schema
arrow::Result<std::unique_ptr<SchemaResult>> GetSchema(
const FlightCallOptions& options, const FlightDescriptor& descriptor);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetSchema(const FlightCallOptions& options, const FlightDescriptor& descriptor,
std::unique_ptr<SchemaResult>* schema_result);
arrow::Result<std::unique_ptr<SchemaResult>> GetSchema(
const FlightDescriptor& descriptor) {
return GetSchema({}, descriptor);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetSchema(const FlightDescriptor& descriptor,
std::unique_ptr<SchemaResult>* schema_result) {
return GetSchema({}, descriptor).Value(schema_result);
}
/// \brief List all available flights known to the server
/// \return Arrow result with an iterator that returns a FlightInfo for each flight
arrow::Result<std::unique_ptr<FlightListing>> ListFlights();
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status ListFlights(std::unique_ptr<FlightListing>* listing);
/// \brief List available flights given indicated filter criteria
/// \param[in] options Per-RPC options
/// \param[in] criteria the filter criteria (opaque)
/// \return Arrow result with an iterator that returns a FlightInfo for each flight
arrow::Result<std::unique_ptr<FlightListing>> ListFlights(
const FlightCallOptions& options, const Criteria& criteria);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status ListFlights(const FlightCallOptions& options, const Criteria& criteria,
std::unique_ptr<FlightListing>* listing);
/// \brief Given a flight ticket and schema, request to be sent the
/// stream. Returns record batch stream reader
/// \param[in] options Per-RPC options
/// \param[in] ticket The flight ticket to use
/// \return Arrow result with the returned RecordBatchReader
arrow::Result<std::unique_ptr<FlightStreamReader>> DoGet(
const FlightCallOptions& options, const Ticket& ticket);
arrow::Result<std::unique_ptr<FlightStreamReader>> DoGet(const Ticket& ticket) {
return DoGet({}, ticket);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoGet(const FlightCallOptions& options, const Ticket& ticket,
std::unique_ptr<FlightStreamReader>* stream);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoGet(const Ticket& ticket, std::unique_ptr<FlightStreamReader>* stream) {
return DoGet({}, ticket).Value(stream);
}
/// \brief DoPut return value
struct DoPutResult {
/// \brief a writer to write record batches to
std::unique_ptr<FlightStreamWriter> writer;
/// \brief a reader for application metadata from the server
std::unique_ptr<FlightMetadataReader> reader;
};
/// \brief Upload data to a Flight described by the given
/// descriptor. The caller must call Close() on the returned stream
/// once they are done writing.
///
/// The reader and writer are linked; closing the writer will also
/// close the reader. Use \a DoneWriting to only close the write
/// side of the channel.
///
/// \param[in] options Per-RPC options
/// \param[in] descriptor the descriptor of the stream
/// \param[in] schema the schema for the data to upload
/// \return Arrow result with a DoPutResult struct holding a reader and a writer
arrow::Result<DoPutResult> DoPut(const FlightCallOptions& options,
const FlightDescriptor& descriptor,
const std::shared_ptr<Schema>& schema);
arrow::Result<DoPutResult> DoPut(const FlightDescriptor& descriptor,
const std::shared_ptr<Schema>& schema) {
return DoPut({}, descriptor, schema);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoPut(const FlightCallOptions& options, const FlightDescriptor& descriptor,
const std::shared_ptr<Schema>& schema,
std::unique_ptr<FlightStreamWriter>* writer,
std::unique_ptr<FlightMetadataReader>* reader);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoPut(const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema,
std::unique_ptr<FlightStreamWriter>* writer,
std::unique_ptr<FlightMetadataReader>* reader) {
ARROW_ASSIGN_OR_RAISE(auto output, DoPut({}, descriptor, schema));
*writer = std::move(output.writer);
*reader = std::move(output.reader);
return Status::OK();
}
struct DoExchangeResult {
std::unique_ptr<FlightStreamWriter> writer;
std::unique_ptr<FlightStreamReader> reader;
};
arrow::Result<DoExchangeResult> DoExchange(const FlightCallOptions& options,
const FlightDescriptor& descriptor);
arrow::Result<DoExchangeResult> DoExchange(const FlightDescriptor& descriptor) {
return DoExchange({}, descriptor);
}
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoExchange(const FlightCallOptions& options, const FlightDescriptor& descriptor,
std::unique_ptr<FlightStreamWriter>* writer,
std::unique_ptr<FlightStreamReader>* reader);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status DoExchange(const FlightDescriptor& descriptor,
std::unique_ptr<FlightStreamWriter>* writer,
std::unique_ptr<FlightStreamReader>* reader) {
ARROW_ASSIGN_OR_RAISE(auto output, DoExchange({}, descriptor));
*writer = std::move(output.writer);
*reader = std::move(output.reader);
return Status::OK();
}
/// \brief Explicitly shut down and clean up the client.
///
/// For backwards compatibility, this will be implicitly called by
/// the destructor if not already called, but this gives the
/// application no chance to handle errors, so it is recommended to
/// explicitly close the client.
///
/// \since 8.0.0
Status Close();
private:
FlightClient();
Status CheckOpen() const;
std::unique_ptr<internal::ClientTransport> transport_;
bool closed_;
int64_t write_size_limit_bytes_;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,62 @@
// 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.
#pragma once
#include <string>
#include "arrow/flight/visibility.h"
#include "arrow/status.h"
namespace arrow {
namespace flight {
/// \brief A reader for messages from the server during an
/// authentication handshake.
class ARROW_FLIGHT_EXPORT ClientAuthReader {
public:
virtual ~ClientAuthReader() = default;
virtual Status Read(std::string* response) = 0;
};
/// \brief A writer for messages to the server during an
/// authentication handshake.
class ARROW_FLIGHT_EXPORT ClientAuthSender {
public:
virtual ~ClientAuthSender() = default;
virtual Status Write(const std::string& token) = 0;
};
/// \brief An authentication implementation for a Flight service.
/// Authentication includes both an initial negotiation and a per-call
/// token validation. Implementations may choose to use either or both
/// mechanisms.
class ARROW_FLIGHT_EXPORT ClientAuthHandler {
public:
virtual ~ClientAuthHandler() = default;
/// \brief Authenticate the client on initial connection. The client
/// can send messages to/read responses from the server at any time.
/// \return Status OK if authenticated successfully
virtual Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader* incoming) = 0;
/// \brief Get a per-call token.
/// \param[out] token The token to send to the server.
virtual Status GetToken(std::string* token) = 0;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,33 @@
// 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.
// Middleware implementation for sending and receiving HTTP cookies.
#pragma once
#include <memory>
#include "arrow/flight/client_middleware.h"
namespace arrow {
namespace flight {
/// \brief Returns a ClientMiddlewareFactory that handles sending and receiving cookies.
ARROW_FLIGHT_EXPORT std::shared_ptr<ClientMiddlewareFactory> GetCookieFactory();
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,73 @@
// 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.
// Interfaces for defining middleware for Flight clients. Currently
// experimental.
#pragma once
#include <memory>
#include "arrow/flight/middleware.h"
#include "arrow/flight/visibility.h" // IWYU pragma: keep
#include "arrow/status.h"
namespace arrow {
namespace flight {
/// \brief Client-side middleware for a call, instantiated per RPC.
///
/// Middleware should be fast and must be infallible: there is no way
/// to reject the call or report errors from the middleware instance.
class ARROW_FLIGHT_EXPORT ClientMiddleware {
public:
virtual ~ClientMiddleware() = default;
/// \brief A callback before headers are sent. Extra headers can be
/// added, but existing ones cannot be read.
virtual void SendingHeaders(AddCallHeaders* outgoing_headers) = 0;
/// \brief A callback when headers are received from the server.
virtual void ReceivedHeaders(const CallHeaders& incoming_headers) = 0;
/// \brief A callback after the call has completed.
virtual void CallCompleted(const Status& status) = 0;
};
/// \brief A factory for new middleware instances.
///
/// If added to a client, this will be called for each RPC (including
/// Handshake) to give the opportunity to intercept the call.
///
/// It is guaranteed that all client middleware methods are called
/// from the same thread that calls the RPC method implementation.
class ARROW_FLIGHT_EXPORT ClientMiddlewareFactory {
public:
virtual ~ClientMiddlewareFactory() = default;
/// \brief A callback for the start of a new call.
///
/// \param info Information about the call.
/// \param[out] middleware The middleware instance for this call. If
/// unset, will not add middleware to this call instance from
/// this factory.
virtual void StartCall(const CallInfo& info,
std::unique_ptr<ClientMiddleware>* middleware) = 0;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,73 @@
// 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.
// Interfaces for defining middleware for Flight clients and
// servers. Currently experimental.
#pragma once
#include <map>
#include <memory>
#include <string>
#include <utility>
#include "arrow/flight/visibility.h" // IWYU pragma: keep
#include "arrow/status.h"
#include "arrow/util/string_view.h"
namespace arrow {
namespace flight {
/// \brief Headers sent from the client or server.
///
/// Header values are ordered.
using CallHeaders = std::multimap<util::string_view, util::string_view>;
/// \brief A write-only wrapper around headers for an RPC call.
class ARROW_FLIGHT_EXPORT AddCallHeaders {
public:
virtual ~AddCallHeaders() = default;
/// \brief Add a header to be sent to the client.
virtual void AddHeader(const std::string& key, const std::string& value) = 0;
};
/// \brief An enumeration of the RPC methods Flight implements.
enum class FlightMethod : char {
Invalid = 0,
Handshake = 1,
ListFlights = 2,
GetFlightInfo = 3,
GetSchema = 4,
DoGet = 5,
DoPut = 6,
DoAction = 7,
ListActions = 8,
DoExchange = 9,
};
/// \brief Information about an instance of a Flight RPC.
struct ARROW_FLIGHT_EXPORT CallInfo {
public:
/// \brief The RPC method of this call.
FlightMethod method;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,26 @@
// 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.
// Often-used headers, for precompiling.
// If updating this header, please make sure you check compilation speed
// before checking in. Adding headers which are not used extremely often
// may incur a slowdown, since it makes the precompiled header heavier to load.
#include "arrow/flight/client.h"
#include "arrow/flight/server.h"
#include "arrow/flight/types.h"
#include "arrow/pch.h"

View File

@@ -0,0 +1,31 @@
// 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.
// Internal header. Platform-specific definitions for Flight.
#pragma once
#ifdef _MSC_VER
// The protobuf documentation says that C4251 warnings when using the
// library are spurious and suppressed when the build the library and
// compiler, but must be also suppressed in downstream projects
#pragma warning(disable : 4251)
#endif // _MSC_VER
#include "arrow/util/config.h" // IWYU pragma: keep

View File

@@ -0,0 +1,306 @@
// 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.
// Interfaces to use for defining Flight RPC servers. API should be considered
// experimental for now
#pragma once
#include <chrono>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "arrow/flight/server_auth.h"
#include "arrow/flight/type_fwd.h"
#include "arrow/flight/types.h" // IWYU pragma: keep
#include "arrow/flight/visibility.h" // IWYU pragma: keep
#include "arrow/ipc/dictionary.h"
#include "arrow/ipc/options.h"
#include "arrow/record_batch.h"
namespace arrow {
class Schema;
class Status;
namespace flight {
/// \brief Interface that produces a sequence of IPC payloads to be sent in
/// FlightData protobuf messages
class ARROW_FLIGHT_EXPORT FlightDataStream {
public:
virtual ~FlightDataStream();
virtual std::shared_ptr<Schema> schema() = 0;
/// \brief Compute FlightPayload containing serialized RecordBatch schema
virtual arrow::Result<FlightPayload> GetSchemaPayload() = 0;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetSchemaPayload(FlightPayload* payload) {
return GetSchemaPayload().Value(payload);
}
// When the stream is completed, the last payload written will have null
// metadata
virtual arrow::Result<FlightPayload> Next() = 0;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status Next(FlightPayload* payload) { return Next().Value(payload); }
};
/// \brief A basic implementation of FlightDataStream that will provide
/// a sequence of FlightData messages to be written to a stream
class ARROW_FLIGHT_EXPORT RecordBatchStream : public FlightDataStream {
public:
/// \param[in] reader produces a sequence of record batches
/// \param[in] options IPC options for writing
explicit RecordBatchStream(
const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options = ipc::IpcWriteOptions::Defaults());
~RecordBatchStream() override;
// inherit deprecated API
using FlightDataStream::GetSchemaPayload;
using FlightDataStream::Next;
std::shared_ptr<Schema> schema() override;
arrow::Result<FlightPayload> GetSchemaPayload() override;
arrow::Result<FlightPayload> Next() override;
private:
class RecordBatchStreamImpl;
std::unique_ptr<RecordBatchStreamImpl> impl_;
};
/// \brief A reader for IPC payloads uploaded by a client. Also allows
/// reading application-defined metadata via the Flight protocol.
class ARROW_FLIGHT_EXPORT FlightMessageReader : public MetadataRecordBatchReader {
public:
/// \brief Get the descriptor for this upload.
virtual const FlightDescriptor& descriptor() const = 0;
};
/// \brief A writer for application-specific metadata sent back to the
/// client during an upload.
class ARROW_FLIGHT_EXPORT FlightMetadataWriter {
public:
virtual ~FlightMetadataWriter();
/// \brief Send a message to the client.
virtual Status WriteMetadata(const Buffer& app_metadata) = 0;
};
/// \brief A writer for IPC payloads to a client. Also allows sending
/// application-defined metadata via the Flight protocol.
///
/// This class offers more control compared to FlightDataStream,
/// including the option to write metadata without data and the
/// ability to interleave reading and writing.
class ARROW_FLIGHT_EXPORT FlightMessageWriter : public MetadataRecordBatchWriter {
public:
virtual ~FlightMessageWriter() = default;
};
/// \brief Call state/contextual data.
class ARROW_FLIGHT_EXPORT ServerCallContext {
public:
virtual ~ServerCallContext() = default;
/// \brief The name of the authenticated peer (may be the empty string)
virtual const std::string& peer_identity() const = 0;
/// \brief The peer address (not validated)
virtual const std::string& peer() const = 0;
/// \brief Look up a middleware by key. Do not maintain a reference
/// to the object beyond the request body.
/// \return The middleware, or nullptr if not found.
virtual ServerMiddleware* GetMiddleware(const std::string& key) const = 0;
/// \brief Check if the current RPC has been cancelled (by the client, by
/// a network error, etc.).
virtual bool is_cancelled() const = 0;
};
class ARROW_FLIGHT_EXPORT FlightServerOptions {
public:
explicit FlightServerOptions(const Location& location_);
~FlightServerOptions();
/// \brief The host & port (or domain socket path) to listen on.
/// Use port 0 to bind to an available port.
Location location;
/// \brief The authentication handler to use.
std::shared_ptr<ServerAuthHandler> auth_handler;
/// \brief A list of TLS certificate+key pairs to use.
std::vector<CertKeyPair> tls_certificates;
/// \brief Enable mTLS and require that the client present a certificate.
bool verify_client;
/// \brief If using mTLS, the PEM-encoded root certificate to use.
std::string root_certificates;
/// \brief A list of server middleware to apply, along with a key to
/// identify them by.
///
/// Middleware are always applied in the order provided. Duplicate
/// keys are an error.
std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
middleware;
/// \brief An optional memory manager to control where to allocate incoming data.
std::shared_ptr<MemoryManager> memory_manager;
/// \brief A Flight implementation-specific callback to customize
/// transport-specific options.
///
/// Not guaranteed to be called. The type of the parameter is
/// specific to the Flight implementation. Users should take care to
/// link to the same transport implementation as Flight to avoid
/// runtime problems. See "Using Arrow C++ in your own project" in
/// the documentation for more details.
std::function<void(void*)> builder_hook;
};
/// \brief Skeleton RPC server implementation which can be used to create
/// custom servers by implementing its abstract methods
class ARROW_FLIGHT_EXPORT FlightServerBase {
public:
FlightServerBase();
virtual ~FlightServerBase();
// Lifecycle methods.
/// \brief Initialize a Flight server listening at the given location.
/// This method must be called before any other method.
/// \param[in] options The configuration for this server.
Status Init(const FlightServerOptions& options);
/// \brief Get the port that the Flight server is listening on.
/// This method must only be called after Init(). Will return a
/// non-positive value if no port exists (e.g. when listening on a
/// domain socket).
int port() const;
/// \brief Get the address that the Flight server is listening on.
/// This method must only be called after Init().
Location location() const;
/// \brief Set the server to stop when receiving any of the given signal
/// numbers.
/// This method must be called before Serve().
Status SetShutdownOnSignals(const std::vector<int> sigs);
/// \brief Start serving.
/// This method blocks until either Shutdown() is called or one of the signals
/// registered in SetShutdownOnSignals() is received.
Status Serve();
/// \brief Query whether Serve() was interrupted by a signal.
/// This method must be called after Serve() has returned.
///
/// \return int the signal number that interrupted Serve(), if any, otherwise 0
int GotSignal() const;
/// \brief Shut down the server. Can be called from signal handler or another
/// thread while Serve() blocks. Optionally a deadline can be set. Once the
/// the deadline expires server will wait until remaining running calls
/// complete.
///
Status Shutdown(const std::chrono::system_clock::time_point* deadline = NULLPTR);
/// \brief Block until server is terminated with Shutdown.
Status Wait();
// Implement these methods to create your own server. The default
// implementations will return a not-implemented result to the client
/// \brief Retrieve a list of available fields given an optional opaque
/// criteria
/// \param[in] context The call context.
/// \param[in] criteria may be null
/// \param[out] listings the returned listings iterator
/// \return Status
virtual Status ListFlights(const ServerCallContext& context, const Criteria* criteria,
std::unique_ptr<FlightListing>* listings);
/// \brief Retrieve the schema and an access plan for the indicated
/// descriptor
/// \param[in] context The call context.
/// \param[in] request may be null
/// \param[out] info the returned flight info provider
/// \return Status
virtual Status GetFlightInfo(const ServerCallContext& context,
const FlightDescriptor& request,
std::unique_ptr<FlightInfo>* info);
/// \brief Retrieve the schema for the indicated descriptor
/// \param[in] context The call context.
/// \param[in] request may be null
/// \param[out] schema the returned flight schema provider
/// \return Status
virtual Status GetSchema(const ServerCallContext& context,
const FlightDescriptor& request,
std::unique_ptr<SchemaResult>* schema);
/// \brief Get a stream of IPC payloads to put on the wire
/// \param[in] context The call context.
/// \param[in] request an opaque ticket
/// \param[out] stream the returned stream provider
/// \return Status
virtual Status DoGet(const ServerCallContext& context, const Ticket& request,
std::unique_ptr<FlightDataStream>* stream);
/// \brief Process a stream of IPC payloads sent from a client
/// \param[in] context The call context.
/// \param[in] reader a sequence of uploaded record batches
/// \param[in] writer send metadata back to the client
/// \return Status
virtual Status DoPut(const ServerCallContext& context,
std::unique_ptr<FlightMessageReader> reader,
std::unique_ptr<FlightMetadataWriter> writer);
/// \brief Process a bidirectional stream of IPC payloads
/// \param[in] context The call context.
/// \param[in] reader a sequence of uploaded record batches
/// \param[in] writer send data back to the client
/// \return Status
virtual Status DoExchange(const ServerCallContext& context,
std::unique_ptr<FlightMessageReader> reader,
std::unique_ptr<FlightMessageWriter> writer);
/// \brief Execute an action, return stream of zero or more results
/// \param[in] context The call context.
/// \param[in] action the action to execute, with type and body
/// \param[out] result the result iterator
/// \return Status
virtual Status DoAction(const ServerCallContext& context, const Action& action,
std::unique_ptr<ResultStream>* result);
/// \brief Retrieve the list of available actions
/// \param[in] context The call context.
/// \param[out] actions a vector of available action types
/// \return Status
virtual Status ListActions(const ServerCallContext& context,
std::vector<ActionType>* actions);
private:
struct Impl;
std::unique_ptr<Impl> impl_;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,78 @@
// 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.
/// \brief Server-side APIs to implement authentication for Flight.
#pragma once
#include <string>
#include "arrow/flight/visibility.h"
#include "arrow/status.h"
namespace arrow {
namespace flight {
/// \brief A reader for messages from the client during an
/// authentication handshake.
class ARROW_FLIGHT_EXPORT ServerAuthReader {
public:
virtual ~ServerAuthReader() = default;
virtual Status Read(std::string* token) = 0;
};
/// \brief A writer for messages to the client during an
/// authentication handshake.
class ARROW_FLIGHT_EXPORT ServerAuthSender {
public:
virtual ~ServerAuthSender() = default;
virtual Status Write(const std::string& message) = 0;
};
/// \brief An authentication implementation for a Flight service.
/// Authentication includes both an initial negotiation and a per-call
/// token validation. Implementations may choose to use either or both
/// mechanisms.
/// An implementation may need to track some state, e.g. a mapping of
/// client tokens to authenticated identities.
class ARROW_FLIGHT_EXPORT ServerAuthHandler {
public:
virtual ~ServerAuthHandler();
/// \brief Authenticate the client on initial connection. The server
/// can send and read responses from the client at any time.
virtual Status Authenticate(ServerAuthSender* outgoing, ServerAuthReader* incoming) = 0;
/// \brief Validate a per-call client token.
/// \param[in] token The client token. May be the empty string if
/// the client does not provide a token.
/// \param[out] peer_identity The identity of the peer, if this
/// authentication method supports it.
/// \return Status OK if the token is valid, any other status if
/// validation failed
virtual Status IsValid(const std::string& token, std::string* peer_identity) = 0;
};
/// \brief An authentication mechanism that does nothing.
class ARROW_FLIGHT_EXPORT NoOpAuthHandler : public ServerAuthHandler {
public:
~NoOpAuthHandler() override;
Status Authenticate(ServerAuthSender* outgoing, ServerAuthReader* incoming) override;
Status IsValid(const std::string& token, std::string* peer_identity) override;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,83 @@
// 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.
// Interfaces for defining middleware for Flight servers. Currently
// experimental.
#pragma once
#include <memory>
#include <string>
#include "arrow/flight/middleware.h"
#include "arrow/flight/visibility.h" // IWYU pragma: keep
#include "arrow/status.h"
namespace arrow {
namespace flight {
/// \brief Server-side middleware for a call, instantiated per RPC.
///
/// Middleware should be fast and must be infallible: there is no way
/// to reject the call or report errors from the middleware instance.
class ARROW_FLIGHT_EXPORT ServerMiddleware {
public:
virtual ~ServerMiddleware() = default;
/// \brief Unique name of middleware, used as alternative to RTTI
/// \return the string name of the middleware
virtual std::string name() const = 0;
/// \brief A callback before headers are sent. Extra headers can be
/// added, but existing ones cannot be read.
virtual void SendingHeaders(AddCallHeaders* outgoing_headers) = 0;
/// \brief A callback after the call has completed.
virtual void CallCompleted(const Status& status) = 0;
};
/// \brief A factory for new middleware instances.
///
/// If added to a server, this will be called for each RPC (including
/// Handshake) to give the opportunity to intercept the call.
///
/// It is guaranteed that all server middleware methods are called
/// from the same thread that calls the RPC method implementation.
class ARROW_FLIGHT_EXPORT ServerMiddlewareFactory {
public:
virtual ~ServerMiddlewareFactory() = default;
/// \brief A callback for the start of a new call.
///
/// Return a non-OK status to reject the call with the given status.
///
/// \param info Information about the call.
/// \param incoming_headers Headers sent by the client for this call.
/// Do not retain a reference to this object.
/// \param[out] middleware The middleware instance for this call. If
/// null, no middleware will be added to this call instance from
/// this factory.
/// \return Status A non-OK status will reject the call with the
/// given status. Middleware previously in the chain will have
/// their CallCompleted callback called. Other middleware
/// factories will not be called.
virtual Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
std::shared_ptr<ServerMiddleware>* middleware) = 0;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,278 @@
// 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.
// Common test definitions for Flight. Individual transport
// implementations can instantiate these tests.
//
// While Googletest's value-parameterized tests would be a more
// natural way to do this, they cause runtime issues on MinGW/MSVC
// (Googletest thinks the test suite has been defined twice).
#pragma once
#include <gtest/gtest.h>
#include <functional>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>
#include "arrow/flight/server.h"
#include "arrow/flight/types.h"
#include "arrow/util/macros.h"
namespace arrow {
namespace flight {
class ARROW_FLIGHT_EXPORT FlightTest : public testing::Test {
protected:
virtual std::string transport() const = 0;
};
/// Common tests of startup/shutdown
class ARROW_FLIGHT_EXPORT ConnectivityTest : public FlightTest {
public:
// Test methods
void TestGetPort();
void TestBuilderHook();
void TestShutdown();
void TestShutdownWithDeadline();
void TestBrokenConnection();
};
#define ARROW_FLIGHT_TEST_CONNECTIVITY(FIXTURE) \
static_assert(std::is_base_of<ConnectivityTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from ConnectivityTest"); \
TEST_F(FIXTURE, GetPort) { TestGetPort(); } \
TEST_F(FIXTURE, BuilderHook) { TestBuilderHook(); } \
TEST_F(FIXTURE, Shutdown) { TestShutdown(); } \
TEST_F(FIXTURE, ShutdownWithDeadline) { TestShutdownWithDeadline(); } \
TEST_F(FIXTURE, BrokenConnection) { TestBrokenConnection(); }
/// Common tests of data plane methods
class ARROW_FLIGHT_EXPORT DataTest : public FlightTest {
public:
void SetUp();
void TearDown();
Status ConnectClient();
// Test methods
void TestDoGetInts();
void TestDoGetFloats();
void TestDoGetDicts();
void TestDoGetLargeBatch();
void TestOverflowServerBatch();
void TestOverflowClientBatch();
void TestDoExchange();
void TestDoExchangeNoData();
void TestDoExchangeWriteOnlySchema();
void TestDoExchangeGet();
void TestDoExchangePut();
void TestDoExchangeEcho();
void TestDoExchangeTotal();
void TestDoExchangeError();
void TestDoExchangeConcurrency();
void TestDoExchangeUndrained();
void TestIssue5095();
private:
void CheckDoGet(
const FlightDescriptor& descr, const RecordBatchVector& expected_batches,
std::function<void(const std::vector<FlightEndpoint>&)> check_endpoints);
void CheckDoGet(const Ticket& ticket, const RecordBatchVector& expected_batches);
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
};
#define ARROW_FLIGHT_TEST_DATA(FIXTURE) \
static_assert(std::is_base_of<DataTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from DataTest"); \
TEST_F(FIXTURE, TestDoGetInts) { TestDoGetInts(); } \
TEST_F(FIXTURE, TestDoGetFloats) { TestDoGetFloats(); } \
TEST_F(FIXTURE, TestDoGetDicts) { TestDoGetDicts(); } \
TEST_F(FIXTURE, TestDoGetLargeBatch) { TestDoGetLargeBatch(); } \
TEST_F(FIXTURE, TestOverflowServerBatch) { TestOverflowServerBatch(); } \
TEST_F(FIXTURE, TestOverflowClientBatch) { TestOverflowClientBatch(); } \
TEST_F(FIXTURE, TestDoExchange) { TestDoExchange(); } \
TEST_F(FIXTURE, TestDoExchangeNoData) { TestDoExchangeNoData(); } \
TEST_F(FIXTURE, TestDoExchangeWriteOnlySchema) { TestDoExchangeWriteOnlySchema(); } \
TEST_F(FIXTURE, TestDoExchangeGet) { TestDoExchangeGet(); } \
TEST_F(FIXTURE, TestDoExchangePut) { TestDoExchangePut(); } \
TEST_F(FIXTURE, TestDoExchangeEcho) { TestDoExchangeEcho(); } \
TEST_F(FIXTURE, TestDoExchangeTotal) { TestDoExchangeTotal(); } \
TEST_F(FIXTURE, TestDoExchangeError) { TestDoExchangeError(); } \
TEST_F(FIXTURE, TestDoExchangeConcurrency) { TestDoExchangeConcurrency(); } \
TEST_F(FIXTURE, TestDoExchangeUndrained) { TestDoExchangeUndrained(); } \
TEST_F(FIXTURE, TestIssue5095) { TestIssue5095(); }
/// \brief Specific tests of DoPut.
class ARROW_FLIGHT_EXPORT DoPutTest : public FlightTest {
public:
void SetUp();
void TearDown();
void CheckBatches(const FlightDescriptor& expected_descriptor,
const RecordBatchVector& expected_batches);
void CheckDoPut(const FlightDescriptor& descr, const std::shared_ptr<Schema>& schema,
const RecordBatchVector& batches);
// Test methods
void TestInts();
void TestFloats();
void TestEmptyBatch();
void TestDicts();
void TestLargeBatch();
void TestSizeLimit();
void TestUndrained();
private:
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
};
#define ARROW_FLIGHT_TEST_DO_PUT(FIXTURE) \
static_assert(std::is_base_of<DoPutTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from DoPutTest"); \
TEST_F(FIXTURE, TestInts) { TestInts(); } \
TEST_F(FIXTURE, TestFloats) { TestFloats(); } \
TEST_F(FIXTURE, TestEmptyBatch) { TestEmptyBatch(); } \
TEST_F(FIXTURE, TestDicts) { TestDicts(); } \
TEST_F(FIXTURE, TestLargeBatch) { TestLargeBatch(); } \
TEST_F(FIXTURE, TestSizeLimit) { TestSizeLimit(); } \
TEST_F(FIXTURE, TestUndrained) { TestUndrained(); }
class ARROW_FLIGHT_EXPORT AppMetadataTestServer : public FlightServerBase {
public:
virtual ~AppMetadataTestServer() = default;
Status DoGet(const ServerCallContext& context, const Ticket& request,
std::unique_ptr<FlightDataStream>* data_stream) override;
Status DoPut(const ServerCallContext& context,
std::unique_ptr<FlightMessageReader> reader,
std::unique_ptr<FlightMetadataWriter> writer) override;
};
/// \brief Tests of app_metadata in data plane methods.
class ARROW_FLIGHT_EXPORT AppMetadataTest : public FlightTest {
public:
void SetUp();
void TearDown();
// Test methods
void TestDoGet();
void TestDoGetDictionaries();
void TestDoPut();
void TestDoPutDictionaries();
void TestDoPutReadMetadata();
private:
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
};
#define ARROW_FLIGHT_TEST_APP_METADATA(FIXTURE) \
static_assert(std::is_base_of<AppMetadataTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from AppMetadataTest"); \
TEST_F(FIXTURE, TestDoGet) { TestDoGet(); } \
TEST_F(FIXTURE, TestDoGetDictionaries) { TestDoGetDictionaries(); } \
TEST_F(FIXTURE, TestDoPut) { TestDoPut(); } \
TEST_F(FIXTURE, TestDoPutDictionaries) { TestDoPutDictionaries(); } \
TEST_F(FIXTURE, TestDoPutReadMetadata) { TestDoPutReadMetadata(); }
/// \brief Tests of IPC options in data plane methods.
class ARROW_FLIGHT_EXPORT IpcOptionsTest : public FlightTest {
public:
void SetUp();
void TearDown();
// Test methods
void TestDoGetReadOptions();
void TestDoPutWriteOptions();
void TestDoExchangeClientWriteOptions();
void TestDoExchangeClientWriteOptionsBegin();
void TestDoExchangeServerWriteOptions();
private:
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
};
#define ARROW_FLIGHT_TEST_IPC_OPTIONS(FIXTURE) \
static_assert(std::is_base_of<IpcOptionsTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from IpcOptionsTest"); \
TEST_F(FIXTURE, TestDoGetReadOptions) { TestDoGetReadOptions(); } \
TEST_F(FIXTURE, TestDoPutWriteOptions) { TestDoPutWriteOptions(); } \
TEST_F(FIXTURE, TestDoExchangeClientWriteOptions) { \
TestDoExchangeClientWriteOptions(); \
} \
TEST_F(FIXTURE, TestDoExchangeClientWriteOptionsBegin) { \
TestDoExchangeClientWriteOptionsBegin(); \
} \
TEST_F(FIXTURE, TestDoExchangeServerWriteOptions) { \
TestDoExchangeServerWriteOptions(); \
}
/// \brief Tests of data plane methods with CUDA memory.
///
/// If not built with ARROW_CUDA, tests are no-ops.
class ARROW_FLIGHT_EXPORT CudaDataTest : public FlightTest {
public:
void SetUp() override;
void TearDown() override;
// Test methods
void TestDoGet();
void TestDoPut();
void TestDoExchange();
private:
class Impl;
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
std::shared_ptr<Impl> impl_;
};
#define ARROW_FLIGHT_TEST_CUDA_DATA(FIXTURE) \
static_assert(std::is_base_of<CudaDataTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from CudaDataTest"); \
TEST_F(FIXTURE, TestDoGet) { TestDoGet(); } \
TEST_F(FIXTURE, TestDoPut) { TestDoPut(); } \
TEST_F(FIXTURE, TestDoExchange) { TestDoExchange(); }
/// \brief Tests of error handling.
class ARROW_FLIGHT_EXPORT ErrorHandlingTest : public FlightTest {
public:
void SetUp() override;
void TearDown() override;
// Test methods
void TestGetFlightInfo();
private:
std::unique_ptr<FlightClient> client_;
std::unique_ptr<FlightServerBase> server_;
};
#define ARROW_FLIGHT_TEST_ERROR_HANDLING(FIXTURE) \
static_assert(std::is_base_of<ErrorHandlingTest, FIXTURE>::value, \
ARROW_STRINGIFY(FIXTURE) " must inherit from ErrorHandlingTest"); \
TEST_F(FIXTURE, TestGetFlightInfo) { TestGetFlightInfo(); }
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,260 @@
// 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.
#pragma once
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <cstdint>
#include <memory>
#include <string>
#include <thread>
#include <utility>
#include <vector>
#include "arrow/status.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/testing/util.h"
#include "arrow/util/make_unique.h"
#include "arrow/flight/client.h"
#include "arrow/flight/client_auth.h"
#include "arrow/flight/server.h"
#include "arrow/flight/server_auth.h"
#include "arrow/flight/types.h"
#include "arrow/flight/visibility.h"
namespace boost {
namespace process {
class child;
} // namespace process
} // namespace boost
namespace arrow {
namespace flight {
// ----------------------------------------------------------------------
// Helpers to compare values for equality
inline void AssertEqual(const FlightInfo& expected, const FlightInfo& actual) {
ipc::DictionaryMemo expected_memo;
ipc::DictionaryMemo actual_memo;
ASSERT_OK_AND_ASSIGN(auto ex_schema, expected.GetSchema(&expected_memo));
ASSERT_OK_AND_ASSIGN(auto actual_schema, actual.GetSchema(&actual_memo));
AssertSchemaEqual(*ex_schema, *actual_schema);
ASSERT_EQ(expected.total_records(), actual.total_records());
ASSERT_EQ(expected.total_bytes(), actual.total_bytes());
ASSERT_EQ(expected.descriptor(), actual.descriptor());
ASSERT_THAT(actual.endpoints(), ::testing::ContainerEq(expected.endpoints()));
}
// ----------------------------------------------------------------------
// Fixture to use for running test servers
class ARROW_FLIGHT_EXPORT TestServer {
public:
explicit TestServer(const std::string& executable_name)
: executable_name_(executable_name), port_(::arrow::GetListenPort()) {}
TestServer(const std::string& executable_name, int port)
: executable_name_(executable_name), port_(port) {}
TestServer(const std::string& executable_name, const std::string& unix_sock)
: executable_name_(executable_name), unix_sock_(unix_sock) {}
void Start(const std::vector<std::string>& extra_args);
void Start() { Start({}); }
int Stop();
bool IsRunning();
int port() const;
const std::string& unix_sock() const;
private:
std::string executable_name_;
int port_;
std::string unix_sock_;
std::shared_ptr<::boost::process::child> server_process_;
};
/// \brief Create a simple Flight server for testing
ARROW_FLIGHT_EXPORT
std::unique_ptr<FlightServerBase> ExampleTestServer();
// Helper to initialize a server and matching client with callbacks to
// populate options.
template <typename T, typename... Args>
Status MakeServer(const Location& location, std::unique_ptr<FlightServerBase>* server,
std::unique_ptr<FlightClient>* client,
std::function<Status(FlightServerOptions*)> make_server_options,
std::function<Status(FlightClientOptions*)> make_client_options,
Args&&... server_args) {
*server = arrow::internal::make_unique<T>(std::forward<Args>(server_args)...);
FlightServerOptions server_options(location);
RETURN_NOT_OK(make_server_options(&server_options));
RETURN_NOT_OK((*server)->Init(server_options));
std::string uri =
location.scheme() + "://127.0.0.1:" + std::to_string((*server)->port());
ARROW_ASSIGN_OR_RAISE(auto real_location, Location::Parse(uri));
FlightClientOptions client_options = FlightClientOptions::Defaults();
RETURN_NOT_OK(make_client_options(&client_options));
return FlightClient::Connect(real_location, client_options).Value(client);
}
// Helper to initialize a server and matching client with callbacks to
// populate options.
template <typename T, typename... Args>
Status MakeServer(std::unique_ptr<FlightServerBase>* server,
std::unique_ptr<FlightClient>* client,
std::function<Status(FlightServerOptions*)> make_server_options,
std::function<Status(FlightClientOptions*)> make_client_options,
Args&&... server_args) {
ARROW_ASSIGN_OR_RAISE(auto location, Location::ForGrpcTcp("localhost", 0));
return MakeServer<T>(location, server, client, std::move(make_server_options),
std::move(make_client_options),
std::forward<Args>(server_args)...);
}
// ----------------------------------------------------------------------
// A FlightDataStream that numbers the record batches
/// \brief A basic implementation of FlightDataStream that will provide
/// a sequence of FlightData messages to be written to a stream
class ARROW_FLIGHT_EXPORT NumberingStream : public FlightDataStream {
public:
explicit NumberingStream(std::unique_ptr<FlightDataStream> stream);
std::shared_ptr<Schema> schema() override;
arrow::Result<FlightPayload> GetSchemaPayload() override;
arrow::Result<FlightPayload> Next() override;
private:
int counter_;
std::shared_ptr<FlightDataStream> stream_;
};
// ----------------------------------------------------------------------
// Example data for test-server and unit tests
ARROW_FLIGHT_EXPORT
std::shared_ptr<Schema> ExampleIntSchema();
ARROW_FLIGHT_EXPORT
std::shared_ptr<Schema> ExampleStringSchema();
ARROW_FLIGHT_EXPORT
std::shared_ptr<Schema> ExampleDictSchema();
ARROW_FLIGHT_EXPORT
std::shared_ptr<Schema> ExampleLargeSchema();
ARROW_FLIGHT_EXPORT
Status ExampleIntBatches(RecordBatchVector* out);
ARROW_FLIGHT_EXPORT
Status ExampleFloatBatches(RecordBatchVector* out);
ARROW_FLIGHT_EXPORT
Status ExampleDictBatches(RecordBatchVector* out);
ARROW_FLIGHT_EXPORT
Status ExampleNestedBatches(RecordBatchVector* out);
ARROW_FLIGHT_EXPORT
Status ExampleLargeBatches(RecordBatchVector* out);
ARROW_FLIGHT_EXPORT
arrow::Result<std::shared_ptr<RecordBatch>> VeryLargeBatch();
ARROW_FLIGHT_EXPORT
std::vector<FlightInfo> ExampleFlightInfo();
ARROW_FLIGHT_EXPORT
std::vector<ActionType> ExampleActionTypes();
ARROW_FLIGHT_EXPORT
Status MakeFlightInfo(const Schema& schema, const FlightDescriptor& descriptor,
const std::vector<FlightEndpoint>& endpoints, int64_t total_records,
int64_t total_bytes, FlightInfo::Data* out);
// ----------------------------------------------------------------------
// A pair of authentication handlers that check for a predefined password
// and set the peer identity to a predefined username.
class ARROW_FLIGHT_EXPORT TestServerAuthHandler : public ServerAuthHandler {
public:
explicit TestServerAuthHandler(const std::string& username,
const std::string& password);
~TestServerAuthHandler() override;
Status Authenticate(ServerAuthSender* outgoing, ServerAuthReader* incoming) override;
Status IsValid(const std::string& token, std::string* peer_identity) override;
private:
std::string username_;
std::string password_;
};
class ARROW_FLIGHT_EXPORT TestServerBasicAuthHandler : public ServerAuthHandler {
public:
explicit TestServerBasicAuthHandler(const std::string& username,
const std::string& password);
~TestServerBasicAuthHandler() override;
Status Authenticate(ServerAuthSender* outgoing, ServerAuthReader* incoming) override;
Status IsValid(const std::string& token, std::string* peer_identity) override;
private:
BasicAuth basic_auth_;
};
class ARROW_FLIGHT_EXPORT TestClientAuthHandler : public ClientAuthHandler {
public:
explicit TestClientAuthHandler(const std::string& username,
const std::string& password);
~TestClientAuthHandler() override;
Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader* incoming) override;
Status GetToken(std::string* token) override;
private:
std::string username_;
std::string password_;
};
class ARROW_FLIGHT_EXPORT TestClientBasicAuthHandler : public ClientAuthHandler {
public:
explicit TestClientBasicAuthHandler(const std::string& username,
const std::string& password);
~TestClientBasicAuthHandler() override;
Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader* incoming) override;
Status GetToken(std::string* token) override;
private:
BasicAuth basic_auth_;
std::string token_;
};
ARROW_FLIGHT_EXPORT
Status ExampleTlsCertificates(std::vector<CertKeyPair>* out);
ARROW_FLIGHT_EXPORT
Status ExampleTlsCertificateRoot(CertKeyPair* out);
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,275 @@
// 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.
/// \file
/// Internal (but not private) interface for implementing
/// alternate network transports in Flight.
///
/// \warning EXPERIMENTAL. Subject to change.
///
/// To implement a transport, implement ServerTransport and
/// ClientTransport, and register the desired URI schemes with
/// TransportRegistry. Flight takes care of most of the per-RPC
/// details; transports only handle connections and providing a I/O
/// stream implementation (TransportDataStream).
///
/// On the server side:
///
/// 1. Applications subclass FlightServerBase and override RPC handlers.
/// 2. FlightServerBase::Init will look up and create a ServerTransport
/// based on the scheme of the Location given to it.
/// 3. The ServerTransport will start the actual server. (For instance,
/// for gRPC, it creates a gRPC server and registers a gRPC service.)
/// That server will handle connections.
/// 4. The transport should forward incoming calls to the server to the RPC
/// handlers defined on ServerTransport, which implements the actual
/// RPC handler using the interfaces here. Any I/O the RPC handler needs
/// to do is managed by transport-specific implementations of
/// TransportDataStream.
/// 5. ServerTransport calls FlightServerBase for the actual application
/// logic.
///
/// On the client side:
///
/// 1. Applications create a FlightClient with a Location.
/// 2. FlightClient will look up and create a ClientTransport based on
/// the scheme of the Location given to it.
/// 3. When calling a method on FlightClient, FlightClient will delegate to
/// the ClientTransport. There is some indirection, e.g. for DoGet,
/// FlightClient only requests that the ClientTransport start the
/// call and provide it with an I/O stream. The "Flight implementation"
/// itself still lives in FlightClient.
#pragma once
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "arrow/flight/type_fwd.h"
#include "arrow/flight/visibility.h"
#include "arrow/type_fwd.h"
#include "arrow/util/optional.h"
namespace arrow {
namespace ipc {
class Message;
}
namespace flight {
class FlightStatusDetail;
namespace internal {
/// Internal, not user-visible type used for memory-efficient reads
struct FlightData {
/// Used only for puts, may be null
std::unique_ptr<FlightDescriptor> descriptor;
/// Non-length-prefixed Message header as described in format/Message.fbs
std::shared_ptr<Buffer> metadata;
/// Application-defined metadata
std::shared_ptr<Buffer> app_metadata;
/// Message body
std::shared_ptr<Buffer> body;
/// Open IPC message from the metadata and body
::arrow::Result<std::unique_ptr<ipc::Message>> OpenMessage();
};
/// \brief A transport-specific interface for reading/writing Arrow data.
///
/// New transports will implement this to read/write IPC payloads to
/// the underlying stream.
class ARROW_FLIGHT_EXPORT TransportDataStream {
public:
virtual ~TransportDataStream() = default;
/// \brief Attempt to read the next FlightData message.
///
/// \return success true if data was populated, false if there was
/// an error. For clients, the error can be retrieved from
/// Finish(Status).
virtual bool ReadData(FlightData* data);
/// \brief Attempt to write a FlightPayload.
///
/// \param[in] payload The data to write.
/// \return true if the message was accepted by the transport, false
/// if not (e.g. due to client/server disconnect), Status if there
/// was an error (e.g. with the payload itself).
virtual arrow::Result<bool> WriteData(const FlightPayload& payload);
/// \brief Indicate that there are no more writes on this stream.
///
/// This is only a hint for the underlying transport and may not
/// actually do anything.
virtual Status WritesDone();
};
/// \brief A transport-specific interface for reading/writing Arrow
/// data for a client.
class ARROW_FLIGHT_EXPORT ClientDataStream : public TransportDataStream {
public:
/// \brief Attempt to read a non-data message.
///
/// Only implemented for DoPut; mutually exclusive with
/// ReadData(FlightData*).
virtual bool ReadPutMetadata(std::shared_ptr<Buffer>* out);
/// \brief Attempt to cancel the call.
///
/// This is only a hint and may not take effect immediately. The
/// client should still finish the call with Finish(Status) as usual.
virtual void TryCancel() {}
/// \brief Finish the call, reporting the server-sent status and/or
/// any client-side errors as appropriate.
///
/// Implies WritesDone() and DoFinish().
///
/// \param[in] st A client-side status to combine with the
/// server-side error. That is, if an error occurs on the
/// client-side, call Finish(Status) to finish the server-side
/// call, get the server-side status, and merge the statuses
/// together so context is not lost.
Status Finish(Status st);
protected:
/// \brief End the call, returning the final server status.
///
/// For implementors: should imply WritesDone() (even if it does not
/// directly call it).
///
/// Implies WritesDone().
virtual Status DoFinish() = 0;
};
/// An implementation of a Flight client for a particular transport.
///
/// Transports should override the methods they are capable of
/// supporting. The default method implementations return an error.
class ARROW_FLIGHT_EXPORT ClientTransport {
public:
virtual ~ClientTransport() = default;
/// Initialize the client.
virtual Status Init(const FlightClientOptions& options, const Location& location,
const arrow::internal::Uri& uri) = 0;
/// Close the client. Once this returns, the client is no longer usable.
virtual Status Close() = 0;
virtual Status Authenticate(const FlightCallOptions& options,
std::unique_ptr<ClientAuthHandler> auth_handler);
virtual arrow::Result<std::pair<std::string, std::string>> AuthenticateBasicToken(
const FlightCallOptions& options, const std::string& username,
const std::string& password);
virtual Status DoAction(const FlightCallOptions& options, const Action& action,
std::unique_ptr<ResultStream>* results);
virtual Status ListActions(const FlightCallOptions& options,
std::vector<ActionType>* actions);
virtual Status GetFlightInfo(const FlightCallOptions& options,
const FlightDescriptor& descriptor,
std::unique_ptr<FlightInfo>* info);
virtual arrow::Result<std::unique_ptr<SchemaResult>> GetSchema(
const FlightCallOptions& options, const FlightDescriptor& descriptor);
virtual Status ListFlights(const FlightCallOptions& options, const Criteria& criteria,
std::unique_ptr<FlightListing>* listing);
virtual Status DoGet(const FlightCallOptions& options, const Ticket& ticket,
std::unique_ptr<ClientDataStream>* stream);
virtual Status DoPut(const FlightCallOptions& options,
std::unique_ptr<ClientDataStream>* stream);
virtual Status DoExchange(const FlightCallOptions& options,
std::unique_ptr<ClientDataStream>* stream);
};
/// A registry of transport implementations.
class ARROW_FLIGHT_EXPORT TransportRegistry {
public:
using ClientFactory = std::function<arrow::Result<std::unique_ptr<ClientTransport>>()>;
using ServerFactory = std::function<arrow::Result<std::unique_ptr<ServerTransport>>(
FlightServerBase*, std::shared_ptr<MemoryManager> memory_manager)>;
TransportRegistry();
~TransportRegistry();
arrow::Result<std::unique_ptr<ClientTransport>> MakeClient(
const std::string& scheme) const;
arrow::Result<std::unique_ptr<ServerTransport>> MakeServer(
const std::string& scheme, FlightServerBase* base,
std::shared_ptr<MemoryManager> memory_manager) const;
Status RegisterClient(const std::string& scheme, ClientFactory factory);
Status RegisterServer(const std::string& scheme, ServerFactory factory);
private:
class Impl;
std::unique_ptr<Impl> impl_;
};
/// \brief Get the registry of transport implementations.
ARROW_FLIGHT_EXPORT
TransportRegistry* GetDefaultTransportRegistry();
//------------------------------------------------------------
// Error propagation helpers
/// \brief Abstract status code as per the Flight specification.
enum class TransportStatusCode {
kOk = 0,
kUnknown = 1,
kInternal = 2,
kInvalidArgument = 3,
kTimedOut = 4,
kNotFound = 5,
kAlreadyExists = 6,
kCancelled = 7,
kUnauthenticated = 8,
kUnauthorized = 9,
kUnimplemented = 10,
kUnavailable = 11,
};
/// \brief Abstract error status.
///
/// Transport implementations may use side channels (e.g. HTTP
/// trailers) to convey additional information to reconstruct the
/// original C++ status for implementations that can use it.
struct ARROW_FLIGHT_EXPORT TransportStatus {
TransportStatusCode code;
std::string message;
/// \brief Convert a C++ status to an abstract transport status.
static TransportStatus FromStatus(const Status& arrow_status);
/// \brief Reconstruct a string-encoded TransportStatus.
static TransportStatus FromCodeStringAndMessage(const std::string& code_str,
std::string message);
/// \brief Convert an abstract transport status to a C++ status.
Status ToStatus() const;
};
/// \brief Convert the string representation of an Arrow status code
/// back to an Arrow status.
ARROW_FLIGHT_EXPORT
Status ReconstructStatus(const std::string& code_str, const Status& current_status,
util::optional<std::string> message,
util::optional<std::string> detail_message,
util::optional<std::string> detail_bin,
std::shared_ptr<FlightStatusDetail> detail);
} // namespace internal
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,133 @@
// 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.
#pragma once
#include <chrono>
#include <memory>
#include "arrow/flight/transport.h"
#include "arrow/flight/type_fwd.h"
#include "arrow/flight/visibility.h"
#include "arrow/type_fwd.h"
namespace arrow {
namespace ipc {
class Message;
}
namespace flight {
namespace internal {
/// \brief A transport-specific interface for reading/writing Arrow
/// data for a server.
class ARROW_FLIGHT_EXPORT ServerDataStream : public TransportDataStream {
public:
/// \brief Attempt to write a non-data message.
///
/// Only implemented for DoPut; mutually exclusive with
/// WriteData(const FlightPayload&).
virtual Status WritePutMetadata(const Buffer& payload);
};
/// \brief An implementation of a Flight server for a particular
/// transport.
///
/// This class (the transport implementation) implements the underlying
/// server and handles connections/incoming RPC calls. It should forward RPC
/// calls to the RPC handlers defined on this class, which work in terms of
/// the generic interfaces above. The RPC handlers here then forward calls
/// to the underlying FlightServerBase instance that contains the actual
/// application RPC method handlers.
///
/// Used by FlightServerBase to manage the server lifecycle.
class ARROW_FLIGHT_EXPORT ServerTransport {
public:
ServerTransport(FlightServerBase* base, std::shared_ptr<MemoryManager> memory_manager)
: base_(base), memory_manager_(std::move(memory_manager)) {}
virtual ~ServerTransport() = default;
/// \name Server Lifecycle Methods
/// Transports implement these methods to start/shutdown the underlying
/// server.
/// @{
/// \brief Initialize the server.
///
/// This method should launch the server in a background thread, i.e. it
/// should not block. Once this returns, the server should be active.
virtual Status Init(const FlightServerOptions& options,
const arrow::internal::Uri& uri) = 0;
/// \brief Shutdown the server.
///
/// This should wait for active RPCs to finish. Once this returns, the
/// server is no longer listening.
virtual Status Shutdown() = 0;
/// \brief Shutdown the server with a deadline.
///
/// This should wait for active RPCs to finish, or for the deadline to
/// expire. Once this returns, the server is no longer listening.
virtual Status Shutdown(const std::chrono::system_clock::time_point& deadline) = 0;
/// \brief Wait for the server to shutdown (but do not shut down the server).
///
/// Once this returns, the server is no longer listening.
virtual Status Wait() = 0;
/// \brief Get the address the server is listening on, else an empty Location.
virtual Location location() const = 0;
///@}
/// \name RPC Handlers
/// Implementations of RPC handlers for Flight methods using the common
/// interfaces here. Transports should call these methods from their
/// server implementation to handle the actual RPC calls.
///@{
/// \brief Get the FlightServerBase.
///
/// Intended as an escape hatch for now since not all methods have been
/// factored into a transport-agnostic interface.
FlightServerBase* base() const { return base_; }
/// \brief Implement DoGet in terms of a transport-level stream.
///
/// \param[in] context The server context.
/// \param[in] request The request payload.
/// \param[in] stream The transport-specific data stream
/// implementation. Must implement WriteData(const
/// FlightPayload&).
Status DoGet(const ServerCallContext& context, const Ticket& request,
ServerDataStream* stream);
/// \brief Implement DoPut in terms of a transport-level stream.
///
/// \param[in] context The server context.
/// \param[in] stream The transport-specific data stream
/// implementation. Must implement ReadData(FlightData*)
/// and WritePutMetadata(const Buffer&).
Status DoPut(const ServerCallContext& context, ServerDataStream* stream);
/// \brief Implement DoExchange in terms of a transport-level stream.
///
/// \param[in] context The server context.
/// \param[in] stream The transport-specific data stream
/// implementation. Must implement ReadData(FlightData*)
/// and WriteData(const FlightPayload&).
Status DoExchange(const ServerCallContext& context, ServerDataStream* stream);
///@}
protected:
FlightServerBase* base_;
std::shared_ptr<MemoryManager> memory_manager_;
};
} // namespace internal
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,59 @@
// 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.
#pragma once
namespace arrow {
namespace internal {
class Uri;
}
namespace flight {
struct Action;
struct ActionType;
struct BasicAuth;
class ClientAuthHandler;
class ClientMiddleware;
class ClientMiddlewareFactory;
struct Criteria;
class FlightCallOptions;
struct FlightClientOptions;
struct FlightDescriptor;
struct FlightEndpoint;
class FlightInfo;
class FlightListing;
class FlightMetadataReader;
class FlightMetadataWriter;
struct FlightPayload;
class FlightServerBase;
class FlightServerOptions;
class FlightStreamReader;
class FlightStreamWriter;
struct Location;
struct Result;
class ResultStream;
struct SchemaResult;
class ServerCallContext;
class ServerMiddleware;
class ServerMiddlewareFactory;
struct Ticket;
namespace internal {
class ClientTransport;
struct FlightData;
class ServerTransport;
} // namespace internal
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,604 @@
// 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.
// Data structure for Flight RPC. API should be considered experimental for now
#pragma once
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "arrow/flight/visibility.h"
#include "arrow/ipc/options.h"
#include "arrow/ipc/writer.h"
#include "arrow/result.h"
#include "arrow/util/string_view.h"
namespace arrow {
class Buffer;
class RecordBatch;
class Schema;
class Status;
class Table;
namespace ipc {
class DictionaryMemo;
} // namespace ipc
namespace internal {
class Uri;
} // namespace internal
namespace flight {
/// \brief A Flight-specific status code.
enum class FlightStatusCode : int8_t {
/// An implementation error has occurred.
Internal,
/// A request timed out.
TimedOut,
/// A request was cancelled.
Cancelled,
/// We are not authenticated to the remote service.
Unauthenticated,
/// We do not have permission to make this request.
Unauthorized,
/// The remote service cannot handle this request at the moment.
Unavailable,
/// A request failed for some other reason
Failed
};
// Silence warning
// "non dll-interface class RecordBatchReader used as base for dll-interface class"
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4275)
#endif
/// \brief Flight-specific error information in a Status.
class ARROW_FLIGHT_EXPORT FlightStatusDetail : public arrow::StatusDetail {
public:
explicit FlightStatusDetail(FlightStatusCode code) : code_{code} {}
explicit FlightStatusDetail(FlightStatusCode code, std::string extra_info)
: code_{code}, extra_info_(std::move(extra_info)) {}
const char* type_id() const override;
std::string ToString() const override;
/// \brief Get the Flight status code.
FlightStatusCode code() const;
/// \brief Get the extra error info
std::string extra_info() const;
/// \brief Get the human-readable name of the status code.
std::string CodeAsString() const;
/// \brief Set the extra error info
void set_extra_info(std::string extra_info);
/// \brief Try to extract a \a FlightStatusDetail from any Arrow
/// status.
///
/// \return a \a FlightStatusDetail if it could be unwrapped, \a
/// nullptr otherwise
static std::shared_ptr<FlightStatusDetail> UnwrapStatus(const arrow::Status& status);
private:
FlightStatusCode code_;
std::string extra_info_;
};
#ifdef _MSC_VER
#pragma warning(pop)
#endif
/// \brief Make an appropriate Arrow status for the given
/// Flight-specific status.
///
/// \param code The Flight status code.
/// \param message The message for the error.
/// \param extra_info Optional extra binary info for the error (eg protobuf)
ARROW_FLIGHT_EXPORT
Status MakeFlightError(FlightStatusCode code, std::string message,
std::string extra_info = {});
/// \brief A TLS certificate plus key.
struct ARROW_FLIGHT_EXPORT CertKeyPair {
/// \brief The certificate in PEM format.
std::string pem_cert;
/// \brief The key in PEM format.
std::string pem_key;
};
/// \brief A type of action that can be performed with the DoAction RPC.
struct ARROW_FLIGHT_EXPORT ActionType {
/// \brief The name of the action.
std::string type;
/// \brief A human-readable description of the action.
std::string description;
bool Equals(const ActionType& other) const;
friend bool operator==(const ActionType& left, const ActionType& right) {
return left.Equals(right);
}
friend bool operator!=(const ActionType& left, const ActionType& right) {
return !(left == right);
}
};
/// \brief Opaque selection criteria for ListFlights RPC
struct ARROW_FLIGHT_EXPORT Criteria {
/// Opaque criteria expression, dependent on server implementation
std::string expression;
};
/// \brief An action to perform with the DoAction RPC
struct ARROW_FLIGHT_EXPORT Action {
/// The action type
std::string type;
/// The action content as a Buffer
std::shared_ptr<Buffer> body;
};
/// \brief Opaque result returned after executing an action
struct ARROW_FLIGHT_EXPORT Result {
std::shared_ptr<Buffer> body;
};
/// \brief message for simple auth
struct ARROW_FLIGHT_EXPORT BasicAuth {
std::string username;
std::string password;
/// \brief Deserialize this message from its wire-format representation.
static arrow::Result<BasicAuth> Deserialize(arrow::util::string_view serialized);
/// \brief Serialize this message to its wire-format representation.
arrow::Result<std::string> SerializeToString() const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Deserialize(const std::string& serialized, BasicAuth* out);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Serialize(const BasicAuth& basic_auth, std::string* out);
};
/// \brief A request to retrieve or generate a dataset
struct ARROW_FLIGHT_EXPORT FlightDescriptor {
enum DescriptorType {
UNKNOWN = 0, /// Unused
PATH = 1, /// Named path identifying a dataset
CMD = 2 /// Opaque command to generate a dataset
};
/// The descriptor type
DescriptorType type;
/// Opaque value used to express a command. Should only be defined when type
/// is CMD
std::string cmd;
/// List of strings identifying a particular dataset. Should only be defined
/// when type is PATH
std::vector<std::string> path;
bool Equals(const FlightDescriptor& other) const;
/// \brief Get a human-readable form of this descriptor.
std::string ToString() const;
/// \brief Get the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
arrow::Result<std::string> SerializeToString() const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status SerializeToString(std::string* out) const;
/// \brief Parse the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
static arrow::Result<FlightDescriptor> Deserialize(arrow::util::string_view serialized);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Deserialize(const std::string& serialized, FlightDescriptor* out);
// Convenience factory functions
static FlightDescriptor Command(const std::string& c) {
return FlightDescriptor{CMD, c, {}};
}
static FlightDescriptor Path(const std::vector<std::string>& p) {
return FlightDescriptor{PATH, "", p};
}
friend bool operator==(const FlightDescriptor& left, const FlightDescriptor& right) {
return left.Equals(right);
}
friend bool operator!=(const FlightDescriptor& left, const FlightDescriptor& right) {
return !(left == right);
}
};
/// \brief Data structure providing an opaque identifier or credential to use
/// when requesting a data stream with the DoGet RPC
struct ARROW_FLIGHT_EXPORT Ticket {
std::string ticket;
bool Equals(const Ticket& other) const;
friend bool operator==(const Ticket& left, const Ticket& right) {
return left.Equals(right);
}
friend bool operator!=(const Ticket& left, const Ticket& right) {
return !(left == right);
}
/// \brief Get the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
arrow::Result<std::string> SerializeToString() const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status SerializeToString(std::string* out) const;
/// \brief Parse the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
static arrow::Result<Ticket> Deserialize(arrow::util::string_view serialized);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Deserialize(const std::string& serialized, Ticket* out);
};
class FlightClient;
class FlightServerBase;
ARROW_FLIGHT_EXPORT
extern const char* kSchemeGrpc;
ARROW_FLIGHT_EXPORT
extern const char* kSchemeGrpcTcp;
ARROW_FLIGHT_EXPORT
extern const char* kSchemeGrpcUnix;
ARROW_FLIGHT_EXPORT
extern const char* kSchemeGrpcTls;
/// \brief A host location (a URI)
struct ARROW_FLIGHT_EXPORT Location {
public:
/// \brief Initialize a blank location.
Location();
/// \brief Initialize a location by parsing a URI string
static arrow::Result<Location> Parse(const std::string& uri_string);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
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
/// \return Arrow result with the resulting location
static arrow::Result<Location> ForGrpcTcp(const std::string& host, const int port);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status ForGrpcTcp(const std::string& host, const int port, Location* location);
/// \brief Initialize a location for a TLS-enabled, gRPC-based Flight
/// service from a host and port
/// \param[in] host The hostname to connect to
/// \param[in] port The port
/// \return Arrow result with the resulting location
static arrow::Result<Location> ForGrpcTls(const std::string& host, const int port);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status ForGrpcTls(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
/// \return Arrow result with the resulting location
static arrow::Result<Location> ForGrpcUnix(const std::string& path);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status ForGrpcUnix(const std::string& path, Location* location);
/// \brief Initialize a location based on a URI scheme
static arrow::Result<Location> ForScheme(const std::string& scheme,
const std::string& host, const int port);
/// \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;
bool Equals(const Location& other) const;
friend bool operator==(const Location& left, const Location& right) {
return left.Equals(right);
}
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
/// redeemed
struct ARROW_FLIGHT_EXPORT FlightEndpoint {
/// Opaque ticket identify; use with DoGet RPC
Ticket ticket;
/// List of locations where ticket can be redeemed. If the list is empty, the
/// ticket can only be redeemed on the current service where the ticket was
/// generated
std::vector<Location> locations;
bool Equals(const FlightEndpoint& other) const;
friend bool operator==(const FlightEndpoint& left, const FlightEndpoint& right) {
return left.Equals(right);
}
friend bool operator!=(const FlightEndpoint& left, const FlightEndpoint& right) {
return !(left == right);
}
};
/// \brief Staging data structure for messages about to be put on the wire
///
/// This structure corresponds to FlightData in the protocol.
struct ARROW_FLIGHT_EXPORT FlightPayload {
std::shared_ptr<Buffer> descriptor;
std::shared_ptr<Buffer> app_metadata;
ipc::IpcPayload ipc_message;
/// \brief Check that the payload can be written to the wire.
Status Validate() const;
};
/// \brief Schema result returned after a schema request RPC
struct ARROW_FLIGHT_EXPORT SchemaResult {
public:
explicit SchemaResult(std::string schema) : raw_schema_(std::move(schema)) {}
/// \brief return schema
/// \param[in,out] dictionary_memo for dictionary bookkeeping, will
/// be modified
/// \return Arrrow result with the reconstructed Schema
arrow::Result<std::shared_ptr<Schema>> GetSchema(
ipc::DictionaryMemo* dictionary_memo) const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetSchema(ipc::DictionaryMemo* dictionary_memo,
std::shared_ptr<Schema>* out) const;
const std::string& serialized_schema() const { return raw_schema_; }
private:
std::string raw_schema_;
};
/// \brief The access coordinates for retireval of a dataset, returned by
/// GetFlightInfo
class ARROW_FLIGHT_EXPORT FlightInfo {
public:
struct Data {
std::string schema;
FlightDescriptor descriptor;
std::vector<FlightEndpoint> endpoints;
int64_t total_records;
int64_t total_bytes;
};
explicit FlightInfo(const Data& data) : data_(data), reconstructed_schema_(false) {}
explicit FlightInfo(Data&& data)
: data_(std::move(data)), reconstructed_schema_(false) {}
/// \brief Factory method to construct a FlightInfo.
static arrow::Result<FlightInfo> Make(const Schema& schema,
const FlightDescriptor& descriptor,
const std::vector<FlightEndpoint>& endpoints,
int64_t total_records, int64_t total_bytes);
/// \brief Deserialize the Arrow schema of the dataset. Populate any
/// dictionary encoded fields into a DictionaryMemo for
/// bookkeeping
/// \param[in,out] dictionary_memo for dictionary bookkeeping, will
/// be modified
/// \return Arrrow result with the reconstructed Schema
arrow::Result<std::shared_ptr<Schema>> GetSchema(
ipc::DictionaryMemo* dictionary_memo) const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status GetSchema(ipc::DictionaryMemo* dictionary_memo,
std::shared_ptr<Schema>* out) const;
const std::string& serialized_schema() const { return data_.schema; }
/// The descriptor associated with this flight, may not be set
const FlightDescriptor& descriptor() const { return data_.descriptor; }
/// A list of endpoints associated with the flight (dataset). To consume the
/// whole flight, all endpoints must be consumed
const std::vector<FlightEndpoint>& endpoints() const { return data_.endpoints; }
/// The total number of records (rows) in the dataset. If unknown, set to -1
int64_t total_records() const { return data_.total_records; }
/// The total number of bytes in the dataset. If unknown, set to -1
int64_t total_bytes() const { return data_.total_bytes; }
/// \brief Get the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
arrow::Result<std::string> SerializeToString() const;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status SerializeToString(std::string* out) const;
/// \brief Parse the wire-format representation of this type.
///
/// Useful when interoperating with non-Flight systems (e.g. REST
/// services) that may want to return Flight types.
static arrow::Result<std::unique_ptr<FlightInfo>> Deserialize(
arrow::util::string_view serialized);
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
static Status Deserialize(const std::string& serialized,
std::unique_ptr<FlightInfo>* out);
private:
Data data_;
mutable std::shared_ptr<Schema> schema_;
mutable bool reconstructed_schema_;
};
/// \brief An iterator to FlightInfo instances returned by ListFlights.
class ARROW_FLIGHT_EXPORT FlightListing {
public:
virtual ~FlightListing() = default;
/// \brief Retrieve the next FlightInfo from the iterator.
/// \return Arrow result with a single FlightInfo. Set to \a nullptr if there
/// are none left.
virtual arrow::Result<std::unique_ptr<FlightInfo>> Next() = 0;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status Next(std::unique_ptr<FlightInfo>* info);
};
/// \brief An iterator to Result instances returned by DoAction.
class ARROW_FLIGHT_EXPORT ResultStream {
public:
virtual ~ResultStream() = default;
/// \brief Retrieve the next Result from the iterator.
/// \return Arrow result with a single Result. Set to \a nullptr if there are none left.
virtual arrow::Result<std::unique_ptr<Result>> Next() = 0;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status Next(std::unique_ptr<Result>* info);
};
/// \brief A holder for a RecordBatch with associated Flight metadata.
struct ARROW_FLIGHT_EXPORT FlightStreamChunk {
public:
std::shared_ptr<RecordBatch> data;
std::shared_ptr<Buffer> app_metadata;
};
/// \brief An interface to read Flight data with metadata.
class ARROW_FLIGHT_EXPORT MetadataRecordBatchReader {
public:
virtual ~MetadataRecordBatchReader() = default;
/// \brief Get the schema for this stream.
virtual arrow::Result<std::shared_ptr<Schema>> GetSchema() = 0;
/// \brief Get the next message from Flight. If the stream is
/// finished, then the members of \a FlightStreamChunk will be
/// nullptr.
virtual arrow::Result<FlightStreamChunk> Next() = 0;
ARROW_DEPRECATED("Deprecated in 8.0.0. Use Result-returning overload instead.")
Status Next(FlightStreamChunk* next);
/// \brief Consume entire stream as a vector of record batches
virtual arrow::Result<std::vector<std::shared_ptr<RecordBatch>>> ToRecordBatches();
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToRecordBatches instead.")
Status ReadAll(std::vector<std::shared_ptr<RecordBatch>>* batches);
/// \brief Consume entire stream as a Table
virtual arrow::Result<std::shared_ptr<Table>> ToTable();
ARROW_DEPRECATED("Deprecated in 8.0.0. Use ToTable instead.")
Status ReadAll(std::shared_ptr<Table>* table);
};
/// \brief Convert a MetadataRecordBatchReader to a regular RecordBatchReader.
ARROW_FLIGHT_EXPORT
arrow::Result<std::shared_ptr<RecordBatchReader>> MakeRecordBatchReader(
std::shared_ptr<MetadataRecordBatchReader> reader);
/// \brief An interface to write IPC payloads with metadata.
class ARROW_FLIGHT_EXPORT MetadataRecordBatchWriter : public ipc::RecordBatchWriter {
public:
virtual ~MetadataRecordBatchWriter() = default;
/// \brief Begin writing data with the given schema. Only used with \a DoExchange.
virtual Status Begin(const std::shared_ptr<Schema>& schema,
const ipc::IpcWriteOptions& options) = 0;
virtual Status Begin(const std::shared_ptr<Schema>& schema);
virtual Status WriteMetadata(std::shared_ptr<Buffer> app_metadata) = 0;
virtual Status WriteWithMetadata(const RecordBatch& batch,
std::shared_ptr<Buffer> app_metadata) = 0;
};
/// \brief A FlightListing implementation based on a vector of
/// FlightInfo objects.
///
/// This can be iterated once, then it is consumed.
class ARROW_FLIGHT_EXPORT SimpleFlightListing : public FlightListing {
public:
explicit SimpleFlightListing(const std::vector<FlightInfo>& flights);
explicit SimpleFlightListing(std::vector<FlightInfo>&& flights);
arrow::Result<std::unique_ptr<FlightInfo>> Next() override;
private:
int position_;
std::vector<FlightInfo> flights_;
};
/// \brief A ResultStream implementation based on a vector of
/// Result objects.
///
/// This can be iterated once, then it is consumed.
class ARROW_FLIGHT_EXPORT SimpleResultStream : public ResultStream {
public:
explicit SimpleResultStream(std::vector<Result>&& results);
arrow::Result<std::unique_ptr<Result>> Next() override;
private:
std::vector<Result> results_;
size_t position_;
};
} // namespace flight
} // namespace arrow

View File

@@ -0,0 +1,48 @@
// 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.
#pragma once
#if defined(_WIN32) || defined(__CYGWIN__)
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4251)
#else
#pragma GCC diagnostic ignored "-Wattributes"
#endif
#ifdef ARROW_FLIGHT_STATIC
#define ARROW_FLIGHT_EXPORT
#elif defined(ARROW_FLIGHT_EXPORTING)
#define ARROW_FLIGHT_EXPORT __declspec(dllexport)
#else
#define ARROW_FLIGHT_EXPORT __declspec(dllimport)
#endif
#define ARROW_FLIGHT_NO_EXPORT
#else // Not Windows
#ifndef ARROW_FLIGHT_EXPORT
#define ARROW_FLIGHT_EXPORT __attribute__((visibility("default")))
#endif
#ifndef ARROW_FLIGHT_NO_EXPORT
#define ARROW_FLIGHT_NO_EXPORT __attribute__((visibility("hidden")))
#endif
#endif // Non-Windows
#if defined(_MSC_VER)
#pragma warning(pop)
#endif