Fix bug with insert, simplify exchanges logic

src/Storages/RabbitMQ/RabbitMQBlockOutputStream.cpp

@@ -56,6 +56,9 @@ void RabbitMQBlockOutputStream::write(const Block & block)
 void RabbitMQBlockOutputStream::writeSuffix()
 {
     child->writeSuffix();
+
+    if (buffer)
+        buffer->finilizeProducer();
 }
 
 }

src/Storages/RabbitMQ/ReadBufferFromRabbitMQConsumer.cpp

@@ -14,20 +14,9 @@
 namespace DB
 {
 
-namespace ErrorCodes
-{
-    extern const int BAD_ARGUMENTS;
-}
-
 namespace ExchangeType
 {
-    /// Note that default here means default by implementation and not by rabbitmq settings
-    static const String DEFAULT = "default";
-    static const String FANOUT = "fanout";
-    static const String DIRECT = "direct";
-    static const String TOPIC = "topic";
-    static const String HASH = "consistent_hash";
-    static const String HEADERS = "headers";
+    static const String HASH_SUF = "_hash";
 }
 
 static const auto QUEUE_SIZE = 50000; /// Equals capacity of a single rabbitmq queue
@@ -36,34 +25,31 @@ ReadBufferFromRabbitMQConsumer::ReadBufferFromRabbitMQConsumer(
         ChannelPtr consumer_channel_,
         HandlerPtr event_handler_,
         const String & exchange_name_,
+        const AMQP::ExchangeType & exchange_type_,
         const Names & routing_keys_,
         size_t channel_id_,
         Poco::Logger * log_,
         char row_delimiter_,
-        bool bind_by_id_,
+        bool hash_exchange_,
         size_t num_queues_,
-        const String & exchange_type_,
         const String & local_exchange_,
         const std::atomic<bool> & stopped_)
         : ReadBuffer(nullptr, 0)
         , consumer_channel(std::move(consumer_channel_))
         , event_handler(event_handler_)
         , exchange_name(exchange_name_)
+        , exchange_type(exchange_type_)
         , routing_keys(routing_keys_)
         , channel_id(channel_id_)
-        , bind_by_id(bind_by_id_)
+        , hash_exchange(hash_exchange_)
         , num_queues(num_queues_)
-        , exchange_type(exchange_type_)
         , local_exchange(local_exchange_)
-        , local_default_exchange(local_exchange + "_" + ExchangeType::DIRECT)
-        , local_hash_exchange(local_exchange + "_" + ExchangeType::HASH)
+        , local_hash_exchange(local_exchange + ExchangeType::HASH_SUF)
         , log(log_)
         , row_delimiter(row_delimiter_)
         , stopped(stopped_)
         , messages(QUEUE_SIZE * num_queues)
 {
-    exchange_type_set = exchange_type != ExchangeType::DEFAULT;
-
     /* One queue per consumer can handle up to 50000 messages. More queues per consumer can be added.
      * By default there is one queue per consumer.
      */
@@ -86,67 +72,24 @@ ReadBufferFromRabbitMQConsumer::~ReadBufferFromRabbitMQConsumer()
 
 void ReadBufferFromRabbitMQConsumer::initExchange()
 {
-    /* This direct-exchange is used for default implemenation and for INSERT query (so it is always declared). If exchange_type
-     * is not set, then there are only two exchanges - external, defined by the client, and local, unique for each table (default).
-     * This strict division to external and local exchanges is needed to avoid too much complexity with defining exchange_name
-     * for INSERT query producer and, in general, it is better to distinguish them into separate ones.
-     */
-    consumer_channel->declareExchange(local_default_exchange, AMQP::direct).onError([&](const char * message)
-    {
-        local_exchange_declared = false;
-        LOG_ERROR(log, "Failed to declare local direct-exchange. Reason: {}", message);
-    });
-
-    if (!exchange_type_set)
-    {
-        consumer_channel->declareExchange(exchange_name, AMQP::fanout).onError([&](const char * message)
-        {
-            local_exchange_declared = false;
-            LOG_ERROR(log, "Failed to declare default fanout-exchange. Reason: {}", message);
-        });
-
-        /// With fanout exchange the binding key is ignored - a parameter might be arbitrary. All distribution lies on local_exchange.
-        consumer_channel->bindExchange(exchange_name, local_default_exchange, routing_keys[0]).onError([&](const char * message)
-        {
-            local_exchange_declared = false;
-            LOG_ERROR(log, "Failed to bind local direct-exchange to fanout-exchange. Reason: {}", message);
-        });
-
-        return;
-    }
-
-    AMQP::ExchangeType type;
-    if      (exchange_type == ExchangeType::FANOUT)         type = AMQP::ExchangeType::fanout;
-    else if (exchange_type == ExchangeType::DIRECT)         type = AMQP::ExchangeType::direct;
-    else if (exchange_type == ExchangeType::TOPIC)          type = AMQP::ExchangeType::topic;
-    else if (exchange_type == ExchangeType::HASH)           type = AMQP::ExchangeType::consistent_hash;
-    else if (exchange_type == ExchangeType::HEADERS)        type = AMQP::ExchangeType::headers;
-    else throw Exception("Invalid exchange type", ErrorCodes::BAD_ARGUMENTS);
-
     /* Declare client's exchange of the specified type and bind it to hash-exchange (if it is not already hash-exchange), which
-     * will evenly distribute messages between all consumers. (This enables better scaling as without hash-exchange - the only
-     * option to avoid getting the same messages more than once - is having only one consumer with one queue)
+     * will evenly distribute messages between all consumers.
      */
-    consumer_channel->declareExchange(exchange_name, type).onError([&](const char * message)
+    consumer_channel->declareExchange(exchange_name, exchange_type).onError([&](const char * message)
     {
         local_exchange_declared = false;
         LOG_ERROR(log, "Failed to declare client's {} exchange. Reason: {}", exchange_type, message);
     });
 
     /// No need for declaring hash-exchange if there is only one consumer with one queue or exchange type is already hash
-    if (!bind_by_id)
+    if (!hash_exchange || exchange_type == AMQP::ExchangeType::consistent_hash)
         return;
 
-    hash_exchange = true;
-
-    if (exchange_type == ExchangeType::HASH)
-        return;
-
-    /* By default hash exchange distributes messages based on a hash value of a routing key, which must be a string integer. But
-     * in current case we use hash exchange for binding to another exchange of some other type, which needs its own routing keys
-     * of other types: headers, patterns and string-keys. This means that hash property must be changed.
-     */
     {
+        /* By default hash exchange distributes messages based on a hash value of a routing key, which must be a string integer. But
+         * in current case we use hash exchange for binding to another exchange of some other type, which needs its own routing keys
+         * of other types: headers, patterns and string-keys. This means that hash property must be changed.
+         */
         AMQP::Table binding_arguments;
         binding_arguments["hash-property"] = "message_id";
 
@@ -161,7 +104,7 @@ void ReadBufferFromRabbitMQConsumer::initExchange()
 
     /// Then bind client's exchange to sharding exchange (by keys, specified by the client):
 
-    if (exchange_type == ExchangeType::HEADERS)
+    if (exchange_type == AMQP::ExchangeType::headers)
     {
         AMQP::Table binding_arguments;
         std::vector<String> matching;
@@ -181,6 +124,14 @@ void ReadBufferFromRabbitMQConsumer::initExchange()
             LOG_ERROR(log, "Failed to bind local hash exchange to client's exchange. Reason: {}", message);
         });
     }
+    else if (exchange_type == AMQP::ExchangeType::fanout)
+    {
+        consumer_channel->bindExchange(exchange_name, local_hash_exchange, routing_keys[0]).onError([&](const char * message)
+        {
+            local_exchange_declared = false;
+            LOG_ERROR(log, "Failed to bind local hash exchange to client's exchange. Reason: {}", message);
+        });
+    }
     else
     {
         for (const auto & routing_key : routing_keys)
@@ -198,30 +149,31 @@ void ReadBufferFromRabbitMQConsumer::initExchange()
 void ReadBufferFromRabbitMQConsumer::initQueueBindings(const size_t queue_id)
 {
     /// These variables might be updated later from a separate thread in onError callbacks.
-    if (!local_exchange_declared || (exchange_type_set && !local_hash_exchange_declared))
+    if (!local_exchange_declared || (hash_exchange && !local_hash_exchange_declared))
     {
         initExchange();
         local_exchange_declared = true;
         local_hash_exchange_declared = true;
     }
 
-    bool default_bindings_created = false, default_bindings_error = false;
     bool bindings_created = false, bindings_error = false;
 
     consumer_channel->declareQueue(AMQP::exclusive)
     .onSuccess([&](const std::string &  queue_name_, int /* msgcount */, int /* consumercount */)
     {
         queues.emplace_back(queue_name_);
+        LOG_DEBUG(log, "Queue " + queue_name_ + " is declared");
+
         subscribed_queue[queue_name_] = false;
-
-        String binding_key = routing_keys[0];
-
-        /* Every consumer has at least one unique queue. Bind the queues to exchange based on the consumer_channel_id
-         * in case there is one queue per consumer and bind by queue_id in case there is more than 1 queue per consumer.
-         * (queue_id is based on channel_id)
+        /* Subscription can probably be moved back to readPrefix(), but not sure whether it is better in regard to speed, because
+         * if moved there, it must(!) be wrapped inside a channel->onSuccess callback or any other, otherwise
+         * consumer might fail to subscribe and no resubscription will help.
          */
-        if (bind_by_id || hash_exchange)
+        subscribe(queues.back());
+
+        if (hash_exchange)
         {
+            String binding_key;
             if (queues.size() == 1)
             {
                 binding_key = std::to_string(channel_id);
@@ -230,39 +182,67 @@ void ReadBufferFromRabbitMQConsumer::initQueueBindings(const size_t queue_id)
             {
                 binding_key = std::to_string(channel_id + queue_id);
             }
+            /* If exchange_type == hash, then bind directly to this client's exchange (because there is no need for a distributor
+             * exchange as it is already hash-exchange), otherwise hash-exchange is a local distributor exchange.
+             */
+            String current_hash_exchange = exchange_type == AMQP::ExchangeType::consistent_hash ? exchange_name : local_hash_exchange;
+
+            /// If hash-exchange is used for messages distribution, then the binding key is ignored - can be arbitrary.
+            consumer_channel->bindQueue(current_hash_exchange, queue_name_, binding_key)
+            .onSuccess([&]
+            {
+                bindings_created = true;
+            })
+            .onError([&](const char * message)
+            {
+                bindings_error = true;
+                LOG_ERROR(log, "Failed to create queue binding. Reason: {}", message);
+            });
         }
-
-        /// Bind queue to exchange that is used for INSERT query and also for default implementation.
-        consumer_channel->bindQueue(local_default_exchange, queue_name_, binding_key)
-        .onSuccess([&]
+        else if (exchange_type == AMQP::ExchangeType::fanout)
         {
-            default_bindings_created = true;
-        })
-        .onError([&](const char * message)
-        {
-            default_bindings_error = true;
-            LOG_ERROR(log, "Failed to bind to key {}. Reason: {}", binding_key, message);
-        });
-
-        /* Subscription can probably be moved back to readPrefix(), but not sure whether it is better in regard to speed, because
-         * if moved there, it must(!) be wrapped inside a channel->onSuccess callback or any other, otherwise
-         * consumer might fail to subscribe and no resubscription will help.
-         */
-        subscribe(queues.back());
-
-        LOG_DEBUG(log, "Queue " + queue_name_ + " is declared");
-
-        if (exchange_type_set)
-        {
-            if (hash_exchange)
+            consumer_channel->bindQueue(exchange_name, queue_name_, routing_keys[0])
+            .onSuccess([&]
             {
-                /* If exchange_type == hash, then bind directly to this client's exchange (because there is no need for a distributor
-                 * exchange as it is already hash-exchange), otherwise hash-exchange is a local distributor exchange.
-                 */
-                String current_hash_exchange = exchange_type == ExchangeType::HASH ? exchange_name : local_hash_exchange;
+                bindings_created = true;
+            })
+            .onError([&](const char * message)
+            {
+                bindings_error = true;
+                LOG_ERROR(log, "Failed to bind to key. Reason: {}", message);
+            });
+        }
+        else if (exchange_type == AMQP::ExchangeType::headers)
+        {
+            AMQP::Table binding_arguments;
+            std::vector<String> matching;
 
-                /// If hash-exchange is used for messages distribution, then the binding key is ignored - can be arbitrary.
-                consumer_channel->bindQueue(current_hash_exchange, queue_name_, binding_key)
+            /// It is not parsed for the second time - if it was parsed above, then it would never end up here.
+            for (const auto & header : routing_keys)
+            {
+                boost::split(matching, header, [](char c){ return c == '='; });
+                binding_arguments[matching[0]] = matching[1];
+                matching.clear();
+            }
+
+            consumer_channel->bindQueue(exchange_name, queue_name_, routing_keys[0], binding_arguments)
+            .onSuccess([&]
+            {
+                bindings_created = true;
+            })
+            .onError([&](const char * message)
+            {
+                bindings_error = true;
+                LOG_ERROR(log, "Failed to bind queue. Reason: {}", message);
+            });
+        }
+        else
+        {
+            /// Means there is only one queue with one consumer - no even distribution needed - no hash-exchange.
+            for (const auto & routing_key : routing_keys)
+            {
+                /// Binding directly to exchange, specified by the client.
+                consumer_channel->bindQueue(exchange_name, queue_name_, routing_key)
                 .onSuccess([&]
                 {
                     bindings_created = true;
@@ -270,56 +250,14 @@ void ReadBufferFromRabbitMQConsumer::initQueueBindings(const size_t queue_id)
                 .onError([&](const char * message)
                 {
                     bindings_error = true;
-                    LOG_ERROR(log, "Failed to create queue binding to key {}. Reason: {}", binding_key, message);
+                    LOG_ERROR(log, "Failed to bind queue. Reason: {}", message);
                 });
             }
-            else if (exchange_type == ExchangeType::HEADERS)
-            {
-                AMQP::Table binding_arguments;
-                std::vector<String> matching;
-
-                /// It is not parsed for the second time - if it was parsed above, then it would never end up here.
-                for (const auto & header : routing_keys)
-                {
-                    boost::split(matching, header, [](char c){ return c == '='; });
-                    binding_arguments[matching[0]] = matching[1];
-                    matching.clear();
-                }
-
-                consumer_channel->bindQueue(exchange_name, queue_name_, routing_keys[0], binding_arguments)
-                .onSuccess([&]
-                {
-                    bindings_created = true;
-                })
-                .onError([&](const char * message)
-                {
-                    bindings_error = true;
-                    LOG_ERROR(log, "Failed to bind queue to key. Reason: {}", message);
-                });
-            }
-            else
-            {
-                /// Means there is only one queue with one consumer - no even distribution needed - no hash-exchange.
-                for (const auto & routing_key : routing_keys)
-                {
-                    /// Binding directly to exchange, specified by the client.
-                    consumer_channel->bindQueue(exchange_name, queue_name_, routing_key)
-                    .onSuccess([&]
-                    {
-                        bindings_created = true;
-                    })
-                    .onError([&](const char * message)
-                    {
-                        bindings_error = true;
-                        LOG_ERROR(log, "Failed to bind queue to key. Reason: {}", message);
-                    });
-                }
-            }
         }
     })
     .onError([&](const char * message)
     {
-        default_bindings_error = true;
+        bindings_error = true;
         LOG_ERROR(log, "Failed to declare queue on the channel. Reason: {}", message);
     });
 
@@ -327,7 +265,7 @@ void ReadBufferFromRabbitMQConsumer::initQueueBindings(const size_t queue_id)
      * It is important at this moment to make sure that queue bindings are created before any publishing can happen because
      * otherwise messages will be routed nowhere.
      */
-    while ((!default_bindings_created && !default_bindings_error) || (exchange_type_set && !bindings_created && !bindings_error))
+    while (!bindings_created && !bindings_error)
     {
         iterateEventLoop();
     }

src/Storages/RabbitMQ/ReadBufferFromRabbitMQConsumer.h

@@ -26,13 +26,13 @@ public:
             ChannelPtr consumer_channel_,
             HandlerPtr event_handler_,
             const String & exchange_name_,
+            const AMQP::ExchangeType & exchange_type_,
             const Names & routing_keys_,
             size_t channel_id_,
             Poco::Logger * log_,
             char row_delimiter_,
-            bool bind_by_id_,
+            bool hash_exchange_,
             size_t num_queues_,
-            const String & exchange_type_,
             const String & local_exchange_,
             const std::atomic<bool> & stopped_);
 
@@ -48,12 +48,12 @@ private:
     HandlerPtr event_handler;
 
     const String exchange_name;
+    const AMQP::ExchangeType exchange_type;
     const Names routing_keys;
     const size_t channel_id;
-    const bool bind_by_id;
+    const bool hash_exchange;
     const size_t num_queues;
 
-    const String exchange_type;
     const String local_exchange;
     const String local_default_exchange;
     const String local_hash_exchange;
@@ -65,7 +65,6 @@ private:
 
     String default_local_exchange;
     bool local_exchange_declared = false, local_hash_exchange_declared = false;
-    bool exchange_type_set = false, hash_exchange = false;
 
     std::atomic<bool> consumer_error = false;
     std::atomic<size_t> count_subscribed = 0, wait_subscribed;

src/Storages/RabbitMQ/StorageRabbitMQ.cpp

@@ -49,6 +49,16 @@ namespace ErrorCodes
     extern const int CANNOT_CONNECT_RABBITMQ;
 }
 
+namespace ExchangeType
+{
+    /// Note that default here means default by implementation and not by rabbitmq settings
+    static const String DEFAULT = "default";
+    static const String FANOUT = "fanout";
+    static const String DIRECT = "direct";
+    static const String TOPIC = "topic";
+    static const String HASH = "consistent_hash";
+    static const String HEADERS = "headers";
+}
 
 StorageRabbitMQ::StorageRabbitMQ(
         const StorageID & table_id_,
@@ -72,7 +82,6 @@ StorageRabbitMQ::StorageRabbitMQ(
         , row_delimiter(row_delimiter_)
         , num_consumers(num_consumers_)
         , num_queues(num_queues_)
-        , exchange_type(exchange_type_)
         , use_transactional_channel(use_transactional_channel_)
         , log(&Poco::Logger::get("StorageRabbitMQ (" + table_id_.table_name + ")"))
         , parsed_address(parseAddress(global_context.getMacros()->expand(host_port_), 5672))
@@ -107,7 +116,22 @@ StorageRabbitMQ::StorageRabbitMQ(
     heartbeat_task = global_context.getSchedulePool().createTask("RabbitMQHeartbeatTask", [this]{ heartbeatFunc(); });
     heartbeat_task->deactivate();
 
-    bind_by_id = num_consumers > 1 || num_queues > 1;
+    hash_exchange = num_consumers > 1 || num_queues > 1;
+
+    exchange_type_set = exchange_type_ != ExchangeType::DEFAULT;
+    if (exchange_type_set)
+    {
+        if      (exchange_type_ == ExchangeType::FANOUT)         exchange_type = AMQP::ExchangeType::fanout;
+        else if (exchange_type_ == ExchangeType::DIRECT)         exchange_type = AMQP::ExchangeType::direct;
+        else if (exchange_type_ == ExchangeType::TOPIC)          exchange_type = AMQP::ExchangeType::topic;
+        else if (exchange_type_ == ExchangeType::HASH)           exchange_type = AMQP::ExchangeType::consistent_hash;
+        else if (exchange_type_ == ExchangeType::HEADERS)        exchange_type = AMQP::ExchangeType::headers;
+        else throw Exception("Invalid exchange type", ErrorCodes::BAD_ARGUMENTS);
+    }
+    else
+    {
+        exchange_type = AMQP::ExchangeType::fanout;
+    }
 
     auto table_id = getStorageID();
     String table_name = table_id.table_name;
@@ -264,17 +288,17 @@ ConsumerBufferPtr StorageRabbitMQ::createReadBuffer()
     ChannelPtr consumer_channel = std::make_shared<AMQP::TcpChannel>(connection.get());
 
     return std::make_shared<ReadBufferFromRabbitMQConsumer>(
-        consumer_channel, event_handler, exchange_name, routing_keys,
-        next_channel_id, log, row_delimiter, bind_by_id, num_queues,
-        exchange_type, local_exchange_name, stream_cancelled);
+        consumer_channel, event_handler, exchange_name, exchange_type, routing_keys,
+        next_channel_id, log, row_delimiter, hash_exchange, num_queues,
+        local_exchange_name, stream_cancelled);
 }
 
 
 ProducerBufferPtr StorageRabbitMQ::createWriteBuffer()
 {
     return std::make_shared<WriteBufferToRabbitMQProducer>(
-        parsed_address, global_context, login_password, routing_keys[0], local_exchange_name,
-        log, num_consumers * num_queues, bind_by_id, use_transactional_channel,
+        parsed_address, global_context, login_password, routing_keys, exchange_name, exchange_type,
+        log, num_consumers * num_queues, use_transactional_channel,
         row_delimiter ? std::optional<char>{row_delimiter} : std::nullopt, 1, 1024);
 }
 

src/Storages/RabbitMQ/StorageRabbitMQ.h

@@ -76,15 +76,15 @@ private:
 
     Names routing_keys;
     const String exchange_name;
+    AMQP::ExchangeType exchange_type;
     String local_exchange_name;
 
     const String format_name;
     char row_delimiter;
     size_t num_consumers;
     size_t num_created_consumers = 0;
-    bool bind_by_id;
+    bool hash_exchange;
     size_t num_queues;
-    const String exchange_type;
     const bool use_transactional_channel;
 
     Poco::Logger * log;
@@ -99,6 +99,7 @@ private:
     std::mutex mutex;
     std::vector<ConsumerBufferPtr> buffers; /// available buffers for RabbitMQ consumers
 
+    bool exchange_type_set = false;
     size_t next_channel_id = 1; /// Must >= 1 because it is used as a binding key, which has to be > 0
     bool update_channel_id = false;
     std::atomic<bool> loop_started = false;

src/Storages/RabbitMQ/WriteBufferToRabbitMQProducer.cpp

@@ -16,6 +16,7 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int CANNOT_CONNECT_RABBITMQ;
+    extern const int LOGICAL_ERROR;
 }
 
 static const auto QUEUE_SIZE = 50000;
@@ -27,20 +28,20 @@ WriteBufferToRabbitMQProducer::WriteBufferToRabbitMQProducer(
         std::pair<String, UInt16> & parsed_address,
         Context & global_context,
         const std::pair<String, String> & login_password_,
-        const String & routing_key_,
-        const String & exchange_,
+        const Names & routing_keys_,
+        const String & exchange_name_,
+        const AMQP::ExchangeType exchange_type_,
         Poco::Logger * log_,
         size_t num_queues_,
-        bool bind_by_id_,
         bool use_transactional_channel_,
         std::optional<char> delimiter,
         size_t rows_per_message,
         size_t chunk_size_)
         : WriteBuffer(nullptr, 0)
         , login_password(login_password_)
-        , routing_key(routing_key_)
-        , exchange_name(exchange_ + "_direct")
-        , bind_by_id(bind_by_id_)
+        , routing_keys(routing_keys_)
+        , exchange_name(exchange_name_)
+        , exchange_type(exchange_type_)
         , num_queues(num_queues_)
         , use_transactional_channel(use_transactional_channel_)
         , payloads(QUEUE_SIZE * num_queues)
@@ -73,7 +74,6 @@ WriteBufferToRabbitMQProducer::WriteBufferToRabbitMQProducer(
     }
 
     producer_channel = std::make_shared<AMQP::TcpChannel>(connection.get());
-    checkExchange();
 
     /// If publishing should be wrapped in transactions
     if (use_transactional_channel)
@@ -83,6 +83,17 @@ WriteBufferToRabbitMQProducer::WriteBufferToRabbitMQProducer(
 
     writing_task = global_context.getSchedulePool().createTask("RabbitMQWritingTask", [this]{ writingFunc(); });
     writing_task->deactivate();
+
+    if (exchange_type == AMQP::ExchangeType::headers)
+    {
+        std::vector<String> matching;
+        for (const auto & header : routing_keys)
+        {
+            boost::split(matching, header, [](char c){ return c == '='; });
+            key_arguments[matching[0]] = matching[1];
+            matching.clear();
+        }
+    }
 }
 
 
@@ -90,7 +101,7 @@ WriteBufferToRabbitMQProducer::~WriteBufferToRabbitMQProducer()
 {
     stop_loop.store(true);
     writing_task->deactivate();
-    checkExchange();
+    initExchange();
 
     connection->close();
     assert(rows == 0 && chunks.empty());
@@ -133,28 +144,34 @@ void WriteBufferToRabbitMQProducer::writingFunc()
         while (!payloads.empty())
         {
             payloads.pop(payload);
-            next_queue = next_queue % num_queues + 1;
 
-            if (bind_by_id)
+            if (exchange_type == AMQP::ExchangeType::consistent_hash)
             {
+                next_queue = next_queue % num_queues + 1;
                 producer_channel->publish(exchange_name, std::to_string(next_queue), payload);
             }
+            else if (exchange_type == AMQP::ExchangeType::headers)
+            {
+                AMQP::Envelope envelope(payload.data(), payload.size());
+                envelope.setHeaders(key_arguments);
+                producer_channel->publish(exchange_name, "", envelope, key_arguments);
+            }
             else
             {
-                producer_channel->publish(exchange_name, routing_key, payload);
+                producer_channel->publish(exchange_name, routing_keys[0], payload);
             }
         }
+
         iterateEventLoop();
     }
 }
 
 
-void WriteBufferToRabbitMQProducer::checkExchange()
+void WriteBufferToRabbitMQProducer::initExchange()
 {
     std::atomic<bool> exchange_declared = false, exchange_error = false;
 
-    /// The AMQP::passive flag indicates that it should only be checked if there is a valid exchange with the given name.
-    producer_channel->declareExchange(exchange_name, AMQP::direct, AMQP::passive)
+    producer_channel->declareExchange(exchange_name, exchange_type)
     .onSuccess([&]()
     {
         exchange_declared = true;
@@ -162,10 +179,10 @@ void WriteBufferToRabbitMQProducer::checkExchange()
     .onError([&](const char * message)
     {
         exchange_error = true;
-        LOG_ERROR(log, "Exchange for INSERT query was not declared. Reason: {}", message);
+        LOG_ERROR(log, "Exchange error: {}", message);
     });
 
-    /// These variables are updated in a separate thread and starting the loop blocks current thread
+    /// These variables are updated in a separate thread.
     while (!exchange_declared && !exchange_error)
     {
         iterateEventLoop();
@@ -175,9 +192,6 @@ void WriteBufferToRabbitMQProducer::checkExchange()
 
 void WriteBufferToRabbitMQProducer::finilizeProducer()
 {
-    /// This will make sure everything is published
-    checkExchange();
-
     if (use_transactional_channel)
     {
         std::atomic<bool> answer_received = false, wait_rollback = false;

src/Storages/RabbitMQ/WriteBufferToRabbitMQProducer.h

@@ -23,11 +23,11 @@ public:
             std::pair<String, UInt16> & parsed_address,
             Context & global_context,
             const std::pair<String, String> & login_password_,
-            const String & routing_key_,
-            const String & exchange_,
+            const Names & routing_keys_,
+            const String & exchange_name_,
+            const AMQP::ExchangeType exchange_type_,
             Poco::Logger * log_,
             size_t num_queues_,
-            bool bind_by_id_,
             bool use_transactional_channel_,
             std::optional<char> delimiter,
             size_t rows_per_message,
@@ -38,21 +38,22 @@ public:
 
     void countRow();
     void activateWriting() { writing_task->activateAndSchedule(); }
+    void finilizeProducer();
 
 private:
     void nextImpl() override;
-    void checkExchange();
+    void initExchange();
     void iterateEventLoop();
     void writingFunc();
-    void finilizeProducer();
 
     const std::pair<String, String> login_password;
-    const String routing_key;
+    const Names routing_keys;
     const String exchange_name;
-    const bool bind_by_id;
+    AMQP::ExchangeType exchange_type;
     const size_t num_queues;
     const bool use_transactional_channel;
 
+    AMQP::Table key_arguments;
     BackgroundSchedulePool::TaskHolder writing_task;
     std::atomic<bool> stop_loop = false;
 

tests/integration/test_storage_rabbitmq/test.py

@@ -485,7 +485,6 @@ def test_rabbitmq_big_message(rabbitmq_cluster):
 
     while True:
         result = instance.query('SELECT count() FROM test.view')
-        print("Result", result, "Expected", batch_messages * rabbitmq_messages)
         if int(result) == batch_messages * rabbitmq_messages:
             break
 
@@ -552,7 +551,6 @@ def test_rabbitmq_sharding_between_channels_publish(rabbitmq_cluster):
     while True:
         result = instance.query('SELECT count() FROM test.view')
         time.sleep(1)
-        print("Result", result, "Expected", messages_num * threads_num)
         if int(result) == messages_num * threads_num:
             break
 
@@ -778,6 +776,7 @@ def test_rabbitmq_insert(rabbitmq_cluster):
             ENGINE = RabbitMQ
             SETTINGS rabbitmq_host_port = 'rabbitmq1:5672',
                      rabbitmq_exchange_name = 'insert',
+                     rabbitmq_exchange_type = 'direct',
                      rabbitmq_routing_key_list = 'insert1',
                      rabbitmq_format = 'TSV',
                      rabbitmq_row_delimiter = '\\n';
@@ -788,10 +787,64 @@ def test_rabbitmq_insert(rabbitmq_cluster):
     consumer_connection = pika.BlockingConnection(parameters)
 
     consumer = consumer_connection.channel()
-    consumer.exchange_declare(exchange='insert_rabbitmq_direct', exchange_type='direct')
+    consumer.exchange_declare(exchange='insert', exchange_type='direct')
     result = consumer.queue_declare(queue='')
     queue_name = result.method.queue
-    consumer.queue_bind(exchange='insert_rabbitmq_direct', queue=queue_name, routing_key='insert1')
+    consumer.queue_bind(exchange='insert', queue=queue_name, routing_key='insert1')
+
+    values = []
+    for i in range(50):
+        values.append("({i}, {i})".format(i=i))
+    values = ','.join(values)
+
+    while True:
+        try:
+            instance.query("INSERT INTO test.rabbitmq VALUES {}".format(values))
+            break
+        except QueryRuntimeException as e:
+            if 'Local: Timed out.' in str(e):
+                continue
+            else:
+                raise
+
+    insert_messages = []
+    def onReceived(channel, method, properties, body):
+        i = 0
+        insert_messages.append(body.decode())
+        if (len(insert_messages) == 50):
+            channel.stop_consuming()
+
+    consumer.basic_consume(onReceived, queue_name)
+    consumer.start_consuming()
+    consumer_connection.close()
+
+    result = '\n'.join(insert_messages)
+    rabbitmq_check_result(result, True)
+
+
+@pytest.mark.timeout(240)
+def test_rabbitmq_insert_headers_exchange(rabbitmq_cluster):
+    instance.query('''
+        CREATE TABLE test.rabbitmq (key UInt64, value UInt64)
+            ENGINE = RabbitMQ
+            SETTINGS rabbitmq_host_port = 'rabbitmq1:5672',
+                     rabbitmq_exchange_name = 'insert_headers',
+                     rabbitmq_exchange_type = 'headers',
+                     rabbitmq_routing_key_list = 'test=insert,topic=headers',
+                     rabbitmq_format = 'TSV',
+                     rabbitmq_row_delimiter = '\\n';
+    ''')
+
+    credentials = pika.PlainCredentials('root', 'clickhouse')
+    parameters = pika.ConnectionParameters('localhost', 5672, '/', credentials)
+    consumer_connection = pika.BlockingConnection(parameters)
+
+    consumer = consumer_connection.channel()
+    consumer.exchange_declare(exchange='insert_headers', exchange_type='headers')
+    result = consumer.queue_declare(queue='')
+    queue_name = result.method.queue
+    consumer.queue_bind(exchange='insert_headers', queue=queue_name, routing_key="",
+            arguments={'x-match':'all', 'test':'insert', 'topic':'headers'})
 
     values = []
     for i in range(50):
@@ -815,7 +868,6 @@ def test_rabbitmq_insert(rabbitmq_cluster):
         if (len(insert_messages) == 50):
             channel.stop_consuming()
 
-    consumer.basic_qos(prefetch_count=50)
     consumer.basic_consume(onReceived, queue_name)
     consumer.start_consuming()
     consumer_connection.close()
@@ -833,6 +885,8 @@ def test_rabbitmq_many_inserts(rabbitmq_cluster):
         CREATE TABLE test.rabbitmq_many (key UInt64, value UInt64)
             ENGINE = RabbitMQ
             SETTINGS rabbitmq_host_port = 'rabbitmq1:5672',
+                     rabbitmq_exchange_name = 'many_inserts',
+                     rabbitmq_exchange_type = 'direct',
                      rabbitmq_routing_key_list = 'insert2',
                      rabbitmq_format = 'TSV',
                      rabbitmq_row_delimiter = '\\n';
@@ -887,69 +941,6 @@ def test_rabbitmq_many_inserts(rabbitmq_cluster):
     assert int(result) == messages_num * threads_num, 'ClickHouse lost some messages: {}'.format(result)
 
 
-@pytest.mark.timeout(240)
-def test_rabbitmq_sharding_between_channels_and_queues_insert(rabbitmq_cluster):
-    instance.query('''
-        DROP TABLE IF EXISTS test.view_sharding;
-        DROP TABLE IF EXISTS test.consumer_sharding;
-        CREATE TABLE test.rabbitmq_sharding (key UInt64, value UInt64)
-            ENGINE = RabbitMQ
-            SETTINGS rabbitmq_host_port = 'rabbitmq1:5672',
-                     rabbitmq_num_consumers = 5,
-                     rabbitmq_num_queues = 2,
-                     rabbitmq_format = 'TSV',
-                     rabbitmq_row_delimiter = '\\n';
-        CREATE TABLE test.view_sharding (key UInt64, value UInt64)
-            ENGINE = MergeTree
-            ORDER BY key
-            SETTINGS old_parts_lifetime=5, cleanup_delay_period=2, cleanup_delay_period_random_add=3;
-        CREATE MATERIALIZED VIEW test.consumer_sharding TO test.view_sharding AS
-            SELECT * FROM test.rabbitmq_sharding;
-    ''')
-
-    messages_num = 10000
-    def insert():
-        values = []
-        for i in range(messages_num):
-            values.append("({i}, {i})".format(i=i))
-        values = ','.join(values)
-
-        while True:
-            try:
-                instance.query("INSERT INTO test.rabbitmq_sharding VALUES {}".format(values))
-                break
-            except QueryRuntimeException as e:
-                if 'Local: Timed out.' in str(e):
-                    continue
-                else:
-                    raise
-
-    threads = []
-    threads_num = 20
-    for _ in range(threads_num):
-        threads.append(threading.Thread(target=insert))
-    for thread in threads:
-        time.sleep(random.uniform(0, 1))
-        thread.start()
-
-    while True:
-        result = instance.query('SELECT count() FROM test.view_sharding')
-        time.sleep(1)
-        if int(result) == messages_num * threads_num:
-            break
-
-    instance.query('''
-        DROP TABLE IF EXISTS test.rabbitmq_sharding;
-        DROP TABLE IF EXISTS test.consumer_sharding;
-        DROP TABLE IF EXISTS test.view_sharding;
-    ''')
-
-    for thread in threads:
-        thread.join()
-
-    assert int(result) == messages_num * threads_num, 'ClickHouse lost some messages: {}'.format(result)
-
-
 @pytest.mark.timeout(420)
 def test_rabbitmq_overloaded_insert(rabbitmq_cluster):
     instance.query('''
@@ -958,6 +949,9 @@ def test_rabbitmq_overloaded_insert(rabbitmq_cluster):
         CREATE TABLE test.rabbitmq_overload (key UInt64, value UInt64)
             ENGINE = RabbitMQ
             SETTINGS rabbitmq_host_port = 'rabbitmq1:5672',
+                     rabbitmq_exchange_name = 'over',
+                     rabbitmq_exchange_type = 'direct',
+                     rabbitmq_routing_key_list = 'over',
                      rabbitmq_num_consumers = 10,
                      rabbitmq_format = 'TSV',
                      rabbitmq_row_delimiter = '\\n';