Add BeginInsert/InsertData/EndInsert flow

Add a new pattern for "prepared inserts". It works like this:

*   Call `BeginInsert` with an `INSERT` query with optional columns
    and ending in `VALUES`. No values should be included in the string.
*   It returns a `Block` pre-configured with columns as
    declared in the `INSERT` statement
*   Add data to the block and periodically call `InsertData` to insert
    data and clear the block.
*   Call `EndInsert()` or just let the `Client` object go out of scope
    to signal the server that it's done inserting.

This allows one to send smaller batches of blocks, thereby using less
memory, but still in a single ClickHouse `INSERT` operation.

Expected to be useful in the Postgres foreign data wrapper insert API,
where multiple rows can be inserted at once but its API handles
one-at-a-time insertion. It will also support the FDW COPY API, which
can submit huge batches of data to insert, as well.

Author: theory

clickhouse/block.h

@@ -85,10 +85,10 @@ class Block {
         return columns_.at(idx).name;
     }
 
-    /// Convinience method to wipe out all rows from all columns
+    /// Convenience method to wipe out all rows from all columns
     void Clear();
 
-    /// Convinience method to do Reserve() on all columns
+    /// Convenience method to do Reserve() on all columns
     void Reserve(size_t new_cap);
 
     /// Reference to column by index in the block.

clickhouse/client.cpp

