set token failure on failure

bruce-szalwinski-he · May 5, 2024 · 6cb6aa3 · 6cb6aa3
1 parent 6244d01
commit 6cb6aa3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git a/examples/rdkafka/oauth_token_refresher.rb b/examples/rdkafka/oauth_token_refresher.rb
@@ -6,15 +6,21 @@
 class OAuthTokenRefresher
   def refresh_token(client_name)
     print "refreshing token\n"
+    client = Producer.from_name(client_name)
     signer = AwsMskIamSaslSigner::MSKTokenProvider.new(region: ENV.fetch("AWS_REGION", nil))
     token = signer.generate_auth_token
 
-    client = Producer.from_name(client_name)
-    client.oauthbearer_set_token(
-      token: token.token,
-      lifetime_ms: token.expiration_time_ms,
-      principal_name: "kafka-cluster"
-    )
+    if token
+      client.oauthbearer_set_token(
+        token: token.token,
+        lifetime_ms: token.expiration_time_ms,
+        principal_name: "kafka-cluster"
+      )
+    else
+      client.oauthbearer_set_token_failure(
+        "Failed to generate token."
+      )
+    end
   end
 end
 

diff --git a/examples/waterdrop/oauth_token_refresher.rb b/examples/waterdrop/oauth_token_refresher.rb
@@ -4,13 +4,20 @@
 
 class OAuthTokenRefresher
   def on_oauthbearer_token_refresh(event)
+    print "refreshing token\n"
+
     signer = AwsMskIamSaslSigner::MSKTokenProvider.new(region: ENV.fetch("AWS_REGION", nil))
     token = signer.generate_auth_token
-
-    event[:bearer].oauthbearer_set_token(
-      token: token.token,
-      lifetime_ms: token.expiration_time_ms,
-      principal_name: "kafka-cluster"
-    )
+    if token
+      event[:bearer].oauthbearer_set_token(
+        token: token.token,
+        lifetime_ms: token.expiration_time_ms,
+        principal_name: "kafka-cluster"
+      )
+    else
+      event[:bearer].oauthbearer_set_token_failure(
+        "Failed to generate token."
+      )
+    end
   end
 end