feat: Kafka external streams support SCRAM-SHA-256/512 (#454)

Author: zliang-min

src/KafkaLog/KafkaWAL.cpp

@@ -365,24 +365,14 @@ void KafkaWAL::initProducerHandle()
         {"message.max.bytes", std::to_string(settings->message_max_bytes)},
         {"message.timeout.ms", std::to_string(settings->message_timeout_ms)},
         /// Protocol used to communicate with brokers.
-        {"security.protocol", settings->auth.security_protocol.c_str()},
         {"topic.metadata.refresh.interval.ms", std::to_string(settings->topic_metadata_refresh_interval_ms)},
         {"compression.codec", settings->compression_codec.c_str()},
     };
 
     if (!settings->debug.empty())
         producer_params.emplace_back("debug", settings->debug);
 
-    if (boost::iequals(settings->auth.security_protocol, "SASL_PLAINTEXT")
-        || boost::iequals(settings->auth.security_protocol, "SASL_SSL"))
-    {
-        producer_params.emplace_back("sasl.mechanisms", "PLAIN");
-        producer_params.emplace_back("sasl.username", settings->auth.username.c_str());
-        producer_params.emplace_back("sasl.password", settings->auth.password.c_str());
-    }
-
-    if (boost::iequals(settings->auth.security_protocol, "SASL_SSL") && !settings->auth.ssl_ca_cert_file.empty())
-        producer_params.emplace_back("ssl.ca.location", settings->auth.ssl_ca_cert_file.c_str());
+    settings->auth.populateConfigs(producer_params);
 
     auto cb_setup = [](rd_kafka_conf_t * kconf) {
         rd_kafka_conf_set_stats_cb(kconf, &KafkaWALStats::logStats);

src/KafkaLog/KafkaWALConsumer.cpp

@@ -98,19 +98,9 @@ void KafkaWALConsumer::initHandle()
         /// ensuring no on-the-wire or on-disk corruption to the messages occurred
         std::make_pair("check.crcs", std::to_string(settings->check_crcs)),
         std::make_pair("statistics.interval.ms", std::to_string(settings->statistic_internal_ms)),
-        std::make_pair("security.protocol", settings->auth.security_protocol.c_str()),
     };
 
-    if (boost::iequals(settings->auth.security_protocol, "SASL_PLAINTEXT")
-        || boost::iequals(settings->auth.security_protocol, "SASL_SSL"))
-    {
-        consumer_params.emplace_back("sasl.mechanisms", "PLAIN");
-        consumer_params.emplace_back("sasl.username", settings->auth.username.c_str());
-        consumer_params.emplace_back("sasl.password", settings->auth.password.c_str());
-    }
-
-    if (boost::iequals(settings->auth.security_protocol, "SASL_SSL") && !settings->auth.ssl_ca_cert_file.empty())
-        consumer_params.emplace_back("ssl.ca.location", settings->auth.ssl_ca_cert_file.c_str());
+    settings->auth.populateConfigs(consumer_params);
 
     if (!settings->debug.empty())
         consumer_params.emplace_back("debug", settings->debug);
