Based on StreamInsert code but using binary columnar format for inserting a stream of rows with optional transformations from stdin to a DB table. More...

#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <cstring>
#include <iostream>
#include <iterator>
#include <string>
#include "Logger/Logger.h"
#include "RowToColumnLoader.h"
#include "Shared/ThriftClient.h"
#include "Shared/clean_boost_regex.hpp"
#include "Shared/sqltypes.h"
#include <chrono>
#include <thread>
#include <boost/program_options.hpp>
#include <librdkafka/rdkafkacpp.h>

Include dependency graph for KafkaImporter.cpp:

Go to the source code of this file.

Classes
class	RebalanceCb

class	ConsumeCb

class	EventCb

struct	stuff

Macros
#define	MAX_FIELD_LEN 20000

Functions
bool	msg_consume (RdKafka::Message *message, RowToColumnLoader &row_loader, import_export::CopyParams copy_params, const std::map< std::string, std::pair< std::unique_ptr< boost::regex >, std::unique_ptr< std::string >>> &transformations, const bool remove_quotes)

void	kafka_insert (RowToColumnLoader &row_loader, const std::map< std::string, std::pair< std::unique_ptr< boost::regex >, std::unique_ptr< std::string >>> &transformations, const import_export::CopyParams &copy_params, const bool remove_quotes, std::string group_id, std::string topic, std::string brokers)

int	main (int argc, char **argv)

Variables
bool	print_error_data = false

bool	print_transformation = false

static bool	run = true

static bool	exit_eof = false

static int	eof_cnt = 0

static int	partition_cnt = 0

static long	msg_cnt = 0

static int64_t	msg_bytes = 0

Detailed Description

Based on StreamInsert code but using binary columnar format for inserting a stream of rows with optional transformations from stdin to a DB table.

Definition in file KafkaImporter.cpp.

Macro Definition Documentation

#define MAX_FIELD_LEN 20000

Definition at line 44 of file KafkaImporter.cpp.

Referenced by msg_consume().

Function Documentation

void kafka_insert	(	RowToColumnLoader &	row_loader,
		const std::map< std::string, std::pair< std::unique_ptr< boost::regex >, std::unique_ptr< std::string >>> &	transformations,
		const import_export::CopyParams &	copy_params,
		const bool	remove_quotes,
		std::string	group_id,
		std::string	topic,
		std::string	brokers
	)

Definition at line 291 of file KafkaImporter.cpp.

References import_export::CopyParams::batch_size, File_Namespace::create(), RowToColumnLoader::do_load(), logger::ERROR, logger::FATAL, logger::INFO, LOG, msg_bytes, msg_cnt, msg_consume(), run_benchmark_import::rows_loaded, run, and setup::version.