@@ -161,6 +161,12 @@ class Client::Impl {
 
     void Insert(const std::string& table_name, const std::string& query_id, const Block& block);
 
+    std::unique_ptr<Block> BeginInsert(Query query);
+
+    void InsertData(Block& block);
+
+    void EndInsert();
+
     void Ping();
 
     void ResetConnection();
@@ -175,6 +181,7 @@ class Client::Impl {
     bool Handshake();
 
     bool ReceivePacket(uint64_t* server_packet = nullptr);
+    bool ReceivePreparePackets(uint64_t* server_packet = nullptr);
 
     void SendQuery(const Query& query, bool finalize = true);
     void FinalizeQuery();
@@ -208,6 +215,7 @@ class Client::Impl {
     }
 
 private:
+    bool inserting;
     /// In case of network errors tries to reconnect to server and
     /// call fuc several times.
     void RetryGuard(std::function<void()> func);
@@ -280,10 +288,13 @@ Client::Impl::Impl(const ClientOptions& opts,
     }
 }
 
-Client::Impl::~Impl()
-{ }
+Client::Impl::~Impl() {
+    // Wrap up an insert if one is in progress.
+    EndInsert();
+}
 
 void Client::Impl::ExecuteQuery(Query query) {
+    assert(!inserting);
     EnsureNull en(static_cast<QueryEvents*>(&query), &events_);
 
     if (options_.ping_before_query) {
@@ -299,6 +310,7 @@ void Client::Impl::ExecuteQuery(Query query) {
 
 
 void Client::Impl::SelectWithExternalData(Query query, const ExternalTables& external_tables) {
+    assert(!inserting);
     if (server_info_.revision < DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES) {
        throw UnimplementedError("This version of ClickHouse server doesn't support temporary tables");
     }
@@ -362,6 +374,7 @@ std::string NameToQueryString(const std::string &input)
 }
 
 void Client::Impl::Insert(const std::string& table_name, const std::string& query_id, const Block& block) {
+    assert(!inserting);
     if (options_.ping_before_query) {
         RetryGuard([this]() { Ping(); });
     }
@@ -397,10 +410,24 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer
     }
 
     // Send data.
+    inserting = true;
     SendData(block);
-    // Send empty block as marker of
-    // end of data.
+    EndInsert();
+}
+
+void Client::Impl::InsertData(Block& block) {
+    assert(inserting);
+    block.RefreshRowCount();
+    SendData(block);
+    block.Clear();
+}
+
+void Client::Impl::EndInsert() {
+    if (!inserting) return;
+
+    // Send empty block as marker of end of data.
     SendData(Block());
+    inserting = false;
 
     // Wait for EOS.
     uint64_t eos_packet{0};
@@ -416,6 +443,7 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer
 }
 
 void Client::Impl::Ping() {
+    assert(!inserting);
     WireFormat::WriteUInt64(*output_, ClientCodes::Ping);
     output_->Flush();
 
@@ -429,6 +457,7 @@ void Client::Impl::Ping() {
 
 void Client::Impl::ResetConnection() {
     InitializeStreams(socket_factory_->connect(options_, current_endpoint_.value()));
+    inserting = false;
 
     if (!Handshake()) {
         throw ProtocolError("fail to connect to " + options_.host);
@@ -648,6 +677,78 @@ bool Client::Impl::ReceivePacket(uint64_t* server_packet) {
     }
 }
 
+bool Client::Impl::ReceivePreparePackets(uint64_t* server_packet) {
+    uint64_t packet_type = 0;
+
+    while (true) {
+        if (!WireFormat::ReadVarint64(*input_, &packet_type)) {
+            throw std::runtime_error("unexpected package type " +
+                    std::to_string((int)packet_type) + " for insert query");
+        }
+        if (server_packet) {
+            *server_packet = packet_type;
+        }
+
+        switch (packet_type) {
+        case ServerCodes::Data: {
+            if (!ReceiveData()) {
+                throw ProtocolError("can't read data packet from input stream");
+            }
+            return true;
+        }
+
+        case ServerCodes::Exception: {
+            ReceiveException();
+            return false;
+        }
+
+        case ServerCodes::ProfileInfo:
+        case ServerCodes::Progress:
+        case ServerCodes::Pong:
+        case ServerCodes::Hello:
+            continue;
+
+        case ServerCodes::Log: {
+            // log tag
+            if (!WireFormat::SkipString(*input_)) {
+                return false;
+            }
+            Block block;
+
+            // Use uncompressed stream since log blocks usually contain only one row
+            if (!ReadBlock(*input_, &block)) {
+                return false;
+            }
+
+            if (events_) {
+                events_->OnServerLog(block);
+            }
+            continue;
+        }
+
+        case ServerCodes::TableColumns: {
+            // external table name
+            if (!WireFormat::SkipString(*input_)) {
+                return false;
+            }
+
+            //  columns metadata
+            if (!WireFormat::SkipString(*input_)) {
+                return false;
+            }
+            continue;
+        }
+
+        // No others expected.
+        case ServerCodes::EndOfStream:
+        case ServerCodes::ProfileEvents:
+        default:
+            throw UnimplementedError("unimplemented " + std::to_string((int)packet_type));
+            break;
+        }
+    }
+}
+
 bool Client::Impl::ReadBlock(InputStream& input, Block* block) {
     // Additional information about block.
     if (server_info_.revision >= DBMS_MIN_REVISION_WITH_BLOCK_INFO) {
@@ -923,7 +1024,6 @@ void Client::Impl::FinalizeQuery() {
     output_->Flush();
 }
 
-
 void Client::Impl::WriteBlock(const Block& block, OutputStream& output) {
     // Additional information about block.
     if (server_info_.revision >= DBMS_MIN_REVISION_WITH_BLOCK_INFO) {
@@ -1063,7 +1163,7 @@ void Client::Impl::RetryGuard(std::function<void()> func) {
             }
         }
     }
-    // Connectiong with current_endpoint_ are broken.
+    // Connections with current_endpoint_ are broken.
     // Trying to establish  with the another one from the list.
     size_t connection_attempts_count = GetConnectionAttempts();
     for (size_t i = 0; i < connection_attempts_count;)
@@ -1085,6 +1185,43 @@ void Client::Impl::RetryGuard(std::function<void()> func) {
     }
 }
 
+std::unique_ptr<Block> Client::Impl::BeginInsert(Query query) {
+    assert(!inserting);
+    // Arrange a query callback to extract a block that corresponds to the
+    // query columns.
+    auto block = std::make_unique<Block>();
+    query.OnData([&block](const Block& b) {
+        for (Block::Iterator bi(b); bi.IsValid(); bi.Next()) {
+            // Create the ClickHouse column type.
+            clickhouse::ColumnRef col = bi.Column();
+            auto chtype = col->Type();
+            if (chtype->GetCode() == Type::LowCardinality) {
+                chtype = col->As<ColumnLowCardinality>()->GetNestedType();
+            }
+            block->AppendColumn(bi.Name(), clickhouse::CreateColumnByType(chtype->GetName()));
+        }
+
+        return true;
+    });
+
+
+    EnsureNull en(static_cast<QueryEvents*>(&query), &events_);
+
+    if (options_.ping_before_query) {
+        RetryGuard([this]() { Ping(); });
+    }
+
+    SendQuery(query.GetText());
+
+    // Receive data packet but keep the query/connection open.
+    if (!ReceivePreparePackets()) {
+       throw std::runtime_error("fail to receive data packet");
+    }
+
+    inserting = true;
+    return block;
+}
+
 Client::Client(const ClientOptions& opts)
     : options_(opts)
     , impl_(new Impl(opts))
@@ -1149,6 +1286,27 @@ void Client::Insert(const std::string& table_name, const std::string& query_id,
     impl_->Insert(table_name, query_id, block);
 }
 
+std::unique_ptr<Block> Client::BeginInsert(const std::string& query) {
+    return impl_->BeginInsert(Query(query));
+}
+
+std::unique_ptr<Block> Client::BeginInsert(const std::string& query, const std::string& query_id) {
+    return impl_->BeginInsert(Query(query, query_id));
+}
+
+void Client::InsertData(Block& block) {
+    impl_->InsertData(block);
+}
+
+void Client::EndInsert(Block& block) {
+    impl_->InsertData(block);
+    impl_->EndInsert();
+}
+
+void Client::EndInsert() {
+    impl_->EndInsert();
+}
+
 void Client::Ping() {
     impl_->Ping();
 }

clickhouse/client.h

@@ -273,6 +273,45 @@ class Client {
     void Insert(const std::string& table_name, const Block& block);
     void Insert(const std::string& table_name, const std::string& query_id, const Block& block);
 
+    /**
+     * @brief Begin and continuously send data for an \p INSERT query
+     * @param query The \p INSERT query must end with \p VALUES but contain no values
+     * @return The Block to use to send data to insert
+     *
+     * Use this API to insert multiple blocks of data into a table. Example:
+     *
+     * \code{.cpp}
+     * // Start the insertion.
+     * auto block = client->BeginInsert("INSERT INTO foo (id, name) VALUES");
+     * 
+     * // Grab the columns from the inserter's block.
+     * auto block = inserter->GetBlock();
+     * auto col1 = (*block)[0]->As<ColumnUInt64>();
+     * auto col2 = (*block)[1]->As<ColumnString>();
+     *
+     * // Add a couple of records to the block.
+     * col1.Append(1);
+     * col1.Append(2);
+     * col2.Append("holden");
+     * col2.Append("naomi");
+     *
+     * Send those records.
+     * client->InsertData(*block);
+     *
+     * // Add another record.
+     * col1.Append(3);
+     * col2.Append("amos");
+     *
+     * // Send it and finish.
+     * client->EndInsert(*block);
+     * \endcode
+     */
+    std::unique_ptr<Block> BeginInsert(const std::string& query);
+    std::unique_ptr<Block> BeginInsert(const std::string& query, const std::string& query_id);
+    void InsertData(Block& block);
+    void EndInsert();
+    void EndInsert(Block& block);
+
     /// Ping server for aliveness.
     void Ping();