@@ -141,11 +131,11 @@ int32_t KafkaWALConsumer::addSubscriptions(const TopicPartitionOffsets & partiti
 
     for (const auto & partition : partitions_)
     {
-        auto new_partition = rd_kafka_topic_partition_list_add(topic_partitions.get(), partition.topic.c_str(), partition.partition);
+        auto * new_partition = rd_kafka_topic_partition_list_add(topic_partitions.get(), partition.topic.c_str(), partition.partition);
         new_partition->offset = partition.offset;
     }
 
-    auto err = rd_kafka_incremental_assign(consumer_handle.get(), topic_partitions.get());
+    auto * err = rd_kafka_incremental_assign(consumer_handle.get(), topic_partitions.get());
     if (err)
     {
         LOG_ERROR(log, "Failed to assign partitions incrementally, error={}", rd_kafka_error_string(err));
@@ -168,7 +158,7 @@ int32_t KafkaWALConsumer::removeSubscriptions(const TopicPartitionOffsets & part
         rd_kafka_topic_partition_list_add(topic_partitions.get(), partition.topic.c_str(), partition.partition);
     }
 
-    auto err = rd_kafka_incremental_unassign(consumer_handle.get(), topic_partitions.get());
+    auto * err = rd_kafka_incremental_unassign(consumer_handle.get(), topic_partitions.get());
     if (err)
     {
         LOG_ERROR(log, "Failed to unassign partitions incrementally, error={}", rd_kafka_error_string(err));
@@ -193,7 +183,7 @@ ConsumeResult KafkaWALConsumer::consume(uint32_t count, int32_t timeout_ms, std:
 
     for (uint32_t i = 0; i < count; ++i)
     {
-        auto rkmessage = rd_kafka_consumer_poll(consumer_handle.get(), timeout_ms);
+        auto * rkmessage = rd_kafka_consumer_poll(consumer_handle.get(), timeout_ms);
         if (likely(rkmessage))
         {
             if (likely(rkmessage->err == RD_KAFKA_RESP_ERR_NO_ERROR))
@@ -254,7 +244,7 @@ int32_t KafkaWALConsumer::commit(const TopicPartitionOffsets & tpos)
 
     for (const auto & tpo : tpos)
     {
-        auto partition_offset = rd_kafka_topic_partition_list_add(topic_partition_list.get(), tpo.topic.c_str(), tpo.partition);
+        auto * partition_offset = rd_kafka_topic_partition_list_add(topic_partition_list.get(), tpo.topic.c_str(), tpo.partition);
 
         /// rd_kafka_offsets_store commits `offset` as it is. We add 1 to the offset
         /// to keep the same semantic as rd_kafka_offset_store

src/KafkaLog/KafkaWALPool.h

@@ -31,7 +31,7 @@ class KafkaWALPool : private boost::noncopyable
 
     KafkaWALSimpleConsumerPtr getOrCreateStreaming(const String & cluster_id);
 
-    KafkaWALSimpleConsumerPtr getOrCreateStreamingExternal(const String & brokers, const KafkaWALAuth & auth, int32_t fetch_max_wait_ms = 200);
+    KafkaWALSimpleConsumerPtr getOrCreateStreamingExternal(const String & brokers, const KafkaWALAuth & auth, int32_t fetch_wait_max_ms = 200);
 
     std::vector<KafkaWALClusterPtr> clusters(const KafkaWALContext & ctx) const;
 

src/KafkaLog/KafkaWALSettings.h

@@ -2,9 +2,7 @@
 
 #include <boost/algorithm/string/join.hpp>
 
-#include <string>
-#include <vector>
-
+#include <boost/algorithm/string/predicate.hpp>
 #include <fmt/format.h>
 
 namespace klog
@@ -14,6 +12,7 @@ struct KafkaWALAuth
     std::string security_protocol;
     std::string username;
     std::string password;
+    std::string sasl_mechanism;
     std::string ssl_ca_cert_file;
 
     bool operator==(const KafkaWALAuth & o) const
@@ -23,6 +22,31 @@ struct KafkaWALAuth
             && password == o.password
             && ssl_ca_cert_file == o.ssl_ca_cert_file;
     }
+
+    bool usesSASL() const
+    {
+        return boost::istarts_with(security_protocol, "SASL_");
+    }
+
+    /// "SASL_SSL" or "SSL"
+    bool usesSecureConnection() const
+    {
+        return boost::iends_with(security_protocol, "SSL");
+    }
+
+    void populateConfigs(std::vector<std::pair<std::string, std::string>> & params) const
+    {
+        params.emplace_back("security.protocol", security_protocol);
+        if (usesSASL())
+        {
+            params.emplace_back("sasl.mechanism", sasl_mechanism);
+            params.emplace_back("sasl.username", username);
+            params.emplace_back("sasl.password", password);
+        }
+
+        if (usesSecureConnection() && !ssl_ca_cert_file.empty())
+            params.emplace_back("ssl.ca.location", ssl_ca_cert_file);
+        }
 };
 
 struct KafkaWALSettings
@@ -34,18 +58,12 @@ struct KafkaWALSettings
     /// Global settings for both producer and consumer /// global metrics
     /// comma separated host/port: host1:port,host2:port,...
     std::string brokers;
-    KafkaWALAuth auth = {
-        .security_protocol = "plaintext",
-        .username = "",
-        .password = "",
-        .ssl_ca_cert_file = ""
-    };
-    /// FIXME, SASL, SSL etc support
+    KafkaWALAuth auth = { .security_protocol = "plaintext" };
 
     int32_t message_max_bytes = 1000000;
     int32_t topic_metadata_refresh_interval_ms = 300000;
     int32_t statistic_internal_ms = 30000;
-    std::string debug = "";
+    std::string debug;
 
     /////////////////////////////////////////////////////
 

src/KafkaLog/KafkaWALSimpleConsumer.cpp

@@ -87,24 +87,14 @@ void KafkaWALSimpleConsumer::initHandle()
         {"enable.partition.eof", "false"},
         {"queued.min.messages", std::to_string(settings->queued_min_messages)},
         {"queued.max.messages.kbytes", std::to_string(settings->queued_max_messages_kbytes)},
-        {"security.protocol", settings->auth.security_protocol.c_str()},
     };
 
     if (!settings->debug.empty())
     {
         consumer_params.emplace_back("debug", settings->debug);
     }
 
-    if (boost::iequals(settings->auth.security_protocol, "SASL_PLAINTEXT")
-        || boost::iequals(settings->auth.security_protocol, "SASL_SSL"))
-    {
-        consumer_params.emplace_back("sasl.mechanisms", "PLAIN");
-        consumer_params.emplace_back("sasl.username", settings->auth.username.c_str());
-        consumer_params.emplace_back("sasl.password", settings->auth.password.c_str());
-    }
-
-    if (boost::iequals(settings->auth.security_protocol, "SASL_SSL") && !settings->auth.ssl_ca_cert_file.empty())
-        consumer_params.emplace_back("ssl.ca.location", settings->auth.ssl_ca_cert_file.c_str());
+    settings->auth.populateConfigs(consumer_params);
 
     auto cb_setup = [](rd_kafka_conf_t * kconf) { /// STYLE_CHECK_ALLOW_BRACE_SAME_LINE_LAMBDA
         rd_kafka_conf_set_stats_cb(kconf, &KafkaWALStats::logStats);

src/Storages/ExternalStream/ExternalStreamSettings.h

@@ -15,6 +15,7 @@ class ASTStorage;
     M(String, security_protocol, "plaintext", "The protocol to connection external logstore", 0) \
     M(String, username, "", "The username of external logstore", 0) \
     M(String, password, "", "The password of external logstore", 0) \
+    M(String, sasl_mechanism, "PLAIN", "SASL mechanism to use for authentication. Supported: PLAIN, SCRAM-SHA-256, SCRAM-SHA-512. Default to PLAIN when SASL is enabled.", 0) \
     M(String, ssl_ca_cert_file, "", "The path of ssl ca cert file", 0) \
     M(String, properties, "", "A semi-colon-separated key-value pairs for configuring the kafka client used by the external stream. A key-value pair is separated by a equal sign. Example: 'client.id=my-client-id;group.id=my-group-id'. Note, not all properties are supported, please check the document for supported properties.", 0) \
     M(String, sharding_expr, "", "An expression which will be evaluated on each row of data returned by the query to calculate the an integer which will be used to determine the ID of the partition to which the row of data will be sent. If not set, data are sent to any partition randomly.", 0) \

src/Storages/ExternalStream/Kafka/Kafka.cpp

@@ -32,6 +32,12 @@ Kafka::Kafka(IStorage * storage, std::unique_ptr<ExternalStreamSettings> setting
     , data_format(StorageExternalStreamImpl::dataFormat())
     , log(&Poco::Logger::get("External-" + settings->topic.value))
     , engine_args(engine_args_)
+    , auth_info(std::make_unique<klog::KafkaWALAuth>(
+        settings->security_protocol.value,
+        settings->username.value,
+        settings->password.value,
+        settings->sasl_mechanism.value,
+        settings->ssl_ca_cert_file.value))
     , external_stream_counter(external_stream_counter_)
 {
     assert(settings->type.value == StreamTypes::KAFKA || settings->type.value == StreamTypes::REDPANDA);
@@ -145,6 +151,11 @@ void Kafka::cacheVirtualColumnNamesAndTypes()
     virtual_column_names_and_types.push_back(NameAndTypePair(ProtonConsts::RESERVED_EVENT_SEQUENCE_ID, std::make_shared<DataTypeInt64>()));
 }
 
+klog::KafkaWALSimpleConsumerPtr Kafka::getConsumer(int32_t fetch_wait_max_ms) const
+{
+    return klog::KafkaWALPool::instance(nullptr).getOrCreateStreamingExternal(settings->brokers.value, *auth_info, fetch_wait_max_ms);
+}
+
 std::vector<Int64> Kafka::getOffsets(const SeekToInfoPtr & seek_to_info, const std::vector<int32_t> & shards_to_query) const
 {
     assert(seek_to_info);
@@ -155,15 +166,6 @@ std::vector<Int64> Kafka::getOffsets(const SeekToInfoPtr & seek_to_info, const s
     }
     else
     {
-        klog::KafkaWALAuth auth{
-            .security_protocol = securityProtocol(),
-            .username = username(),
-            .password = password(),
-            .ssl_ca_cert_file = sslCaCertFile(),
-        };
-
-        auto consumer = klog::KafkaWALPool::instance(nullptr).getOrCreateStreamingExternal(settings->brokers.value, auth);
-
         std::vector<klog::PartitionTimestamp> partition_timestamps;
         partition_timestamps.reserve(shards_to_query.size());
         auto seek_timestamps{seek_to_info->getSeekPoints()};
@@ -172,7 +174,7 @@ std::vector<Int64> Kafka::getOffsets(const SeekToInfoPtr & seek_to_info, const s
         for (auto [shard, timestamp] : std::ranges::views::zip(shards_to_query, seek_timestamps))
             partition_timestamps.emplace_back(shard, timestamp);
 
-        return consumer->offsetsForTimestamps(settings->topic.value, partition_timestamps);
+        return getConsumer()->offsetsForTimestamps(settings->topic.value, partition_timestamps);
     }
 }
 
@@ -241,16 +243,7 @@ void Kafka::validate(const std::vector<int32_t> & shards_to_query)
         if (shards == 0)
         {
             /// We haven't describe the topic yet
-            klog::KafkaWALAuth auth = {
-                .security_protocol = settings->security_protocol.value,
-                .username = settings->username.value,
-                .password = settings->password.value,
-                .ssl_ca_cert_file = settings->ssl_ca_cert_file.value,
-            };
-
-            auto consumer = klog::KafkaWALPool::instance(nullptr).getOrCreateStreamingExternal(settings->brokers.value, auth);
-
-            auto result = consumer->describe(settings->topic.value);
+            auto result = getConsumer()->describe(settings->topic.value);
             if (result.err != ErrorCodes::OK)
                 throw Exception(ErrorCodes::RESOURCE_NOT_FOUND, "{} topic doesn't exist", settings->topic.value);
 

src/Storages/ExternalStream/Kafka/Kafka.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <KafkaLog/KafkaWALCommon.h>
+#include <KafkaLog/KafkaWALSimpleConsumer.h>
 #include <Storages/ExternalStream/ExternalStreamSettings.h>
 #include <Storages/ExternalStream/StorageExternalStreamImpl.h>
 #include <Storages/ExternalStream/ExternalStreamCounter.h>
@@ -38,13 +39,11 @@ class Kafka final : public StorageExternalStreamImpl
     const String & brokers() const { return settings->brokers.value; }
     const String & dataFormat() const override { return data_format; }
     const String & topic() const { return settings->topic.value; }
-    const String & securityProtocol() const { return settings->security_protocol.value; }
-    const String & username() const { return settings->username.value; }
-    const String & password() const { return settings->password.value; }
-    const String & sslCaCertFile() const { return settings->ssl_ca_cert_file.value; }
     const klog::KConfParams & properties() const { return kafka_properties; }
+    const klog::KafkaWALAuth & auth() const noexcept { return *auth_info; }
     bool hasCustomShardingExpr() const { return !engine_args.empty(); }
     const ASTPtr & shardingExprAst() const { assert(!engine_args.empty()); return engine_args[0]; }
+    klog::KafkaWALSimpleConsumerPtr getConsumer(int32_t fetch_wait_max_ms = 200) const;
 
 private:
     void calculateDataFormat(const IStorage * storage);
@@ -63,6 +62,7 @@ class Kafka final : public StorageExternalStreamImpl
     NamesAndTypesList virtual_column_names_and_types;
     klog::KConfParams kafka_properties;
     const ASTs engine_args;
+    const std::unique_ptr<klog::KafkaWALAuth> auth_info;
 
     std::mutex shards_mutex;
     int32_t shards = 0;

src/Storages/ExternalStream/Kafka/KafkaSink.cpp

@@ -174,17 +174,7 @@ KafkaSink::KafkaSink(const Kafka * kafka, const Block & header, ContextPtr conte
 
     /// properies from settings have higher priority
     producer_params.emplace_back("bootstrap.servers", kafka->brokers());
-    producer_params.emplace_back("security.protocol", kafka->securityProtocol());
-    if (boost::iequals(kafka->securityProtocol(), "SASL_PLAINTEXT")
-        || boost::iequals(kafka->securityProtocol(), "SASL_SSL"))
-    {
-        producer_params.emplace_back("sasl.mechanisms", "PLAIN");
-        producer_params.emplace_back("sasl.username", kafka->username());
-        producer_params.emplace_back("sasl.password", kafka->password());
-    }
-
-    if (boost::iequals(kafka->securityProtocol(), "SASL_SSL") && !kafka->sslCaCertFile().empty())
-        producer_params.emplace_back("ssl.ca.location", kafka->sslCaCertFile());
+    kafka->auth().populateConfigs(producer_params);
 
     auto * conf = rd_kafka_conf_new();
     char errstr[512]{'\0'};

src/Storages/ExternalStream/Kafka/KafkaSource.cpp

@@ -278,13 +278,7 @@ void KafkaSource::initConsumer(const Kafka * kafka)
         consume_ctx.auto_offset_reset = "earliest";
 
     consume_ctx.enforce_offset = true;
-    klog::KafkaWALAuth auth = {
-        .security_protocol = kafka->securityProtocol(),
-        .username = kafka->username(),
-        .password = kafka->password(),
-        .ssl_ca_cert_file = kafka->sslCaCertFile()
-    };
-    consumer = klog::KafkaWALPool::instance(nullptr).getOrCreateStreamingExternal(kafka->brokers(), auth, record_consume_timeout_ms);
+    consumer = kafka->getConsumer(record_consume_timeout_ms);
     consumer->initTopicHandle(consume_ctx);
 }