Referenced by main().

                        {
   std::string errstr;
   std::string topic_str;
   std::string mode;
   std::string debug;
   std::vector<std::string> topics;
   bool do_conf_dump = false;
   int use_ccb = 0;
 
   RebalanceCb ex_rebalance_cb;
 
   /*
    * Create configuration objects
    */
   RdKafka::Conf* conf = RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL);
   RdKafka::Conf* tconf = RdKafka::Conf::create(RdKafka::Conf::CONF_TOPIC);
 
   conf->set("rebalance_cb", &ex_rebalance_cb, errstr);
 
   if (conf->set("group.id", group_id, errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << "could not set  group.id " << errstr;
   }
 
   if (conf->set("compression.codec", "none", errstr) != /* can also be gzip or snappy */
       RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   if (conf->set("statistics.interval.ms", "1000", errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
   if (conf->set("enable.auto.commit", "false", errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   if (tconf->set("auto.offset.reset", "earliest", errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   if (tconf->set("enable.auto.commit", "false", errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   do_conf_dump = true;
 
   topics.push_back(topic);
 
   LOG(INFO) << "Version " << RdKafka::version_str().c_str();
   LOG(INFO) << RdKafka::version();
   LOG(INFO) << RdKafka::get_debug_contexts().c_str();
 
   conf->set("metadata.broker.list", brokers, errstr);
 
   // debug = "none";
   if (!debug.empty()) {
     if (conf->set("debug", debug, errstr) != RdKafka::Conf::CONF_OK) {
       LOG(FATAL) << errstr;
     }
   }
 
   ConsumeCb consume_cb;
   use_ccb = 0;
   if (use_ccb) {
     if (conf->set("consume_cb", &consume_cb, errstr) != RdKafka::Conf::CONF_OK) {
       LOG(FATAL) << errstr;
     }
     // need to set the opaque pointer here for the callbacks
     //        rd_kafka_conf_set_opaque(conf, this);
   }
 
   EventCb ex_event_cb;
   if (conf->set("event_cb", &ex_event_cb, errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   if (conf->set("default_topic_conf", tconf, errstr) != RdKafka::Conf::CONF_OK) {
     LOG(FATAL) << errstr;
   }
 
   if (do_conf_dump) {
     int pass;
 
     for (pass = 0; pass < 2; pass++) {
       std::list<std::string>* dump;
       if (pass == 0) {
         dump = conf->dump();
         LOG(INFO) << "# Global config";
         LOG(INFO) << "===============";
       } else {
         dump = tconf->dump();
         LOG(INFO) << "# Topic config";
         LOG(INFO) << "===============";
       }
 
       for (std::list<std::string>::iterator it = dump->begin(); it != dump->end();) {
         std::string ts = *it;
         it++;
         LOG(INFO) << ts << " = " << *it;
         it++;
       }
       LOG(INFO) << "Dump config finished";
     }
   }
   LOG(INFO) << "FULL Dump config finished";
 
   delete tconf;
 
   /*
    * Create consumer using accumulated global configuration.
    */
   RdKafka::KafkaConsumer* consumer = RdKafka::KafkaConsumer::create(conf, errstr);
   if (!consumer) {
     LOG(ERROR) << "Failed to create consumer: " << errstr;
   }
 
   delete conf;
 
   LOG(INFO) << " Created consumer " << consumer->name();
 
   /*
    * Subscribe to topics
    */
   RdKafka::ErrorCode err = consumer->subscribe(topics);
   if (err) {
     LOG(FATAL) << "Failed to subscribe to " << topics.size()
                << " topics: " << RdKafka::err2str(err);
   }
 
   /*
    * Consume messages
    */
   size_t recv_rows = 0;
   int skipped = 0;
   int rows_loaded = 0;
   while (run) {
     RdKafka::Message* msg = consumer->consume(10000);
     if (msg->err() == RdKafka::ERR_NO_ERROR) {
       if (!use_ccb) {
         bool added =
             msg_consume(msg, row_loader, copy_params, transformations, remove_quotes);
         if (added) {
           recv_rows++;
           if (recv_rows == copy_params.batch_size) {
             recv_rows = 0;
             row_loader.do_load(rows_loaded, skipped, copy_params);
             // make sure we now commit that we are up to here to cover the mesages we just
             // loaded
             consumer->commitSync();
           }
         } else {
           // LOG(ERROR) << " messsage was skipped ";
           skipped++;
         }
       }
     }
     delete msg;
   }
 
   /*
    * Stop consumer
    */
   consumer->close();
   delete consumer;
 
   LOG(INFO) << "Consumed " << msg_cnt << " messages (" << msg_bytes << " bytes)";
   LOG(FATAL) << "Consumer shut down, probably due to an error please review logs";
 };

Here is the call graph for this function:

Here is the caller graph for this function:

int main	(	int	argc,
		char **	argv
	)

Definition at line 476 of file KafkaImporter.cpp.

References BINARY, BINARY_SSL, logger::LogOptions::get_options(), HTTP, HTTPS, logger::init(), kafka_insert(), logger::LogOptions::max_files_, anonymous_namespace{Utm.h}::n, print_error_data, print_transformation, foreign_storage::anonymous_namespace{LogFileBufferParser.cpp}::remove_quotes(), run_benchmark_import::required, and run.

                                 {
   std::string server_host("localhost");  // default to localhost
   int port = 6274;                       // default port number
   bool http = false;
   bool https = false;
   bool skip_host_verify = false;
   std::string ca_cert_name{""};
   std::string table_name;
   std::string db_name;
   std::string user_name;
   std::string passwd;
   std::string group_id;
   std::string topic;
   std::string brokers;
   std::string delim_str(","), nulls("\\N"), line_delim_str("\n"), quoted("false");
   size_t batch_size = 10000;
   size_t retry_count = 10;
   size_t retry_wait = 5;
   bool remove_quotes = false;
   std::vector<std::string> xforms;
   std::map<std::string,
            std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>>
       transformations;
   ThriftConnectionType conn_type;
 
   namespace po = boost::program_options;
 
   po::options_description desc("Options");
   desc.add_options()("help,h", "Print help messages ");
   desc.add_options()(
       "table", po::value<std::string>(&table_name)->required(), "Table Name");
   desc.add_options()(
       "database", po::value<std::string>(&db_name)->required(), "Database Name");
   desc.add_options()(
       "user,u", po::value<std::string>(&user_name)->required(), "User Name");
   desc.add_options()(
       "passwd,p", po::value<std::string>(&passwd)->required(), "User Password");
   desc.add_options()("host",
                      po::value<std::string>(&server_host)->default_value(server_host),
                      "HeavyDB Server Hostname");
   desc.add_options()(
       "port", po::value<int>(&port)->default_value(port), "HeavyDB Server Port Number");
   desc.add_options()("http",
                      po::bool_switch(&http)->default_value(http)->implicit_value(true),
                      "Use HTTP transport");
   desc.add_options()("https",
                      po::bool_switch(&https)->default_value(https)->implicit_value(true),
                      "Use HTTPS transport");
   desc.add_options()("skip-verify",
                      po::bool_switch(&skip_host_verify)
                          ->default_value(skip_host_verify)
                          ->implicit_value(true),
                      "Don't verify SSL certificate validity");
   desc.add_options()(
       "ca-cert",
       po::value<std::string>(&ca_cert_name)->default_value(ca_cert_name),
       "Path to trusted server certificate. Initiates an encrypted connection");
   desc.add_options()("delim",
                      po::value<std::string>(&delim_str)->default_value(delim_str),
                      "Field delimiter");
   desc.add_options()("null", po::value<std::string>(&nulls), "NULL string");
   desc.add_options()("line", po::value<std::string>(&line_delim_str), "Line delimiter");
   desc.add_options()(
       "quoted",
       po::value<std::string>(&quoted),
       "Whether the source contains quoted fields (true/false, default false)");
   desc.add_options()("batch",
                      po::value<size_t>(&batch_size)->default_value(batch_size),
                      "Insert batch size");
   desc.add_options()("retry_count",
                      po::value<size_t>(&retry_count)->default_value(retry_count),
                      "Number of time to retry an insert");
   desc.add_options()("retry_wait",
                      po::value<size_t>(&retry_wait)->default_value(retry_wait),
                      "wait in secs between retries");
   desc.add_options()("transform,t",
                      po::value<std::vector<std::string>>(&xforms)->multitoken(),
                      "Column Transformations");
   desc.add_options()("print_error", "Print Error Rows");
   desc.add_options()("print_transform", "Print Transformations");
   desc.add_options()("topic",
                      po::value<std::string>(&topic)->required(),
                      "Kafka topic to consume from ");
   desc.add_options()("group-id",
                      po::value<std::string>(&group_id)->required(),
                      "Group id this consumer is part of");
   desc.add_options()("brokers",
                      po::value<std::string>(&brokers)->required(),
                      "list of kafka brokers for topic");
 
   po::positional_options_description positionalOptions;
   positionalOptions.add("table", 1);
   positionalOptions.add("database", 1);
 
   logger::LogOptions log_options(argv[0]);
   log_options.max_files_ = 0;  // stderr only
   desc.add(log_options.get_options());
 
   po::variables_map vm;
 
   try {
     po::store(po::command_line_parser(argc, argv)
                   .options(desc)
                   .positional(positionalOptions)
                   .run(),
               vm);
     if (vm.count("help")) {
       std::cout << "Usage: <table name> <database name> {-u|--user} <user> {-p|--passwd} "
                    "<password> [{--host} "
                    "<hostname>][--port <port number>][--delim <delimiter>][--null <null "
                    "string>][--line <line "
                    "delimiter>][--batch <batch size>][{-t|--transform} transformation "
                    "[--quoted <true|false>] "
                    "...][--retry_count <num_of_retries>] [--retry_wait <wait in "
                    "secs>][--print_error][--print_transform]\n\n";
       std::cout << desc << std::endl;
       return 0;
     }
     if (vm.count("print_error")) {
       print_error_data = true;
     }
     if (vm.count("print_transform")) {
       print_transformation = true;
     }
 
     po::notify(vm);
   } catch (boost::program_options::error& e) {
     std::cerr << "Usage Error: " << e.what() << std::endl;
     return 1;
   }
 
   logger::init(log_options);
 
   if (http) {
     conn_type = ThriftConnectionType::HTTP;
   } else if (https) {
     conn_type = ThriftConnectionType::HTTPS;
   } else if (!ca_cert_name.empty()) {
     conn_type = ThriftConnectionType::BINARY_SSL;
   } else {
     conn_type = ThriftConnectionType::BINARY;
   }
 
   char delim = delim_str[0];
   if (delim == '\\') {
     if (delim_str.size() < 2 ||
         (delim_str[1] != 'x' && delim_str[1] != 't' && delim_str[1] != 'n')) {
       std::cerr << "Incorrect delimiter string: " << delim_str << std::endl;
       return 1;
     }
     if (delim_str[1] == 't') {
       delim = '\t';
     } else if (delim_str[1] == 'n') {
       delim = '\n';
     } else {
       std::string d(delim_str);
       d[0] = '0';
       delim = (char)std::stoi(d, nullptr, 16);
     }
   }
   if (isprint(delim)) {
     std::cout << "Field Delimiter: " << delim << std::endl;
   } else if (delim == '\t') {
     std::cout << "Field Delimiter: "
               << "\\t" << std::endl;
   } else if (delim == '\n') {
     std::cout << "Field Delimiter: "
               << "\\n"
               << std::endl;
   } else {
     std::cout << "Field Delimiter: \\x" << std::hex << (int)delim << std::endl;
   }
   char line_delim = line_delim_str[0];
   if (line_delim == '\\') {
     if (line_delim_str.size() < 2 ||
         (line_delim_str[1] != 'x' && line_delim_str[1] != 't' &&
          line_delim_str[1] != 'n')) {
       std::cerr << "Incorrect delimiter string: " << line_delim_str << std::endl;
       return 1;
     }
     if (line_delim_str[1] == 't') {
       line_delim = '\t';
     } else if (line_delim_str[1] == 'n') {
       line_delim = '\n';
     } else {
       std::string d(line_delim_str);
       d[0] = '0';
       line_delim = (char)std::stoi(d, nullptr, 16);
     }
   }
   if (isprint(line_delim)) {
     std::cout << "Line Delimiter: " << line_delim << std::endl;
   } else if (line_delim == '\t') {
     std::cout << "Line Delimiter: "
               << "\\t" << std::endl;
   } else if (line_delim == '\n') {
     std::cout << "Line Delimiter: "
               << "\\n"
               << std::endl;
   } else {
     std::cout << "Line Delimiter: \\x" << std::hex << (int)line_delim << std::endl;
   }
   std::cout << "Null String: " << nulls << std::endl;
   std::cout << "Insert Batch Size: " << std::dec << batch_size << std::endl;
 
   if (quoted == "true") {
     remove_quotes = true;
   }
 
   for (auto& t : xforms) {
     auto n = t.find_first_of(':');
     if (n == std::string::npos) {
       std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/"
                 << std::endl;
       return 1;
     }
     std::string col_name = t.substr(0, n);
     if (t.size() < n + 3 || t[n + 1] != 's' || t[n + 2] != '/') {
       std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/"
                 << std::endl;
       return 1;
     }
     auto n1 = n + 3;
     auto n2 = t.find_first_of('/', n1);
     if (n2 == std::string::npos) {
       std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/"
                 << std::endl;
       return 1;
     }
     std::string regex_str = t.substr(n1, n2 - n1);
     n1 = n2 + 1;
     n2 = t.find_first_of('/', n1);
     if (n2 == std::string::npos) {
       std::cerr << "Transformation format: <column name>:s/<regex pattern>/<fmt string>/"
                 << std::endl;
       return 1;
     }
     std::string fmt_str = t.substr(n1, n2 - n1);
     std::cout << "transform " << col_name << ": s/" << regex_str << "/" << fmt_str << "/"
               << std::endl;
     transformations[col_name] =
         std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>(
             std::unique_ptr<boost::regex>(new boost::regex(regex_str)),
             std::unique_ptr<std::string>(new std::string(fmt_str)));
   }
 
   import_export::CopyParams copy_params(
       delim, nulls, line_delim, batch_size, retry_count, retry_wait);
   RowToColumnLoader row_loader(
       ThriftClientConnection(
           server_host, port, conn_type, skip_host_verify, ca_cert_name, ca_cert_name),
       user_name,
       passwd,
       db_name,
       table_name);
 
   kafka_insert(
       row_loader, transformations, copy_params, remove_quotes, group_id, topic, brokers);
   return 0;
 }

Here is the call graph for this function:

bool msg_consume	(	RdKafka::Message *	message,
		RowToColumnLoader &	row_loader,
		import_export::CopyParams	copy_params,
		const std::map< std::string, std::pair< std::unique_ptr< boost::regex >, std::unique_ptr< std::string >>> &	transformations,
		const bool	remove_quotes
	)

Definition at line 84 of file KafkaImporter.cpp.

References RowToColumnLoader::convert_string_to_column(), import_export::CopyParams::delimiter, eof_cnt, logger::ERROR, exit_eof, field(), RowToColumnLoader::get_row_descriptor(), import_export::CopyParams::line_delim, LOG, MAX_FIELD_LEN, msg_bytes, msg_cnt, import_export::CopyParams::null_str, partition_cnt, print_error_data, RowToColumnLoader::print_row_with_delim(), print_transformation, run, and VLOG.

Referenced by kafka_insert().

                                            {
   switch (message->err()) {
     case RdKafka::ERR__TIMED_OUT:
       VLOG(1) << " Timed out";
       break;
 
     case RdKafka::ERR_NO_ERROR: { /* Real message */
       msg_cnt++;
       msg_bytes += message->len();
       VLOG(1) << "Read msg at offset " << message->offset();
       RdKafka::MessageTimestamp ts;
       ts = message->timestamp();
       if (ts.type != RdKafka::MessageTimestamp::MSG_TIMESTAMP_NOT_AVAILABLE) {
         std::string tsname = "?";
         if (ts.type == RdKafka::MessageTimestamp::MSG_TIMESTAMP_CREATE_TIME) {
           tsname = "create time";
         } else if (ts.type == RdKafka::MessageTimestamp::MSG_TIMESTAMP_LOG_APPEND_TIME) {
           tsname = "log append time";
         }
         VLOG(1) << "Timestamp: " << tsname << " " << ts.timestamp << std::endl;
       }
 
       std::vector<char> buffer(message->len() + 2);
       sprintf(buffer.data(),
               "%.*s\n",
               static_cast<int>(message->len()),
               static_cast<const char*>(message->payload()));
       VLOG(1) << "Full Message received is :'" << buffer.data() << "'";
 
       char field[MAX_FIELD_LEN];
       size_t field_i = 0;
 
       bool backEscape = false;
 
       auto row_desc = row_loader.get_row_descriptor();
 
       std::vector<
           const std::pair<std::unique_ptr<boost::regex>, std::unique_ptr<std::string>>*>
           xforms(row_desc.size());
       for (size_t i = 0; i < row_desc.size(); i++) {
         auto it = transformations.find(row_desc[i].col_name);
         if (it != transformations.end()) {
           xforms[i] = &(it->second);
         } else {
           xforms[i] = nullptr;
         }
       }
 
       std::vector<TStringValue>
           row;  // used to store each row as we move through the stream
 
       for (auto iit : buffer) {
         if (iit == copy_params.delimiter || iit == copy_params.line_delim) {
           bool end_of_field = (iit == copy_params.delimiter);
           bool end_of_row;
           if (end_of_field) {
             end_of_row = false;
           } else {
             end_of_row = (row_desc[row.size()].col_type.type != TDatumType::STR) ||
                          (row.size() == row_desc.size() - 1);
             if (!end_of_row) {
               size_t l = copy_params.null_str.size();
               if (field_i >= l &&
                   strncmp(field + field_i - l, copy_params.null_str.c_str(), l) == 0) {
                 end_of_row = true;
               }
             }
           }
           if (!end_of_field && !end_of_row) {
             // not enough columns yet and it is a string column
             // treat the line delimiter as part of the string
             field[field_i++] = iit;
           } else {
             field[field_i] = '\0';
             field_i = 0;
             TStringValue ts;
             ts.str_val = std::string(field);
             ts.is_null = (ts.str_val.empty() || ts.str_val == copy_params.null_str);
             auto xform = row.size() < row_desc.size() ? xforms[row.size()] : nullptr;
             if (!ts.is_null && xform != nullptr) {
               if (print_transformation) {
                 std::cout << "\ntransforming\n" << ts.str_val << "\nto\n";
               }
               ts.str_val =
                   boost::regex_replace(ts.str_val, *xform->first, *xform->second);
               if (ts.str_val.empty()) {
                 ts.is_null = true;
               }
               if (print_transformation) {
                 std::cout << ts.str_val << std::endl;
               }
             }
 
             row.push_back(ts);  // add column value to row
             if (end_of_row || (row.size() > row_desc.size())) {
               break;  // found row
             }
           }
         } else {
           if (iit == '\\') {
             backEscape = true;
           } else if (backEscape || !remove_quotes || iit != '\"') {
             field[field_i++] = iit;
             backEscape = false;
           }
           // else if unescaped double-quote, continue without adding the
           // character to the field string.
         }
         if (field_i >= MAX_FIELD_LEN) {
           field[MAX_FIELD_LEN - 1] = '\0';
           std::cerr << "String too long for buffer." << std::endl;
           if (print_error_data) {
             std::cerr << field << std::endl;
           }
           field_i = 0;
           break;
         }
       }
       if (row.size() == row_desc.size()) {
         // add the new data in the column format
         bool record_loaded = row_loader.convert_string_to_column(row, copy_params);
         if (!record_loaded) {
           // record could not be parsed correctly consider it skipped
           return false;
         } else {
           return true;
         }
       } else {
         if (print_error_data) {
           std::cerr << "Incorrect number of columns for row: ";
           std::cerr << row_loader.print_row_with_delim(row, copy_params) << std::endl;
           return false;
         }
       }
     }
 
     case RdKafka::ERR__PARTITION_EOF:
       /* Last message */
       if (exit_eof && ++eof_cnt == partition_cnt) {
         LOG(ERROR) << "%% EOF reached for all " << partition_cnt << " partition(s)";
         run = false;
       }
       break;
 
     case RdKafka::ERR__UNKNOWN_TOPIC:
     case RdKafka::ERR__UNKNOWN_PARTITION:
       LOG(ERROR) << "Consume failed: " << message->errstr() << std::endl;
       run = false;
       break;
 
     default:
       /* Errors */
       LOG(ERROR) << "Consume failed: " << message->errstr();
       run = false;
   }
   return false;
 };