#pragma once

#include <chrono>

#include <memory>

#include <string>

#include <utility>

#include <variant>

#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/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 {

typedef std::chrono::duration<double, std::chrono::seconds::period> TimeoutDuration;

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;

};

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, std::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();

};

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);

};

#ifdef _MSC_VER

#pragma warning(push)

#pragma warning(disable : 4275)

#endif

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

class ARROW_FLIGHT_EXPORT FlightMetadataReader {

public:

virtual ~FlightMetadataReader();

/// \brief Read a message from the server.

virtual Status ReadMetadata(std::shared_ptr<Buffer>* out) = 0;

};

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 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,

}

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,

}

/// \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