[WIP] Arrow Flight Server

programs/server/Server.cpp

@@ -64,6 +64,9 @@
 #include <Server/MySQLHandlerFactory.h>
 #include <Server/PostgreSQLHandlerFactory.h>
 
+#include <Server/IRoutineServer.h>
+#include <Server/HTTPServer.h>
+#include <Server/TCPServer.h>
 
 #if !defined(ARCADIA_BUILD)
 #   include "config_core.h"
@@ -78,6 +81,10 @@
 #    include <Common/hasLinuxCapability.h>
 #endif
 
+#if USE_ARROW
+#   include <Server/ArrowFlightServer.h>
+#endif
+
 #if USE_SSL
 #    include <Poco/Net/Context.h>
 #    include <Poco/Net/SecureServerSocket.h>
@@ -802,7 +809,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
         http_params->setTimeout(settings.http_receive_timeout);
         http_params->setKeepAliveTimeout(keep_alive_timeout);
 
-        std::vector<std::unique_ptr<Poco::Net::TCPServer>> servers;
+        std::vector<std::unique_ptr<IRoutineServer>> servers;
 
         std::vector<std::string> listen_hosts = DB::getMultipleValuesFromConfig(config(), "", "listen_host");
 
@@ -907,8 +914,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port);
                 socket.setReceiveTimeout(settings.http_receive_timeout);
                 socket.setSendTimeout(settings.http_send_timeout);
