140 lines
6.1 KiB
C++
140 lines
6.1 KiB
C++
// Copyright 2022 Memgraph Ltd.
|
|
//
|
|
// Use of this software is governed by the Business Source License
|
|
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
|
|
// License, and you may not use this file except in compliance with the Business Source License.
|
|
//
|
|
// As of the Change Date specified in that file, in accordance with
|
|
// the Business Source License, use of this software will be governed
|
|
// by the Apache License, Version 2.0, included in the file
|
|
// licenses/APL.txt.
|
|
|
|
#include "query/stream/sources.hpp"
|
|
|
|
#include <json/json.hpp>
|
|
|
|
#include "integrations/constants.hpp"
|
|
|
|
namespace query::stream {
|
|
KafkaStream::KafkaStream(std::string stream_name, StreamInfo stream_info,
|
|
ConsumerFunction<integrations::kafka::Message> consumer_function) {
|
|
integrations::kafka::ConsumerInfo consumer_info{
|
|
.consumer_name = std::move(stream_name),
|
|
.topics = std::move(stream_info.topics),
|
|
.consumer_group = std::move(stream_info.consumer_group),
|
|
.bootstrap_servers = std::move(stream_info.bootstrap_servers),
|
|
.batch_interval = stream_info.common_info.batch_interval,
|
|
.batch_size = stream_info.common_info.batch_size,
|
|
.public_configs = stream_info.configs,
|
|
.private_configs = stream_info.credentials,
|
|
};
|
|
consumer_.emplace(std::move(consumer_info), std::move(consumer_function));
|
|
};
|
|
|
|
KafkaStream::StreamInfo KafkaStream::Info(std::string transformation_name) const {
|
|
const auto &info = consumer_->Info();
|
|
using CredentialsType = decltype(StreamInfo::credentials);
|
|
CredentialsType reducted_credentials;
|
|
std::transform(info.private_configs.begin(), info.private_configs.end(),
|
|
std::inserter(reducted_credentials, reducted_credentials.end()),
|
|
[](const auto &pair) -> CredentialsType::value_type {
|
|
return {pair.first, integrations::kReducted};
|
|
});
|
|
return {{.batch_interval = info.batch_interval,
|
|
.batch_size = info.batch_size,
|
|
.transformation_name = std::move(transformation_name)},
|
|
.topics = info.topics,
|
|
.consumer_group = info.consumer_group,
|
|
.bootstrap_servers = info.bootstrap_servers,
|
|
.configs = info.public_configs,
|
|
.credentials = std::move(reducted_credentials)};
|
|
}
|
|
|
|
void KafkaStream::Start() { consumer_->Start(); }
|
|
void KafkaStream::Stop() { consumer_->Stop(); }
|
|
bool KafkaStream::IsRunning() const { return consumer_->IsRunning(); }
|
|
|
|
void KafkaStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
|
|
const ConsumerFunction<integrations::kafka::Message> &consumer_function) const {
|
|
consumer_->Check(timeout, batch_limit, consumer_function);
|
|
}
|
|
|
|
utils::BasicResult<std::string> KafkaStream::SetStreamOffset(const int64_t offset) {
|
|
return consumer_->SetConsumerOffsets(offset);
|
|
}
|
|
|
|
namespace {
|
|
const std::string kTopicsKey{"topics"};
|
|
const std::string kConsumerGroupKey{"consumer_group"};
|
|
const std::string kBoostrapServers{"bootstrap_servers"};
|
|
const std::string kConfigs{"configs"};
|
|
const std::string kCredentials{"credentials"};
|
|
|
|
const std::unordered_map<std::string, std::string> kDefaultConfigsMap;
|
|
} // namespace
|
|
|
|
void to_json(nlohmann::json &data, KafkaStream::StreamInfo &&info) {
|
|
data[kCommonInfoKey] = std::move(info.common_info);
|
|
data[kTopicsKey] = std::move(info.topics);
|
|
data[kConsumerGroupKey] = info.consumer_group;
|
|
data[kBoostrapServers] = std::move(info.bootstrap_servers);
|
|
data[kConfigs] = std::move(info.configs);
|
|
data[kCredentials] = std::move(info.credentials);
|
|
}
|
|
|
|
void from_json(const nlohmann::json &data, KafkaStream::StreamInfo &info) {
|
|
data.at(kCommonInfoKey).get_to(info.common_info);
|
|
data.at(kTopicsKey).get_to(info.topics);
|
|
data.at(kConsumerGroupKey).get_to(info.consumer_group);
|
|
data.at(kBoostrapServers).get_to(info.bootstrap_servers);
|
|
// These values might not be present in the persisted JSON object
|
|
info.configs = data.value(kConfigs, kDefaultConfigsMap);
|
|
info.credentials = data.value(kCredentials, kDefaultConfigsMap);
|
|
}
|
|
|
|
PulsarStream::PulsarStream(std::string stream_name, StreamInfo stream_info,
|
|
ConsumerFunction<integrations::pulsar::Message> consumer_function) {
|
|
integrations::pulsar::ConsumerInfo consumer_info{.batch_size = stream_info.common_info.batch_size,
|
|
.batch_interval = stream_info.common_info.batch_interval,
|
|
.topics = std::move(stream_info.topics),
|
|
.consumer_name = std::move(stream_name),
|
|
.service_url = std::move(stream_info.service_url)};
|
|
|
|
consumer_.emplace(std::move(consumer_info), std::move(consumer_function));
|
|
};
|
|
|
|
PulsarStream::StreamInfo PulsarStream::Info(std::string transformation_name) const {
|
|
const auto &info = consumer_->Info();
|
|
return {{.batch_interval = info.batch_interval,
|
|
.batch_size = info.batch_size,
|
|
.transformation_name = std::move(transformation_name)},
|
|
.topics = info.topics,
|
|
.service_url = info.service_url};
|
|
}
|
|
|
|
void PulsarStream::Start() { consumer_->Start(); }
|
|
void PulsarStream::Stop() { consumer_->Stop(); }
|
|
bool PulsarStream::IsRunning() const { return consumer_->IsRunning(); }
|
|
|
|
void PulsarStream::Check(std::optional<std::chrono::milliseconds> timeout, std::optional<int64_t> batch_limit,
|
|
const ConsumerFunction<Message> &consumer_function) const {
|
|
consumer_->Check(timeout, batch_limit, consumer_function);
|
|
}
|
|
|
|
namespace {
|
|
const std::string kServiceUrl{"service_url"};
|
|
} // namespace
|
|
|
|
void to_json(nlohmann::json &data, PulsarStream::StreamInfo &&info) {
|
|
data[kCommonInfoKey] = std::move(info.common_info);
|
|
data[kTopicsKey] = std::move(info.topics);
|
|
data[kServiceUrl] = std::move(info.service_url);
|
|
}
|
|
|
|
void from_json(const nlohmann::json &data, PulsarStream::StreamInfo &info) {
|
|
data.at(kCommonInfoKey).get_to(info.common_info);
|
|
data.at(kTopicsKey).get_to(info.topics);
|
|
data.at(kServiceUrl).get_to(info.service_url);
|
|
}
|
|
} // namespace query::stream
|