custreamz oauth callback for kafka (librdkafka)

conda/recipes/libcudf_kafka/meta.yaml

@@ -26,7 +26,7 @@ requirements:
     - cmake >=3.20.1
   host:
     - libcudf {{version}}
-    - librdkafka >=1.6.0,<1.7.0a0
+    - librdkafka >=1.7.0
   run:
     - {{ pin_compatible('librdkafka', max_pin='x.x') }} #TODO: librdkafka should be automatically included here by run_exports but is not
 

cpp/libcudf_kafka/CMakeLists.txt

@@ -44,6 +44,9 @@ rapids_cpm_init()
 include(cmake/thirdparty/get_cudf.cmake)
 include(cmake/thirdparty/get_rdkafka.cmake)
 
+# Locate Python Development headers
+find_package(Python3 COMPONENTS Interpreter Development)
+
 # # GTests if enabled
 if (BUILD_TESTS)
     # GoogleTest
@@ -57,7 +60,8 @@ endif()
 ###################################################################################################
 # - library target --------------------------------------------------------------------------------
 add_library(cudf_kafka SHARED
-    src/kafka_consumer.cpp)
+    src/kafka_consumer.cpp
+    src/kafka_callback.cpp)
 
 ###################################################################################################
 # - include paths ---------------------------------------------------------------------------------
