-
Notifications
You must be signed in to change notification settings - Fork 915
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
custreamz oauth callback for kafka (librdkafka) (#9486)
Previously it was impossible to use custreamz with oauth enabled Kafka brokers. This PR adds a feature so that the user can supply a Python function which is invoked to get the oauth token, from a http endpoint for example, and then supply that token to librdkafka to be used in both the initial connection to kafka and also subsequently as the token becomes stale. This closes #9410 Authors: - Jeremy Dyer (https://github.com/jdye64) - AJ Schmidt (https://github.com/ajschmidt8) Approvers: - Robert Maynard (https://github.com/robertmaynard) - Vyas Ramasubramani (https://github.com/vyasr) - Ram (Ramakrishna Prabhu) (https://github.com/rgsl888prabhu) - AJ Schmidt (https://github.com/ajschmidt8) URL: #9486
- Loading branch information
Showing
21 changed files
with
295 additions
and
77 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,4 @@ | ||
| # Copyright (c) 2020, NVIDIA CORPORATION. | ||
| # Copyright (c) 2020-2022, NVIDIA CORPORATION. | ||
|
|
||
| # This assumes the script is executed from the root of the repo directory | ||
| ./build.sh -v cudf_kafka |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,4 @@ | ||
| # Copyright (c) 2020, NVIDIA CORPORATION. | ||
| # Copyright (c) 2020-2022, NVIDIA CORPORATION. | ||
|
|
||
| # This assumes the script is executed from the root of the repo directory | ||
| ./build.sh -v custreamz |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,71 @@ | ||
| /* | ||
| * Copyright (c) 2021-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 <cudf/io/datasource.hpp> | ||
|
|
||
| #include <librdkafka/rdkafkacpp.h> | ||
|
|
||
| #include <map> | ||
| #include <memory> | ||
| #include <string> | ||
|
|
||
| namespace cudf { | ||
| namespace io { | ||
| namespace external { | ||
| namespace kafka { | ||
|
|
||
| /** | ||
| * @brief Python Callback function wrapper type used for Kafka OAuth events | ||
| * | ||
| * The KafkaConsumer calls the `kafka_oauth_callback_wrapper_type` when the existing | ||
| * oauth token is considered expired by the KafkaConsumer. Typically that | ||
| * means this will be invoked a single time when the KafkaConsumer is created | ||
| * to get the initial token and then intermediately as the token becomes | ||
| * expired. | ||
| * | ||
| * The callback function signature is: | ||
| * `std::map<std::string, std::string> kafka_oauth_callback_wrapper_type(void*)` | ||
| * | ||
| * The callback function returns a std::map<std::string, std::string>, | ||
| * where the std::map consists of the Oauth token and its | ||
| * linux epoch expiration time. Generally the token and expiration | ||
| * time is retrieved from an external service by the callback. | ||
| * Ex: [token, token_expiration_in_epoch] | ||
| */ | ||
| using kafka_oauth_callback_wrapper_type = std::map<std::string, std::string> (*)(void*); | ||
| using python_callable_type = void*; | ||
|
|
||
| /** | ||
| * @brief Callback to retrieve OAuth token from external source. Invoked when | ||
| * token refresh is required. | ||
| */ | ||
| class python_oauth_refresh_callback : public RdKafka::OAuthBearerTokenRefreshCb { | ||
| public: | ||
| python_oauth_refresh_callback(kafka_oauth_callback_wrapper_type callback_wrapper, | ||
| python_callable_type python_callable); | ||
|
|
||
| void oauthbearer_token_refresh_cb(RdKafka::Handle* handle, const std::string& oauthbearer_config); | ||
|
|
||
| private: | ||
| kafka_oauth_callback_wrapper_type callback_wrapper_; | ||
| python_callable_type python_callable_; | ||
| }; | ||
|
|
||
| } // namespace kafka | ||
| } // namespace external | ||
| } // namespace io | ||
| } // namespace cudf |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright (c) 2021-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. | ||
| */ | ||
| #include "cudf_kafka/kafka_callback.hpp" | ||
|
|
||
| #include <librdkafka/rdkafkacpp.h> | ||
|
|
||
| namespace cudf { | ||
| namespace io { | ||
| namespace external { | ||
| namespace kafka { | ||
|
|
||
| python_oauth_refresh_callback::python_oauth_refresh_callback( | ||
| kafka_oauth_callback_wrapper_type callback_wrapper, python_callable_type python_callable) | ||
| : callback_wrapper_(callback_wrapper), python_callable_(python_callable){}; | ||
|
|
||
| void python_oauth_refresh_callback::oauthbearer_token_refresh_cb( | ||
| RdKafka::Handle* handle, std::string const& oauthbearer_config) | ||
| { | ||
| std::map<std::string, std::string> resp = callback_wrapper_(python_callable_); | ||
|
|
||
| // Build parameters to pass to librdkafka | ||
| std::string token = resp["token"]; | ||
| int64_t token_lifetime_ms = std::stoll(resp["token_expiration_in_epoch"]); | ||
| std::list<std::string> extensions; // currently not supported | ||
| std::string errstr; | ||
| CUDF_EXPECTS( | ||
| RdKafka::ErrorCode::ERR_NO_ERROR == | ||
| handle->oauthbearer_set_token(token, token_lifetime_ms, "kafka", extensions, errstr), | ||
| "Error occurred while setting the oauthbearer token"); | ||
| } | ||
|
|
||
| } // namespace kafka | ||
| } // namespace external | ||
| } // namespace io | ||
| } // namespace cudf |
Oops, something went wrong.