kafka_consumer.hpp

/*
 * Copyright (c) 2020-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include "kafka_callback.hpp"
 
#include <cudf/io/datasource.hpp>
 
#include <librdkafka/rdkafkacpp.h>
 
#include <algorithm>
#include <chrono>
#include <map>
#include <memory>
#include <string>
 
namespace cudf {
namespace io {
namespace external {
namespace kafka {
 
class kafka_consumer : public cudf::io::datasource {
 public:
  kafka_consumer(std::map<std::string, std::string> configs,
                 python_callable_type python_callable,
                 kafka_oauth_callback_wrapper_type callable_wrapper);
 
  kafka_consumer(std::map<std::string, std::string> configs,
                 python_callable_type python_callable,
                 kafka_oauth_callback_wrapper_type callable_wrapper,
                 std::string const& topic_name,
                 int partition,
                 int64_t start_offset,
                 int64_t end_offset,
                 int batch_timeout,
                 std::string const& delimiter);
 
  std::unique_ptr<cudf::io::datasource::buffer> host_read(size_t offset, size_t size) override;
 
  size_t size() const override;
 
  size_t host_read(size_t offset, size_t size, uint8_t* dst) override;
 
  void commit_offset(std::string const& topic, int partition, int64_t offset);
 
  std::map<std::string, int64_t> get_watermark_offset(std::string const& topic,
                                                      int partition,
                                                      int timeout,
                                                      bool cached);
 
  std::map<std::string, std::string> current_configs();
 
  int64_t get_committed_offset(std::string const& topic, int partition);
 
  std::map<std::string, std::vector<int32_t>> list_topics(std::string specific_topic);
 
  void close(int timeout);
 
  void unsubscribe();
 
  virtual ~kafka_consumer(){};
 
 private:
  std::unique_ptr<RdKafka::Conf> kafka_conf;  // RDKafka configuration object
  std::unique_ptr<RdKafka::KafkaConsumer> consumer;
 
  std::map<std::string, std::string> configs;
  python_callable_type python_callable_;
  kafka_oauth_callback_wrapper_type callable_wrapper_;
 
  std::string topic_name;
  int partition;
  int64_t start_offset;
  int64_t end_offset;
  int batch_timeout;
  int default_timeout = 10000;  // milliseconds
  std::string delimiter;
 
  std::string buffer;
 
 private:
  RdKafka::ErrorCode update_consumer_topic_partition_assignment(std::string const& topic,
                                                                int partition,
                                                                int64_t offset);
 
  int64_t now();
 
  void consume_to_buffer();
};
 
}  // namespace kafka
}  // namespace external
}  // namespace io
}  // namespace cudf