@@ -68,7 +72,7 @@ target_include_directories(cudf_kafka
 
 ###################################################################################################
 # - library paths ---------------------------------------------------------------------------------
-target_link_libraries(cudf_kafka PUBLIC cudf::cudf RDKAFKA::RDKAFKA)
+target_link_libraries(cudf_kafka PUBLIC cudf::cudf RDKAFKA::RDKAFKA Python3::Python)
 
 set_target_properties(cudf_kafka
     PROPERTIES BUILD_RPATH                         "\$ORIGIN"

cpp/libcudf_kafka/include/cudf_kafka/kafka_callback.hpp

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2021, 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
+
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#include <librdkafka/rdkafkacpp.h>
+#include <cudf/io/datasource.hpp>
+#include <map>
+#include <memory>
+#include <string>
+
+namespace cudf {
+namespace io {
+namespace external {
+namespace kafka {
+
+/**
+ * @brief Callback to retrieve OAuth token from external source. Invoked when
+ * token refresh is required.
+ */
+class PythonOAuthRefreshCb : public RdKafka::OAuthBearerTokenRefreshCb {
+ public:
+  PythonOAuthRefreshCb(PyObject* callback, PyObject* args);
+
+  void oauthbearer_token_refresh_cb(RdKafka::Handle* handle, const std::string& oauthbearer_config);
+
+ private:
+  PyObject* callback;
+  PyObject* args;
+};
+
+}  // namespace kafka
+}  // namespace external
+}  // namespace io
+}  // namespace cudf

cpp/libcudf_kafka/include/cudf_kafka/kafka_consumer.hpp

@@ -15,13 +15,17 @@
  */
 #pragma once
 
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
 #include <librdkafka/rdkafkacpp.h>
 #include <algorithm>
 #include <chrono>
 #include <cudf/io/datasource.hpp>
 #include <map>
 #include <memory>
 #include <string>
+#include "kafka_callback.hpp"
 
 namespace cudf {
 namespace io {
@@ -49,7 +53,7 @@ class kafka_consumer : public cudf::io::datasource {
    * @param configs key/value pairs of librdkafka configurations that will be
    *                passed to the librdkafka client
    */
-  kafka_consumer(std::map<std::string, std::string> const& configs);
+  kafka_consumer(PyObject* configs);
 
   /**
    * @brief Instantiate a Kafka consumer object. Documentation for librdkafka configurations can be
@@ -66,7 +70,7 @@ class kafka_consumer : public cudf::io::datasource {
    * before batch_timeout, a smaller subset will be returned
    * @param delimiter optional delimiter to insert into the output between kafka messages, Ex: "\n"
    */
-  kafka_consumer(std::map<std::string, std::string> const& configs,
+  kafka_consumer(PyObject* configs,
                  std::string const& topic_name,
                  int partition,
                  int64_t start_offset,
@@ -178,6 +182,12 @@ class kafka_consumer : public cudf::io::datasource {
   std::unique_ptr<RdKafka::Conf> kafka_conf;  // RDKafka configuration object
   std::unique_ptr<RdKafka::KafkaConsumer> consumer;
 
+  // Configurations that can be Python callables. Anything else is expected to be a str
+  const std::vector<std::string> callableConfigs{"oauth_cb"};
+
+  // The Python configuration dict that was used to create this instance
+  PyObject* conf_dict;
+
   std::string topic_name;
   int partition;
   int64_t start_offset;
@@ -193,6 +203,8 @@ class kafka_consumer : public cudf::io::datasource {
                                                                 int partition,
                                                                 int64_t offset);
 
+  void build_validate_configs(PyObject* python_config_dict);
+
   /**
    * Convenience method for getting "now()" in Kafka's standard format
    */

cpp/libcudf_kafka/src/kafka_callback.cpp

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2021, 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.
+ */
+#include "cudf_kafka/kafka_callback.hpp"
+
+namespace cudf {
+namespace io {
+namespace external {
+namespace kafka {
+
+PythonOAuthRefreshCb::PythonOAuthRefreshCb(PyObject* callback, PyObject* args)
+  : callback(callback), args(args){};
+
+void PythonOAuthRefreshCb::oauthbearer_token_refresh_cb(RdKafka::Handle* handle,
+                                                        const std::string& oauthbearer_config)
+{
+  CUDF_EXPECTS(PyCallable_Check(callback), "A Python callable is required");
+
+  // Make sure that we own the GIL
+  PyGILState_STATE state = PyGILState_Ensure();
+  PyObject* result       = PyObject_CallObject(callback, args);
+  Py_XINCREF(result);
+
+  // Set the token in the Kafka context
+  if (result) {
+    CUDF_EXPECTS(PyDict_Check(result),
+                 "cudf_kafka requires a Dictionary response from the Python OAuthRefreshCb with "
+                 "dictionary keys (token, token_lifetime_ms, principal, extensions)");
+
+    // Ensure that expected keys are present from the Python callback response.
+    std::string token = PyUnicode_AsUTF8(PyDict_GetItemString(result, "token"));
+    int64_t token_lifetime_ms =
+      PyLong_AsLongLong(PyDict_GetItemString(result, "token_lifetime_ms"));
+    std::string principal = PyUnicode_AsUTF8(PyDict_GetItemString(result, "principal"));
+    std::list<std::string> extensions;
+    std::string errstr;
+
+    handle->oauthbearer_set_token(token, token_lifetime_ms, principal, extensions, errstr);
+  } else {
+    handle->oauthbearer_set_token_failure("");
+  }
+
+  Py_XDECREF(result);
+  PyGILState_Release(state);
+}
+
+}  // namespace kafka
+}  // namespace external
+}  // namespace io
+}  // namespace cudf

cpp/libcudf_kafka/src/kafka_consumer.cpp

@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 #include "cudf_kafka/kafka_consumer.hpp"
 #include <librdkafka/rdkafkacpp.h>
 #include <chrono>
@@ -24,27 +23,17 @@ namespace io {
 namespace external {
 namespace kafka {
 
-kafka_consumer::kafka_consumer(std::map<std::string, std::string> const& configs)
-  : kafka_conf(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL))
+kafka_consumer::kafka_consumer(PyObject* confdict)
+  : conf_dict(confdict), kafka_conf(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL))
 {
-  for (auto const& key_value : configs) {
-    std::string error_string;
-    CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK ==
-                   kafka_conf->set(key_value.first, key_value.second, error_string),
-                 "Invalid Kafka configuration");
-  }
-
-  // Kafka 0.9 > requires group.id in the configuration
-  std::string conf_val;
-  CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK == kafka_conf->get("group.id", conf_val),
-               "Kafka group.id must be configured");
+  build_validate_configs(confdict);
 
   std::string errstr;
   consumer = std::unique_ptr<RdKafka::KafkaConsumer>(
     RdKafka::KafkaConsumer::create(kafka_conf.get(), errstr));
 }
 
-kafka_consumer::kafka_consumer(std::map<std::string, std::string> const& configs,
+kafka_consumer::kafka_consumer(PyObject* confdict,
                                std::string const& topic_name,
                                int partition,
                                int64_t start_offset,
@@ -56,21 +45,11 @@ kafka_consumer::kafka_consumer(std::map<std::string, std::string> const& configs
     start_offset(start_offset),
     end_offset(end_offset),
     batch_timeout(batch_timeout),
-    delimiter(delimiter)
+    delimiter(delimiter),
+    conf_dict(confdict),
+    kafka_conf(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL))
 {
-  kafka_conf = std::unique_ptr<RdKafka::Conf>(RdKafka::Conf::create(RdKafka::Conf::CONF_GLOBAL));
-
-  for (auto const& key_value : configs) {
-    std::string error_string;
-    CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK ==
-                   kafka_conf->set(key_value.first, key_value.second, error_string),
-                 "Invalid Kafka configuration");
-  }
-
-  // Kafka 0.9 > requires group.id in the configuration
-  std::string conf_val;
-  CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK == kafka_conf->get("group.id", conf_val),
-               "Kafka group.id must be configured");
+  build_validate_configs(confdict);
 
   std::string errstr;
   consumer = std::unique_ptr<RdKafka::KafkaConsumer>(
@@ -81,6 +60,42 @@ kafka_consumer::kafka_consumer(std::map<std::string, std::string> const& configs
   consume_to_buffer();
 }
 
+/**
+ * @brief Builds and validates Kafka C++ configuration object from Python values
+ *
+ * @param kafka_configs
+ *  Python Dict of configuration values and possibly callables for callbacks
+ */
+void kafka_consumer::build_validate_configs(PyObject* python_config_dict)
+{
+  Py_ssize_t pos = 0;
+  PyObject *ko, *vo;
+
+  while (PyDict_Next(python_config_dict, &pos, &ko, &vo)) {
+    CUDF_EXPECTS(PyUnicode_Check(ko), "expected kafka configuration property name as type string");
+    std::string key(PyUnicode_AsUTF8(ko));
+    std::string valueType(Py_TYPE(vo)->tp_name);
+
+    std::string error_string;
+    if (std::find(callableConfigs.begin(), callableConfigs.end(), key) != callableConfigs.end()) {
+      // Properly configure the callable. This is a Python callback for oauth processing
+      PythonOAuthRefreshCb cb(vo, NULL);
+      kafka_conf->set("oauthbearer_token_refresh_cb", &cb, error_string);
+    } else {
+      CUDF_EXPECTS(valueType.compare("str") == 0,
+                   "Only string values are supported for this configuration");
+      CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK ==
+                     kafka_conf->set(key, PyUnicode_AsUTF8(vo), error_string),
+                   "Invalid Kafka configuration provided");
+    }
+  }
+
+  // Kafka 0.9 > requires group.id in the configuration
+  std::string conf_val;
+  CUDF_EXPECTS(RdKafka::Conf::ConfResult::CONF_OK == kafka_conf->get("group.id", conf_val),
+               "Kafka group.id must be configured");
+}
+
 std::unique_ptr<cudf::io::datasource::buffer> kafka_consumer::host_read(size_t offset, size_t size)
 {
   if (offset > buffer.size()) { return 0; }

cpp/libcudf_kafka/tests/CMakeLists.txt

@@ -14,14 +14,17 @@
 # limitations under the License.
 #=============================================================================
 
+# Locate Python Development headers
+find_package(Python3 COMPONENTS Interpreter Development)
+
 ###################################################################################################
 # - compiler function -----------------------------------------------------------------------------
 
 function(ConfigureTest test_name )
     add_executable(${test_name} ${ARGN})
     set_target_properties(${test_name}
         PROPERTIES RUNTIME_OUTPUT_DIRECTORY "$<BUILD_INTERFACE:${CUDA_KAFKA_BINARY_DIR}/gtests>")
-    target_link_libraries(${test_name} PRIVATE GTest::gmock_main GTest::gtest_main cudf_kafka)
+    target_link_libraries(${test_name} PRIVATE GTest::gmock_main GTest::gtest_main cudf_kafka Python3::Python)
 
     add_test(NAME ${test_name} COMMAND ${test_name})
 endfunction()

cpp/libcudf_kafka/tests/kafka_consumer_tests.cpp

@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
 #include <gtest/gtest.h>
 #include <map>
 #include <memory>
@@ -31,8 +34,7 @@ struct KafkaDatasourceTest : public ::testing::Test {
 TEST_F(KafkaDatasourceTest, MissingGroupID)
 {
   // group.id is a required configuration.
-  std::map<std::string, std::string> kafka_configs;
-  kafka_configs.insert({"bootstrap.servers", "localhost:9092"});
+  PyObject* kafka_configs = Py_BuildValue("{s:s}", "bootstrap.servers", "localhost:9092");
 
   EXPECT_THROW(kafka::kafka_consumer kc(kafka_configs, "csv-topic", 0, 0, 3, 5000, "\n"),
                cudf::logic_error);
@@ -41,16 +43,13 @@ TEST_F(KafkaDatasourceTest, MissingGroupID)
 TEST_F(KafkaDatasourceTest, InvalidConfigValues)
 {
   // Give a made up configuration value
-  std::map<std::string, std::string> kafka_configs;
-  kafka_configs.insert({"completely_made_up_config", "wrong"});
+  PyObject* kafka_configs = Py_BuildValue("{s:s}", "completely_made_up_config", "wrong");
 
   EXPECT_THROW(kafka::kafka_consumer kc(kafka_configs, "csv-topic", 0, 0, 3, 5000, "\n"),
                cudf::logic_error);
 
-  kafka_configs.clear();
-
   // Give a good config property with a bad value
-  kafka_configs.insert({"message.max.bytes", "this should be a number not text"});
+  kafka_configs = Py_BuildValue("{s:s}", "message.max.bytes", "his should be a number not text");
   EXPECT_THROW(kafka::kafka_consumer kc(kafka_configs, "csv-topic", 0, 0, 3, 5000, "\n"),
                cudf::logic_error);
 }

python/cudf_kafka/cudf_kafka/_lib/kafka.pxd

@@ -16,9 +16,9 @@ cdef extern from "kafka_consumer.hpp" \
 
     cpdef cppclass kafka_consumer:
 
-        kafka_consumer(map[string, string] configs) except +
+        kafka_consumer(object configs) except +
 
-        kafka_consumer(map[string, string] configs,
+        kafka_consumer(object configs,
                        string topic_name,
                        int32_t partition,
                        int64_t start_offset,