-
-                servers.emplace_back(std::make_unique<Poco::Net::HTTPServer>(
+                servers.emplace_back(std::make_unique<DB::HTTPServer>(
                     createHandlerFactory(*this, async_metrics, "HTTPHandler-factory"), server_pool, socket, http_params));
 
                 LOG_INFO(log, "Listening for http://{}", address.toString());
@@ -922,7 +928,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port, /* secure = */ true);
                 socket.setReceiveTimeout(settings.http_receive_timeout);
                 socket.setSendTimeout(settings.http_send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::HTTPServer>(
+                servers.emplace_back(std::make_unique<DB::HTTPServer>(
                     createHandlerFactory(*this, async_metrics, "HTTPSHandler-factory"), server_pool, socket, http_params));
 
                 LOG_INFO(log, "Listening for https://{}", address.toString());
@@ -940,7 +946,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port);
                 socket.setReceiveTimeout(settings.receive_timeout);
                 socket.setSendTimeout(settings.send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::TCPServer>(
+                servers.emplace_back(std::make_unique<DB::TCPServer>(
                     new TCPHandlerFactory(*this),
                     server_pool,
                     socket,
@@ -957,7 +963,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port, /* secure = */ true);
                 socket.setReceiveTimeout(settings.receive_timeout);
                 socket.setSendTimeout(settings.send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::TCPServer>(
+                servers.emplace_back(std::make_unique<DB::TCPServer>(
                     new TCPHandlerFactory(*this, /* secure= */ true),
                     server_pool,
                     socket,
@@ -977,7 +983,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port);
                 socket.setReceiveTimeout(settings.http_receive_timeout);
                 socket.setSendTimeout(settings.http_send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::HTTPServer>(
+                servers.emplace_back(std::make_unique<DB::HTTPServer>(
                     createHandlerFactory(*this, async_metrics, "InterserverIOHTTPHandler-factory"), server_pool, socket, http_params));
 
                 LOG_INFO(log, "Listening for replica communication (interserver): http://{}", address.toString());
@@ -990,7 +996,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port, /* secure = */ true);
                 socket.setReceiveTimeout(settings.http_receive_timeout);
                 socket.setSendTimeout(settings.http_send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::HTTPServer>(
+                servers.emplace_back(std::make_unique<DB::HTTPServer>(
                     createHandlerFactory(*this, async_metrics, "InterserverIOHTTPSHandler-factory"), server_pool, socket, http_params));
 
                 LOG_INFO(log, "Listening for secure replica communication (interserver): https://{}", address.toString());
@@ -1007,7 +1013,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port, /* secure = */ true);
                 socket.setReceiveTimeout(Poco::Timespan());
                 socket.setSendTimeout(settings.send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::TCPServer>(
+                servers.emplace_back(std::make_unique<DB::TCPServer>(
                     new MySQLHandlerFactory(*this),
                     server_pool,
                     socket,
@@ -1022,7 +1028,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 auto address = socket_bind_listen(socket, listen_host, port, /* secure = */ true);
                 socket.setReceiveTimeout(Poco::Timespan());
                 socket.setSendTimeout(settings.send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::TCPServer>(
+                servers.emplace_back(std::make_unique<DB::TCPServer>(
                     new PostgreSQLHandlerFactory(*this),
                     server_pool,
                     socket,
@@ -1031,14 +1037,28 @@ int Server::main(const std::vector<std::string> & /*args*/)
                 LOG_INFO(log, "Listening for PostgreSQL compatibility protocol: " + address.toString());
             });
 
+            create_server("arrow_flight_port", [&](UInt16 port)
+            {
+#if USE_ARROW
+                auto arrow_flight_server = std::make_unique<DB::ArrowFlightServer>(*this, listen_host, port);
+                std::string location = arrow_flight_server->getLocation();
+                servers.emplace_back(std::move(arrow_flight_server));
+                LOG_INFO(log, "Listening for Arrow Flight compatibility protocol: ", location);
+#else
+                UNUSED(port);
+                throw Exception{"Arrow Flight compatibility protocol is disabled because ClickHouse was built without Arrow support.",
+                                ErrorCodes::SUPPORT_IS_DISABLED};
+#endif
+            });
+
             /// Prometheus (if defined and not setup yet with http_port)
             create_server("prometheus.port", [&](UInt16 port)
             {
                 Poco::Net::ServerSocket socket;
                 auto address = socket_bind_listen(socket, listen_host, port);
                 socket.setReceiveTimeout(settings.http_receive_timeout);
                 socket.setSendTimeout(settings.http_send_timeout);
-                servers.emplace_back(std::make_unique<Poco::Net::HTTPServer>(
+                servers.emplace_back(std::make_unique<DB::HTTPServer>(
                     createHandlerFactory(*this, async_metrics, "PrometheusHandler-factory"), server_pool, socket, http_params));
 
                 LOG_INFO(log, "Listening for Prometheus: http://{}", address.toString());

src/CMakeLists.txt

@@ -332,6 +332,11 @@ if (USE_CAPNP)
     dbms_target_link_libraries (PRIVATE ${CAPNP_LIBRARIES})
 endif ()
 
+if (USE_ARROW)
+    dbms_target_include_directories (SYSTEM BEFORE PRIVATE ${ARROW_INCLUDE_DIR})
+    dbms_target_link_libraries(PRIVATE ${ARROW_LIBRARY})
+endif ()
+
 if (USE_PARQUET)
     dbms_target_link_libraries(PRIVATE ${PARQUET_LIBRARY})
     if (NOT USE_INTERNAL_PARQUET_LIBRARY OR USE_INTERNAL_PARQUET_LIBRARY_NATIVE_CMAKE)

src/Common/CurrentMetrics.cpp

@@ -25,6 +25,7 @@
     M(HTTPConnection, "Number of connections to HTTP server") \
     M(InterserverConnection, "Number of connections from other replicas to fetch parts") \
     M(PostgreSQLConnection, "Number of client connections using PostgreSQL protocol") \
+    M(ArrowFlightConnection, "Number of client connections using Arrow Flight protocol") \
     M(OpenFileForRead, "Number of files open for reading") \
     M(OpenFileForWrite, "Number of files open for writing") \
     M(Read, "Number of read (read, pread, io_getevents, etc.) syscalls in fly") \

src/Core/config_core.h.in

@@ -11,3 +11,4 @@
 #cmakedefine01 USE_SSL
 #cmakedefine01 USE_OPENCL
 #cmakedefine01 USE_LDAP
+#cmakedefine01 USE_ARROW

src/Server/ArrowFlightServer.cpp

@@ -0,0 +1,153 @@
+#include "ArrowFlightServer.h"
+#include <arrow/flight/test_util.h> // FIXME: Remove it before merge
+
+#include <memory>
+
+namespace
+{
+
+arrow::Status GetBatchForFlight(const arrow::flight::Ticket & ticket, std::shared_ptr<arrow::RecordBatchReader> * out) {
+    if (ticket.ticket == "ticket-ints-1") {
+        arrow::flight::BatchVector batches;
+        RETURN_NOT_OK(arrow::flight::ExampleIntBatches(&batches));
+        *out = std::make_shared<arrow::flight::BatchIterator>(batches[0]->schema(), batches);
+        return arrow::Status::OK();
+    } else if (ticket.ticket == "ticket-dicts-1") {
+        arrow::flight::BatchVector batches;
+        RETURN_NOT_OK(arrow::flight::ExampleDictBatches(&batches));
+        *out = std::make_shared<arrow::flight::BatchIterator>(batches[0]->schema(), batches);
+        return arrow::Status::OK();
+    } else {
+        return arrow::Status::NotImplemented("no stream implemented for this ticket");
+    }
+}
+
+} // FIXME: Remove it before merge
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+extern const int UNKNOWN_EXCEPTION;
+}
+
+ArrowFlightServer::ArrowFlightServer(IServer & server_, std::string host, int port)
+    : server(server_)
+{
+    auto parse_location_status = arrow::flight::Location::ForGrpcTcp(host, port, &location);
+    if (!parse_location_status.ok())
+        throw Exception(ErrorCodes::UNKNOWN_EXCEPTION,
+                        "Invalid location {}:{} for Arrow Flight Server: {}",
+                        host,
+                        port,
+                        parse_location_status.ToString());
+}
+
+std::string ArrowFlightServer::getLocation() const {
+    return location.ToString();
+}
+
+void ArrowFlightServer::start() {
+    arrow::flight::FlightServerOptions options(location);
+    auto init_status = Init(options);
+    if (!init_status.ok())
+        throw Exception(ErrorCodes::UNKNOWN_EXCEPTION,
+                        "Failed init Arrow Flight Server: {}",
+                        init_status.ToString());
+
+    // Exit with a clean error code (0) on SIGTERM
+    // ARROW_CHECK_OK(handler->SetShutdownOnSignals({SIGTERM}));
+
+    auto serve_status = Serve();
+    if (!serve_status.ok())
+        throw Exception(ErrorCodes::UNKNOWN_EXCEPTION,
+                        "Failed serve Arrow Flight: {}",
+                        serve_status.ToString());
+}
+
+void ArrowFlightServer::stop() {
+    auto status = Shutdown();
+    if (!status.ok())
+        throw Exception(ErrorCodes::UNKNOWN_EXCEPTION,
+                        "Failed shutdown Arrow Flight: {}",
+                        status.ToString());
+}
+
+int ArrowFlightServer::currentConnections() const {
+    return 0; // FIXME: implement
+}
+
+arrow::Status ArrowFlightServer::ListFlights(
+    const arrow::flight::ServerCallContext & /* context */,
+    const arrow::flight::Criteria * criteria,
+    std::unique_ptr<arrow::flight::FlightListing> * listings)
+{
+    std::vector<arrow::flight::FlightInfo> flights = arrow::flight::ExampleFlightInfo();
+    if (criteria && !criteria->expression.empty()) {
+        // For test purposes, if we get criteria, return no results
+        flights.clear();
+    }
+    *listings = std::unique_ptr<arrow::flight::FlightListing>(new arrow::flight::SimpleFlightListing(flights));
+    return arrow::Status::OK();
+}
+
+arrow::Status ArrowFlightServer::GetFlightInfo(
+    const arrow::flight::ServerCallContext & /* context */,
+    const arrow::flight::FlightDescriptor & request,
+    std::unique_ptr<arrow::flight::FlightInfo> * out_info)
+{
+    // Test that Arrow-C++ status codes can make it through gRPC
+    if (request.type == arrow::flight::FlightDescriptor::DescriptorType::CMD &&
+        request.cmd == "status-outofmemory") {
+        return arrow::Status::OutOfMemory("Sentinel");
+    }
+
+    std::vector<arrow::flight::FlightInfo> flights = arrow::flight::ExampleFlightInfo();
+
+    for (const auto& info : flights) {
+        if (info.descriptor().Equals(request)) {
+            *out_info = std::make_unique<arrow::flight::FlightInfo>(info);
+            return arrow::Status::OK();
+        }
+    }
+    return arrow::Status::Invalid("Flight not found: ", request.ToString());
+}
+
+arrow::Status ArrowFlightServer::GetSchema(
+    const arrow::flight::ServerCallContext & /* context */,
+    const arrow::flight::FlightDescriptor & request,
+    std::unique_ptr<arrow::flight::SchemaResult> * schema)
+{
+    std::vector<arrow::flight::FlightInfo> flights = arrow::flight::ExampleFlightInfo();
+
+    for (const auto& info : flights) {
+        if (info.descriptor().Equals(request)) {
+            *schema = std::make_unique<arrow::flight::SchemaResult>(info.serialized_schema());
+            return arrow::Status::OK();
+        }
+    }
+    return arrow::Status::Invalid("Flight not found: ", request.ToString());
+}
+
+arrow::Status ArrowFlightServer::DoGet(
+    const arrow::flight::ServerCallContext & /* context */,
+    const arrow::flight::Ticket & request,
+    std::unique_ptr<arrow::flight::FlightDataStream> * data_stream)
+{
+    // Test for ARROW-5095
+    if (request.ticket == "ARROW-5095-fail") {
+        return arrow::Status::UnknownError("Server-side error");
+    }
+    if (request.ticket == "ARROW-5095-success") {
+        return arrow::Status::OK();
+    }
+
+    std::shared_ptr<arrow::RecordBatchReader> batch_reader;
+    RETURN_NOT_OK(GetBatchForFlight(request, &batch_reader));
+
+    *data_stream = std::unique_ptr<arrow::flight::FlightDataStream>(new arrow::flight::RecordBatchStream(batch_reader));
+    return arrow::Status::OK();
+}
+
+}

src/Server/ArrowFlightServer.h

@@ -0,0 +1,74 @@
+#pragma once
+
+#include <Common/CurrentMetrics.h>
+#include <Poco/Net/TCPServerConnection.h>
+#include <common/logger_useful.h>
+#include "IServer.h"
+#include "IRoutineServer.h"
+
+#include <arrow/flight/server.h>
+
+namespace CurrentMetrics
+{
+extern const Metric ArrowFlightConnection;
+}
+
+namespace DB
+{
+
+class ArrowFlightServer : public IRoutineServer, public arrow::flight::FlightServerBase
+{
+public:
+    ArrowFlightServer(IServer & server_, std::string host, int port);
+
+    std::string getLocation() const;
+
+    void start() override;
+
+    void stop() override;
+
+    int currentConnections() const override;
+
+    arrow::Status ListFlights(
+        const arrow::flight::ServerCallContext & context,
+        const arrow::flight::Criteria * criteria,
+        std::unique_ptr<arrow::flight::FlightListing> * listings
+    ) override;
+
+    arrow::Status GetFlightInfo(
+        const arrow::flight::ServerCallContext & context,
+        const arrow::flight::FlightDescriptor & request,
+        std::unique_ptr<arrow::flight::FlightInfo> * info
+    ) override;
+
+    arrow::Status GetSchema(
+        const arrow::flight::ServerCallContext & context,
+        const arrow::flight::FlightDescriptor & request,
+        std::unique_ptr<arrow::flight::SchemaResult> * schema
+    ) override;
+
+    arrow::Status DoGet(
+        const arrow::flight::ServerCallContext & context,
+        const arrow::flight::Ticket & request,
+        std::unique_ptr<arrow::flight::FlightDataStream> * data_stream
+    ) override;
+
+//    arrow::Status DoPut(
+//        const arrow::flight::ServerCallContext & context,
+//        std::unique_ptr<arrow::flight::FlightMessageReader> reader,
+//        std::unique_ptr<arrow::flight::FlightMetadataWriter> writer
+//    ) override;
+
+private:
+    Poco::Logger * log = &Poco::Logger::get("ArrowFlightServer");
+
+    IServer & server;
+    arrow::flight::Location location;
+
+    std::shared_ptr<ReadBuffer> in;
+    std::shared_ptr<WriteBuffer> out;
+
+    CurrentMetrics::Increment metric_increment{CurrentMetrics::ArrowFlightConnection};
+};
+
+}