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.
*   When the `PreparedInsert` object goes out of scope it first signals
    the server that it's done sending data.

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);
 
@@ -415,6 +420,23 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer
     }
 }
 
+void Client::Impl::FinishInsert() {
+    // Send empty block as marker of end of data.
+    SendData(Block());
+
+    // Wait for EOS.
+    uint64_t eos_packet{0};
+    while (ReceivePacket(&eos_packet)) {
+        ;
+    }
+
+    if (eos_packet != ServerCodes::EndOfStream && eos_packet != ServerCodes::Exception
+        && eos_packet != ServerCodes::Log && options_.rethrow_exceptions) {
+        throw ProtocolError(std::string{"unexpected packet from server while receiving end of query, expected (expected Exception, EndOfStream or Log, got: "}
+                            + (eos_packet ? std::to_string(eos_packet) : "nothing") + ")");
+    }
+}
+
 void Client::Impl::Ping() {
     WireFormat::WriteUInt64(*output_, ClientCodes::Ping);
     output_->Flush();
@@ -648,6 +670,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) {
@@ -1063,7 +1157,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 +1179,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(col->Type()->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 +1278,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 +1316,26 @@ Client::Version Client::GetVersion() {
     };
 }
 
+Client::PreparedInsert::PreparedInsert(void *c, Block *b) {
+    client = c;
+    block = b;
+}
+
+Client::PreparedInsert::~PreparedInsert() {
+    auto c = (Client::Impl *)(client);
+    c->FinishInsert();
+
+    // Do not delete client as we're effectively its child.
+    if (block) delete block;
+}
+
+Block * Client::PreparedInsert::GetBlock() { return block; }
+
+void Client::PreparedInsert::Execute() {
+    auto c = (Client::Impl *)(client);
+    block->RefreshRowCount();
+    c->SendData(*block);
+    block->Clear();
+}
+
 }

clickhouse/client.h

@@ -243,6 +243,21 @@ class Client {
             std::unique_ptr<SocketFactory> socket_factory);
     ~Client();
 
+    /**
+     * Prepared insert object.
+     */
+    class PreparedInsert {
+    public:
+        Block * GetBlock();
+        void Execute();
+        // 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 +288,43 @@ 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->Execute();
+     * delete inserter; // Finishes insert.
+     * \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();