Add PreparedInsert flow

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

*   Call `PrepareInsert` with an `INSERT` query with optional columns
    and ending in `VALUES`. No values should be included in the string.
*   It returns a `PreparedInsert` object that has two methods:
    *   `Block()` returns a `Block` pre-configured with columns as
        declared in the `INSERT` statement
    *   `Execute()` inserts data from the block then clears it.
*   Call `Finish()` or just let the `PreparedInsert` object go out of
    scope to send any remaining rows and to signal the server that it's
    done.

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);
 
+    PreparedInsert * PrepareInsert(Query query);
+
+    void FinishInsert();
+
+    void SendData(const Block& block);
+
     void Ping();
 
     void ResetConnection();
@@ -175,12 +181,11 @@ 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();
 
-    void SendData(const Block& block);
-
     void SendBlockData(const Block& block);
     void SendExternalData(const ExternalTables& external_tables);
 
@@ -398,8 +403,11 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer
 
     // Send data.
     SendData(block);
-    // Send empty block as marker of
-    // end of data.
+    FinishInsert();
+}
+
+void Client::Impl::FinishInsert() {
+    // Send empty block as marker of end of data.
     SendData(Block());
 
     // Wait for EOS.
@@ -648,6 +656,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 +1003,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 +1142,7 @@ void Client::Impl::RetryGuard(std::function<void()> func) {
             }
         }
     }
-    // Connectiong with current_endpoint_ are broken.
+    // Connecting 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 +1164,41 @@ void Client::Impl::RetryGuard(std::function<void()> func) {
     }
 }
 
+Client::PreparedInsert * Client::Impl::PrepareInsert(Query query) {
+    // Arrange a query callback to extract a block that corresponds to the
+    // query columns.
+    auto block = new 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");
+    }
+
+    return new PreparedInsert(this, block);
+}
+
 Client::Client(const ClientOptions& opts)
     : options_(opts)
     , impl_(new Impl(opts))
@@ -1149,6 +1263,14 @@ void Client::Insert(const std::string& table_name, const std::string& query_id,
     impl_->Insert(table_name, query_id, block);
 }
 
+Client::PreparedInsert * Client::PrepareInsert(const std::string& query) {
+    return impl_->PrepareInsert(Query(query));
+}
+
+Client::PreparedInsert * Client::PrepareInsert(const std::string& query, const std::string& query_id) {
+    return impl_->PrepareInsert(Query(query, query_id));
+}
+
 void Client::Ping() {
     impl_->Ping();
 }
@@ -1179,4 +1301,42 @@ Client::Version Client::GetVersion() {
     };
 }
 
+Client::PreparedInsert::PreparedInsert(void *c, Block *b) {
+    client = c;
+    block = b;
+}
+
+Client::PreparedInsert::~PreparedInsert() {
+    Finish();
+}
+
+Block * Client::PreparedInsert::GetBlock() { return block; }
+
+void Client::PreparedInsert::Execute() {
+    if (!block) {
+         throw std::runtime_error("Cannot execute finished PrepareExecute");
+    }
+    auto c = (Client::Impl *)(client);
+    block->RefreshRowCount();
+    c->SendData(*block);
+    block->Clear();
+}
+
+void Client::PreparedInsert::Finish() {
+    if (!block) return;
+
+    auto c = (Client::Impl *)(client);
+    block->RefreshRowCount();
+    if (block->GetRowCount() > 0) {
+        c->SendData(*block);
+        block->Clear();
+    }
+    c->FinishInsert();
+    // Do not delete client as we're effectively its child.
+    if (block) {
+        delete block;
+        block = nullptr;
+    }
+}
+
 }

clickhouse/client.h

@@ -243,6 +243,28 @@ class Client {
             std::unique_ptr<SocketFactory> socket_factory);
     ~Client();
 
+    /**
+     * Prepared insert object.
+     */
+    class PreparedInsert {
+    public:
+        /// Fetch the Block prepared for the INSERT. Returns a null pointer
+        /// after \b Finish() has been called.
+        Block * GetBlock();
+        /// Insert the data in the block and clear it. Raises an exception if
+        /// \b Finish() has been called.
+        void Execute();
+        /// Insert any remaining data in the block, complete the query, and
+        /// deletes the block. Subsequent calls are no-ops.
+        void Finish();
+        // XXX This shouldn't be public.
+        PreparedInsert(void *client, Block * block);
+        ~PreparedInsert();
+    private:
+        Block * block;
+        void * client;
+    };
+
     /// Intends for execute arbitrary queries.
     void Execute(const Query& query);
 
@@ -273,6 +295,42 @@ 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 Prepare an \p INSERT query.
+     * @param query The INSERT query; must end with \p VALUES but contain no values
+     * @return The prepared insert object
+     *
+     * Use this object to insert multiple blocks of data into a table. Example:
+     *
+     * \code{.cpp}
+     * // Start the insertion.
+     * auto inserter = client->PrepareInsert("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.
+     * inserter->Execute();
+     *
+     * // Add another record.
+     * col1.Append(3);
+     * col2.Append("amos");
+     *
+     * // Send it and finish.
+     * inserter->Finish();
+     * \endcode
+     */
+    Client::PreparedInsert * PrepareInsert(const std::string& query);
+    Client::PreparedInsert * PrepareInsert(const std::string& query, const std::string& query_id);
+
     /// Ping server for aliveness.
     void Ping();