Repository: dpkp/kafka-python
Pull request #2640 (2025-06-04 22:00:58)

Do not ignore metadata response for single topic with error
Improve behavior when single topic is unauthorized.


Patch:
diff --git a/kafka/cluster.py b/kafka/cluster.py
index d6ec82d..ded8c6f 100644
--- a/kafka/cluster.py
+++ b/kafka/cluster.py
@@ -245,13 +245,6 @@ class ClusterMetadata(object):
 
         Returns: None
         """
-        # In the common case where we ask for a single topic and get back an
-        # error, we should fail the future
-        if len(metadata.topics) == 1 and metadata.topics[0][0] != Errors.NoError.errno:
-            error_code, topic = metadata.topics[0][:2]
-            error = Errors.for_code(error_code)(topic)
-            return self.failed_update(error)
-
         if not metadata.brokers:
             log.warning("No broker metadata found in MetadataResponse -- ignoring.")
             return self.failed_update(Errors.MetadataEmptyBrokerList(metadata))
@@ -349,7 +342,15 @@ class ClusterMetadata(object):
         self._last_successful_refresh_ms = now
 
         if f:
-            f.success(self)
+            # In the common case where we ask for a single topic and get back an
+            # error, we should fail the future
+            if len(metadata.topics) == 1 and metadata.topics[0][0] != Errors.NoError.errno:
+                error_code, topic = metadata.topics[0][:2]
+                error = Errors.for_code(error_code)(topic)
+                f.failure(error)
+            else:
+                f.success(self)
+
         log.debug("Updated cluster metadata to %s", self)
 
         for listener in self._listeners:
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index 2416cc3..9401bd8 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -960,9 +960,11 @@ class KafkaProducer(object):
             future.add_both(lambda e, *args: e.set(), metadata_event)
             self._sender.wakeup()
             metadata_event.wait(timer.timeout_ms / 1000)
-            if not metadata_event.is_set():
+            if not future.is_done:
                 raise Errors.KafkaTimeoutError(
                     "Failed to update metadata after %.1f secs." % (max_wait_ms / 1000,))
+            elif future.failed() and not future.retriable():
+                raise future.exception
             elif topic in self._metadata.unauthorized_topics:
                 raise Errors.TopicAuthorizationFailedError(set([topic]))
             else:
diff --git a/kafka/sasl/msk.py b/kafka/sasl/msk.py
index db56b48..7ec0321 100644
--- a/kafka/sasl/msk.py
+++ b/kafka/sasl/msk.py
@@ -4,6 +4,7 @@ import datetime
 import hashlib
 import hmac
 import json
+import logging
 import string
 
 # needed for AWS_MSK_IAM authentication:
@@ -13,10 +14,14 @@ except ImportError:
     # no botocore available, will disable AWS_MSK_IAM mechanism
     BotoSession = None
 
+from kafka.errors import KafkaConfigurationError
 from kafka.sasl.abc import SaslMechanism
 from kafka.vendor.six.moves import urllib
 
 
+log = logging.getLogger(__name__)
+
+
 class SaslMechanismAwsMskIam(SaslMechanism):
     def __init__(self, **config):
         assert BotoSession is not None, 'AWS_MSK_IAM requires the "botocore" package'
@@ -27,22 +32,28 @@ class SaslMechanismAwsMskIam(SaslMechanism):
         self._is_done = False
         self._is_authenticated = False
 
-    def auth_bytes(self):
+    def _build_client(self):
         session = BotoSession()
         credentials = session.get_credentials().get_frozen_credentials()
-        client = AwsMskIamClient(
+        if not session.get_config_variable('region'):
+            raise KafkaConfigurationError('Unable to determine region for AWS MSK cluster. Is AWS_DEFAULT_REGION set?')
+        return AwsMskIamClient(
             host=self.host,
             access_key=credentials.access_key,
             secret_key=credentials.secret_key,
             region=session.get_config_variable('region'),
             token=credentials.token,
         )
+
+    def auth_bytes(self):
+        client = self._build_client()
+        log.debug("Generating auth token for MSK scope: %s", client._scope)
         return client.first_message()
 
     def receive(self, auth_bytes):
         self._is_done = True
         self._is_authenticated = auth_bytes != b''
-        self._auth = auth_bytes.deode('utf-8')
+        self._auth = auth_bytes.decode('utf-8')
 
     def is_done(self):
         return self._is_done


Test patch:
diff --git a/test/sasl/test_msk.py b/test/sasl/test_msk.py
index e9f1325..f3cc46c 100644
--- a/test/sasl/test_msk.py
+++ b/test/sasl/test_msk.py
@@ -2,7 +2,7 @@ import datetime
 import json
 import sys
 
-from kafka.sasl.msk import AwsMskIamClient
+from kafka.sasl.msk import AwsMskIamClient, SaslMechanismAwsMskIam
 
 try:
     from unittest import mock
@@ -69,3 +69,17 @@ def test_aws_msk_iam_client_temporary_credentials():
         'x-amz-security-token': 'XXXXX',
     }
     assert actual == expected
+
+
+def test_aws_msk_iam_sasl_mechanism():
+    with mock.patch('kafka.sasl.msk.BotoSession'):
+        sasl = SaslMechanismAwsMskIam(security_protocol='SASL_SSL', host='localhost')
+        with mock.patch.object(sasl, '_build_client', return_value=client_factory(token=None)):
+            assert sasl.auth_bytes() != b''
+            assert not sasl.is_done()
+            assert not sasl.is_authenticated()
+            sasl.receive(b'foo')
+            assert sasl._auth == 'foo'
+            assert sasl.is_done()
+            assert sasl.is_authenticated()
+            assert sasl.auth_details()
diff --git a/test/test_cluster.py b/test/test_cluster.py
index c57bd8f..730b27c 100644
--- a/test/test_cluster.py
+++ b/test/test_cluster.py
@@ -136,6 +136,21 @@ def test_metadata_v7():
     assert cluster._partitions['topic-1'][0].leader_epoch == 0
 
 
+def test_unauthorized_topic():
+    cluster = ClusterMetadata()
+    assert len(cluster.brokers()) == 0
+
+    cluster.update_metadata(MetadataResponse[0](
+        [(0, 'foo', 12), (1, 'bar', 34)],
+        [(29, 'unauthorized-topic', [])]))  # single topic w/ unauthorized error
+
+    # broker metadata should get updated
+    assert len(cluster.brokers()) == 2
+
+    # topic should be added to unauthorized list
+    assert 'unauthorized-topic' in cluster.unauthorized_topics
+
+
 def test_collect_hosts__happy_path():
     hosts = "127.0.0.1:1234,127.0.0.1"
     results = collect_hosts(hosts)