Refresh OAuth token in the background

So that OAUTHBEARER authentication works even when poll isn't being called regularly (e.g., in a client just used for fetching metadata).
MaterializeInc · Aug 4, 2024 · 8a13235 · 8a13235
1 parent 1e27332
commit 8a13235
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/src/client.rs b/src/client.rs
@@ -28,7 +28,7 @@ use rdkafka_sys::types::*;
 
 use crate::config::{ClientConfig, NativeClientConfig, RDKafkaLogLevel};
 use crate::consumer::RebalanceProtocol;
-use crate::error::{IsError, KafkaError, KafkaResult};
+use crate::error::{IsError, KafkaError, KafkaResult, RDKafkaError};
 use crate::groups::GroupList;
 use crate::log::{debug, error, info, trace, warn};
 use crate::metadata::Metadata;
@@ -247,7 +247,8 @@ impl<C: ClientContext + 'static> Client<C> {
                 rdsys::rd_kafka_conf_set_oauthbearer_token_refresh_cb(
                     native_config.ptr(),
                     Some(native_oauth_refresh_cb::<C>),
-                )
+                );
+                rdkafka_sys::rd_kafka_conf_enable_sasl_queue(native_config.ptr(), 1);
             };
         }
 
@@ -268,6 +269,15 @@ impl<C: ClientContext + 'static> Client<C> {
 
         unsafe { rdsys::rd_kafka_set_log_level(client_ptr, config.log_level as i32) };
 
+        if C::ENABLE_REFRESH_OAUTH_TOKEN {
+            let ret = unsafe {
+                RDKafkaError::from_ptr(rdsys::rd_kafka_sasl_background_callbacks_enable(client_ptr))
+            };
+            if ret.is_error() {
+                return Err(KafkaError::OAuthConfig(ret));
+            }
+        }
+
         Ok(Client {
             native: unsafe { NativeClient::from_ptr(client_ptr) },
             context,

diff --git a/src/error.rs b/src/error.rs
@@ -163,6 +163,8 @@ pub enum KafkaError {
     NoMessageReceived,
     /// Unexpected null pointer
     Nul(ffi::NulError),
+    /// OAuth configuration failed.
+    OAuthConfig(RDKafkaError),
     /// Offset fetch failed.
     OffsetFetch(RDKafkaErrorCode),
     /// End of partition reached.
@@ -218,6 +220,7 @@ impl fmt::Debug for KafkaError {
                 write!(f, "No message received within the given poll interval")
             }
             KafkaError::Nul(_) => write!(f, "FFI null error"),
+            KafkaError::OAuthConfig(err) => write!(f, "KafkaError (OAuth config error: {})", err),
             KafkaError::OffsetFetch(err) => write!(f, "KafkaError (Offset fetch error: {})", err),
             KafkaError::PartitionEOF(part_n) => write!(f, "KafkaError (Partition EOF: {})", part_n),
             KafkaError::PauseResume(ref err) => {
@@ -259,6 +262,7 @@ impl fmt::Display for KafkaError {
                 write!(f, "No message received within the given poll interval")
             }
             KafkaError::Nul(_) => write!(f, "FFI nul error"),
+            KafkaError::OAuthConfig(err) => write!(f, "OAuth config error: {}", err),
             KafkaError::OffsetFetch(err) => write!(f, "Offset fetch error: {}", err),
             KafkaError::PartitionEOF(part_n) => write!(f, "Partition EOF: {}", part_n),
             KafkaError::PauseResume(ref err) => write!(f, "Pause/resume error: {}", err),
@@ -288,6 +292,7 @@ impl Error for KafkaError {
             KafkaError::MetadataFetch(err) => Some(err),
             KafkaError::NoMessageReceived => None,
             KafkaError::Nul(_) => None,
+            KafkaError::OAuthConfig(err) => Some(err),
             KafkaError::OffsetFetch(err) => Some(err),
             KafkaError::PartitionEOF(_) => None,
             KafkaError::PauseResume(_) => None,
@@ -325,6 +330,7 @@ impl KafkaError {
             KafkaError::MetadataFetch(err) => Some(*err),
             KafkaError::NoMessageReceived => None,
             KafkaError::Nul(_) => None,
+            KafkaError::OAuthConfig(err) => Some(err.code()),
             KafkaError::OffsetFetch(err) => Some(*err),
             KafkaError::PartitionEOF(_) => None,
             KafkaError::PauseResume(_) => None,