kafka-python 2.0.1__py2.py3-none-any.whl → 2.0.3__py2.py3-none-any.whl

kafka/__init__.py

@@ -4,7 +4,7 @@ __title__ = 'kafka'
 from kafka.version import __version__
 __author__ = 'Dana Powers'
 __license__ = 'Apache License 2.0'
-__copyright__ = 'Copyright 2016 Dana Powers, David Arthur, and Contributors'
+__copyright__ = 'Copyright 2025 Dana Powers, David Arthur, and Contributors'
 
 # Set default logging handler to avoid "No handler found" warnings.
 import logging

kafka/admin/client.py

@@ -8,7 +8,10 @@ import socket
 from . import ConfigResourceType
 from kafka.vendor import six
 
+from kafka.admin.acl_resource import ACLOperation, ACLPermissionType, ACLFilter, ACL, ResourcePattern, ResourceType, \
+    ACLResourcePatternType
 from kafka.client_async import KafkaClient, selectors
+from kafka.coordinator.protocol import ConsumerProtocolMemberMetadata, ConsumerProtocolMemberAssignment, ConsumerProtocol
 import kafka.errors as Errors
 from kafka.errors import (
     IncompatibleBrokerVersion, KafkaConfigurationError, NotControllerError,
@@ -16,12 +19,13 @@ from kafka.errors import (
 from kafka.metrics import MetricConfig, Metrics
 from kafka.protocol.admin import (
     CreateTopicsRequest, DeleteTopicsRequest, DescribeConfigsRequest, AlterConfigsRequest, CreatePartitionsRequest,
-    ListGroupsRequest, DescribeGroupsRequest, DescribeAclsRequest, CreateAclsRequest, DeleteAclsRequest)
+    ListGroupsRequest, DescribeGroupsRequest, DescribeAclsRequest, CreateAclsRequest, DeleteAclsRequest,
+    DeleteGroupsRequest, DescribeLogDirsRequest
+)
 from kafka.protocol.commit import GroupCoordinatorRequest, OffsetFetchRequest
 from kafka.protocol.metadata import MetadataRequest
-from kafka.structs import TopicPartition, OffsetAndMetadata
-from kafka.admin.acl_resource import ACLOperation, ACLPermissionType, ACLFilter, ACL, ResourcePattern, ResourceType, \
-    ACLResourcePatternType
+from kafka.protocol.types import Array
+from kafka.structs import TopicPartition, OffsetAndMetadata, MemberInformation, GroupInformation
 from kafka.version import __version__
 
 
@@ -142,6 +146,7 @@ class KafkaAdminClient(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     """
     DEFAULT_CONFIG = {
@@ -182,6 +187,7 @@ class KafkaAdminClient(object):
         'metric_reporters': [],
         'metrics_num_samples': 2,
         'metrics_sample_window_ms': 30000,
+        'kafka_client': KafkaClient,
     }
 
     def __init__(self, **configs):
@@ -201,10 +207,12 @@ class KafkaAdminClient(object):
         reporters = [reporter() for reporter in self.config['metric_reporters']]
         self._metrics = Metrics(metric_config, reporters)
 
-        self._client = KafkaClient(metrics=self._metrics,
-                                   metric_group_prefix='admin',
-                                   **self.config)
-        self._client.check_version()
+        self._client = self.config['kafka_client'](
+            metrics=self._metrics,
+            metric_group_prefix='admin',
+            **self.config
+        )
+        self._client.check_version(timeout=(self.config['api_version_auto_timeout_ms'] / 1000))
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:
@@ -271,7 +279,7 @@ class KafkaAdminClient(object):
             response = future.value
             controller_id = response.controller_id
             # verify the controller is new enough to support our requests
-            controller_version = self._client.check_version(controller_id)
+            controller_version = self._client.check_version(controller_id, timeout=(self.config['api_version_auto_timeout_ms'] / 1000))
             if controller_version < (0, 10, 0):
                 raise IncompatibleBrokerVersion(
                     "The controller appears to be running Kafka {}. KafkaAdminClient requires brokers >= 0.10.0.0."
@@ -324,30 +332,37 @@ class KafkaAdminClient(object):
                 .format(response.API_VERSION))
         return response.coordinator_id
 
-    def _find_coordinator_id(self, group_id):
-        """Find the broker node_id of the coordinator of the given group.
+    def _find_coordinator_ids(self, group_ids):
+        """Find the broker node_ids of the coordinators of the given groups.
 
-        Sends a FindCoordinatorRequest message to the cluster. Will block until
-        the FindCoordinatorResponse is received. Any errors are immediately
-        raised.
+        Sends a FindCoordinatorRequest message to the cluster for each group_id.
+        Will block until the FindCoordinatorResponse is received for all groups.
+        Any errors are immediately raised.
 
-        :param group_id: The consumer group ID. This is typically the group
+        :param group_ids: A list of consumer group IDs. This is typically the group
             name as a string.
-        :return: The node_id of the broker that is the coordinator.
+        :return: A dict of {group_id: node_id} where node_id is the id of the
+            broker that is the coordinator for the corresponding group.
         """
-        # Note: Java may change how this is implemented in KAFKA-6791.
-        future = self._find_coordinator_id_send_request(group_id)
-        self._wait_for_futures([future])
-        response = future.value
-        return self._find_coordinator_id_process_response(response)
-
-    def _send_request_to_node(self, node_id, request):
+        groups_futures = {
+            group_id: self._find_coordinator_id_send_request(group_id)
+            for group_id in group_ids
+        }
+        self._wait_for_futures(groups_futures.values())
+        groups_coordinators = {
+            group_id: self._find_coordinator_id_process_response(future.value)
+            for group_id, future in groups_futures.items()
+        }
+        return groups_coordinators
+
+    def _send_request_to_node(self, node_id, request, wakeup=True):
         """Send a Kafka protocol message to a specific broker.
 
         Returns a future that may be polled for status and results.
 
         :param node_id: The broker id to which to send the message.
         :param request: The message to send.
+        :param wakeup: Optional flag to disable thread-wakeup.
         :return: A future object that may be polled for status and results.
         :exception: The exception if the message could not be sent.
         """
@@ -355,7 +370,7 @@ class KafkaAdminClient(object):
             # poll until the connection to broker is ready, otherwise send()
             # will fail with NodeNotReadyError
             self._client.poll()
-        return self._client.send(node_id, request)
+        return self._client.send(node_id, request, wakeup)
 
     def _send_request_to_controller(self, request):
         """Send a Kafka protocol message to the cluster controller.
@@ -682,7 +697,6 @@ class KafkaAdminClient(object):
         self._wait_for_futures([future])
         response = future.value
 
-
         return self._convert_create_acls_response_to_acls(acls, response)
 
     @staticmethod
@@ -1000,22 +1014,47 @@ class KafkaAdminClient(object):
         """Process a DescribeGroupsResponse into a group description."""
         if response.API_VERSION <= 3:
             assert len(response.groups) == 1
-            # TODO need to implement converting the response tuple into
-            # a more accessible interface like a namedtuple and then stop
-            # hardcoding tuple indices here. Several Java examples,
-            # including KafkaAdminClient.java
-            group_description = response.groups[0]
-            error_code = group_description[0]
+            for response_field, response_name in zip(response.SCHEMA.fields, response.SCHEMA.names):
+                if isinstance(response_field, Array):
+                    described_groups_field_schema = response_field.array_of
+                    described_group = response.__dict__[response_name][0]
+                    described_group_information_list = []
+                    protocol_type_is_consumer = False
+                    for (described_group_information, group_information_name, group_information_field) in zip(described_group, described_groups_field_schema.names, described_groups_field_schema.fields):
+                        if group_information_name == 'protocol_type':
+                            protocol_type = described_group_information
+                            protocol_type_is_consumer = (protocol_type == ConsumerProtocol.PROTOCOL_TYPE or not protocol_type)
+                        if isinstance(group_information_field, Array):
+                            member_information_list = []
+                            member_schema = group_information_field.array_of
+                            for members in described_group_information:
+                                member_information = []
+                                for (member, member_field, member_name)  in zip(members, member_schema.fields, member_schema.names):
+                                    if protocol_type_is_consumer:
+                                        if member_name == 'member_metadata' and member:
+                                            member_information.append(ConsumerProtocolMemberMetadata.decode(member))
+                                        elif member_name == 'member_assignment' and member:
+                                            member_information.append(ConsumerProtocolMemberAssignment.decode(member))
+                                        else:
+                                            member_information.append(member)
+                                member_info_tuple = MemberInformation._make(member_information)
+                                member_information_list.append(member_info_tuple)
+                            described_group_information_list.append(member_information_list)
+                        else:
+                            described_group_information_list.append(described_group_information)
+                    # Version 3 of the DescribeGroups API introduced the "authorized_operations" field.
+                    # This will cause the namedtuple to fail.
+                    # Therefore, appending a placeholder of None in it.
+                    if response.API_VERSION <=2:
+                        described_group_information_list.append(None)
+                    group_description = GroupInformation._make(described_group_information_list)
+            error_code = group_description.error_code
             error_type = Errors.for_code(error_code)
             # Java has the note: KAFKA-6789, we can retry based on the error code
             if error_type is not Errors.NoError:
                 raise error_type(
                     "DescribeGroupsResponse failed with response '{}'."
                     .format(response))
-            # TODO Java checks the group protocol type, and if consumer
-            # (ConsumerProtocol.PROTOCOL_TYPE) or empty string, it decodes
-            # the members' partition assignments... that hasn't yet been
-            # implemented here so just return the raw struct results
         else:
             raise NotImplementedError(
                 "Support for DescribeGroupsResponse_v{} has not yet been added to KafkaAdminClient."
@@ -1044,18 +1083,19 @@ class KafkaAdminClient(object):
             partition assignments.
         """
         group_descriptions = []
-        futures = []
-        for group_id in group_ids:
-            if group_coordinator_id is not None:
-                this_groups_coordinator_id = group_coordinator_id
-            else:
-                this_groups_coordinator_id = self._find_coordinator_id(group_id)
-            f = self._describe_consumer_groups_send_request(
+
+        if group_coordinator_id is not None:
+            groups_coordinators = {group_id: group_coordinator_id for group_id in group_ids}
+        else:
+            groups_coordinators = self._find_coordinator_ids(group_ids)
+
+        futures = [
+            self._describe_consumer_groups_send_request(
                 group_id,
-                this_groups_coordinator_id,
+                coordinator_id,
                 include_authorized_operations)
-            futures.append(f)
-
+            for group_id, coordinator_id in groups_coordinators.items()
+        ]
         self._wait_for_futures(futures)
 
         for future in futures:
@@ -1170,7 +1210,7 @@ class KafkaAdminClient(object):
 
         :param response: an OffsetFetchResponse.
         :return: A dictionary composed of TopicPartition keys and
-            OffsetAndMetada values.
+            OffsetAndMetadata values.
         """
         if response.API_VERSION <= 3:
 
@@ -1184,7 +1224,7 @@ class KafkaAdminClient(object):
                         .format(response))
 
             # transform response into a dictionary with TopicPartition keys and
-            # OffsetAndMetada values--this is what the Java AdminClient returns
+            # OffsetAndMetadata values--this is what the Java AdminClient returns
             offsets = {}
             for topic, partitions in response.topics:
                 for partition, offset, metadata, error_code in partitions:
@@ -1227,15 +1267,76 @@ class KafkaAdminClient(object):
             explicitly specified.
         """
         if group_coordinator_id is None:
-            group_coordinator_id = self._find_coordinator_id(group_id)
+            group_coordinator_id = self._find_coordinator_ids([group_id])[group_id]
         future = self._list_consumer_group_offsets_send_request(
                                     group_id, group_coordinator_id, partitions)
         self._wait_for_futures([future])
         response = future.value
         return self._list_consumer_group_offsets_process_response(response)
 
-    # delete groups protocol not yet implemented
-    # Note: send the request to the group's coordinator.
+    def delete_consumer_groups(self, group_ids, group_coordinator_id=None):
+        """Delete Consumer Group Offsets for given consumer groups.
+
+        Note:
+        This does not verify that the group ids actually exist and
+        group_coordinator_id is the correct coordinator for all these groups.
+
+        The result needs checking for potential errors.
+
+        :param group_ids: The consumer group ids of the groups which are to be deleted.
+        :param group_coordinator_id: The node_id of the broker which is the coordinator for
+            all the groups. Use only if all groups are coordinated by the same broker.
+            If set to None, will query the cluster to find the coordinator for every single group.
+            Explicitly specifying this can be useful to prevent
+            that extra network round trips if you already know the group
+            coordinator. Default: None.
+        :return: A list of tuples (group_id, KafkaError)
+        """
+        if group_coordinator_id is not None:
+            futures = [self._delete_consumer_groups_send_request(group_ids, group_coordinator_id)]
+        else:
+            coordinators_groups = defaultdict(list)
+            for group_id, coordinator_id in self._find_coordinator_ids(group_ids).items():
+                coordinators_groups[coordinator_id].append(group_id)
+            futures = [
+                self._delete_consumer_groups_send_request(group_ids, coordinator_id)
+                for coordinator_id, group_ids in coordinators_groups.items()
+            ]
+
+        self._wait_for_futures(futures)
+
+        results = []
+        for f in futures:
+            results.extend(self._convert_delete_groups_response(f.value))
+        return results
+
+    def _convert_delete_groups_response(self, response):
+        if response.API_VERSION <= 1:
+            results = []
+            for group_id, error_code in response.results:
+                results.append((group_id, Errors.for_code(error_code)))
+            return results
+        else:
+            raise NotImplementedError(
+                "Support for DeleteGroupsResponse_v{} has not yet been added to KafkaAdminClient."
+                    .format(response.API_VERSION))
+
+    def _delete_consumer_groups_send_request(self, group_ids, group_coordinator_id):
+        """Send a DeleteGroups request to a broker.
+
+        :param group_ids: The consumer group ids of the groups which are to be deleted.
+        :param group_coordinator_id: The node_id of the broker which is the coordinator for
+            all the groups.
+        :return: A message future
+        """
+        version = self._matching_api_version(DeleteGroupsRequest)
+        if version <= 1:
+            request = DeleteGroupsRequest[version](group_ids)
+        else:
+            raise NotImplementedError(
+                "Support for DeleteGroupsRequest_v{} has not yet been added to KafkaAdminClient."
+                    .format(version))
+        return self._send_request_to_node(group_coordinator_id, request)
 
     def _wait_for_futures(self, futures):
         while not all(future.succeeded() for future in futures):
@@ -1244,3 +1345,19 @@ class KafkaAdminClient(object):
 
                 if future.failed():
                     raise future.exception  # pylint: disable-msg=raising-bad-type
+
+    def describe_log_dirs(self):
+        """Send a DescribeLogDirsRequest request to a broker.
+
+        :return: A message future
+        """
+        version = self._matching_api_version(DescribeLogDirsRequest)
+        if version <= 0:
+            request = DescribeLogDirsRequest[version]()
+            future = self._send_request_to_node(self._client.least_loaded_node(), request)
+            self._wait_for_futures([future])
+        else:
+            raise NotImplementedError(
+                "Support for DescribeLogDirsRequest_v{} has not yet been added to KafkaAdminClient."
+                    .format(version))
+        return future.value

kafka/client_async.py

@@ -2,7 +2,6 @@ from __future__ import absolute_import, division
 
 import collections
 import copy
-import functools
 import logging
 import random
 import socket
@@ -202,10 +201,15 @@ class KafkaClient(object):
             if key in configs:
                 self.config[key] = configs[key]
 
+        # these properties need to be set on top of the initialization pipeline
+        # because they are used when __del__ method is called
+        self._closed = False
+        self._wake_r, self._wake_w = socket.socketpair()
+        self._selector = self.config['selector']()
+
         self.cluster = ClusterMetadata(**self.config)
         self._topics = set()  # empty set will fetch all topic metadata
         self._metadata_refresh_in_progress = False
-        self._selector = self.config['selector']()
         self._conns = Dict()  # object to support weakrefs
         self._api_versions = None
         self._connecting = set()
@@ -213,7 +217,6 @@ class KafkaClient(object):
         self._refresh_on_disconnects = True
         self._last_bootstrap = 0
         self._bootstrap_fails = 0
-        self._wake_r, self._wake_w = socket.socketpair()
         self._wake_r.setblocking(False)
         self._wake_w.settimeout(self.config['wakeup_timeout_ms'] / 1000.0)
         self._wake_lock = threading.Lock()
@@ -227,7 +230,6 @@ class KafkaClient(object):
 
         self._selector.register(self._wake_r, selectors.EVENT_READ)
         self._idle_expiry_manager = IdleConnectionManager(self.config['connections_max_idle_ms'])
-        self._closed = False
         self._sensors = None
         if self.config['metrics']:
             self._sensors = KafkaClientMetrics(self.config['metrics'],

kafka/codec.py

@@ -10,12 +10,18 @@ from kafka.vendor.six.moves import range
 
 _XERIAL_V1_HEADER = (-126, b'S', b'N', b'A', b'P', b'P', b'Y', 0, 1, 1)
 _XERIAL_V1_FORMAT = 'bccccccBii'
+ZSTD_MAX_OUTPUT_SIZE = 1024 * 1024
 
 try:
     import snappy
 except ImportError:
     snappy = None
 
+try:
+    import zstandard as zstd
+except ImportError:
+    zstd = None
+
 try:
     import lz4.frame as lz4
 
@@ -58,6 +64,10 @@ def has_snappy():
     return snappy is not None
 
 
+def has_zstd():
+    return zstd is not None
+
+
 def has_lz4():
     if lz4 is not None:
         return True
@@ -177,7 +187,7 @@ def _detect_xerial_stream(payload):
         The version is the version of this format as written by xerial,
         in the wild this is currently 1 as such we only support v1.
 
-        Compat is there to claim the miniumum supported version that
+        Compat is there to claim the minimum supported version that
         can read a xerial block stream, presently in the wild this is
         1.
     """
@@ -299,3 +309,18 @@ def lz4_decode_old_kafka(payload):
         payload[header_size:]
     ])
     return lz4_decode(munged_payload)
+
+
+def zstd_encode(payload):
+    if not zstd:
+        raise NotImplementedError("Zstd codec is not available")
+    return zstd.ZstdCompressor().compress(payload)
+
+
+def zstd_decode(payload):
+    if not zstd:
+        raise NotImplementedError("Zstd codec is not available")
+    try:
+        return zstd.ZstdDecompressor().decompress(payload)
+    except zstd.ZstdError:
+        return zstd.ZstdDecompressor().decompress(payload, max_output_size=ZSTD_MAX_OUTPUT_SIZE)

kafka/conn.py

@@ -24,9 +24,12 @@ import kafka.errors as Errors
 from kafka.future import Future
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.oauth.abstract import AbstractTokenProvider
-from kafka.protocol.admin import SaslHandShakeRequest
+from kafka.protocol.admin import SaslHandShakeRequest, DescribeAclsRequest_v2, DescribeClientQuotasRequest
 from kafka.protocol.commit import OffsetFetchRequest
+from kafka.protocol.offset import OffsetRequest
+from kafka.protocol.produce import ProduceRequest
 from kafka.protocol.metadata import MetadataRequest
+from kafka.protocol.fetch import FetchRequest
 from kafka.protocol.parser import KafkaProtocol
 from kafka.protocol.types import Int32, Int8
 from kafka.scram import ScramClient
@@ -75,7 +78,7 @@ except ImportError:
 try:
     import gssapi
     from gssapi.raw.misc import GSSError
-except ImportError:
+except (ImportError, OSError):
     #no gssapi available, will disable gssapi mechanism
     gssapi = None
     GSSError = None
@@ -493,7 +496,7 @@ class BrokerConnection(object):
         try:
             self._sock = self._ssl_context.wrap_socket(
                 self._sock,
-                server_hostname=self.host,
+                server_hostname=self.host.rstrip("."),
                 do_handshake_on_connect=False)
         except ssl.SSLError as e:
             log.exception('%s: Failed to wrap socket in SSLContext!', self)
@@ -913,7 +916,7 @@ class BrokerConnection(object):
         with self._lock:
             if self.state is ConnectionStates.DISCONNECTED:
                 return
-            log.info('%s: Closing connection. %s', self, error or '')
+            log.log(logging.ERROR if error else logging.INFO, '%s: Closing connection. %s', self, error or '')
             self._update_reconnect_backoff()
             self._sasl_auth_future = None
             self._protocol = KafkaProtocol(
@@ -1166,6 +1169,14 @@ class BrokerConnection(object):
         # in reverse order. As soon as we find one that works, return it
         test_cases = [
             # format (<broker version>, <needed struct>)
+            ((2, 6, 0), DescribeClientQuotasRequest[0]),
+            ((2, 5, 0), DescribeAclsRequest_v2),
+            ((2, 4, 0), ProduceRequest[8]),
+            ((2, 3, 0), FetchRequest[11]),
+            ((2, 2, 0), OffsetRequest[5]),
+            ((2, 1, 0), FetchRequest[10]),
+            ((2, 0, 0), FetchRequest[8]),
+            ((1, 1, 0), FetchRequest[7]),
             ((1, 0, 0), MetadataRequest[5]),
             ((0, 11, 0), MetadataRequest[4]),
             ((0, 10, 2), OffsetFetchRequest[2]),

kafka/consumer/fetcher.py

@@ -125,7 +125,7 @@ class Fetcher(six.Iterator):
                 log.debug("Sending FetchRequest to node %s", node_id)
                 future = self._client.send(node_id, request, wakeup=False)
                 future.add_callback(self._handle_fetch_response, request, time.time())
-                future.add_errback(log.error, 'Fetch to node %s failed: %s', node_id)
+                future.add_errback(self._handle_fetch_error, node_id)
                 futures.append(future)
         self._fetch_futures.extend(futures)
         self._clean_done_fetch_futures()
@@ -293,7 +293,7 @@ class Fetcher(six.Iterator):
                 # Issue #1780
                 # Recheck partition existence after after a successful metadata refresh
                 if refresh_future.succeeded() and isinstance(future.exception, Errors.StaleMetadata):
-                    log.debug("Stale metadata was raised, and we now have an updated metadata. Rechecking partition existance")
+                    log.debug("Stale metadata was raised, and we now have an updated metadata. Rechecking partition existence")
                     unknown_partition = future.exception.args[0]  # TopicPartition from StaleMetadata
                     if self._client.cluster.leader_for_partition(unknown_partition) is None:
                         log.debug("Removed partition %s from offsets retrieval" % (unknown_partition, ))
@@ -474,7 +474,9 @@ class Fetcher(six.Iterator):
                         self.config['value_deserializer'],
                         tp.topic, record.value)
                     headers = record.headers
-                    header_size = sum(len(h_key.encode("utf-8")) + len(h_val) for h_key, h_val in headers) if headers else -1
+                    header_size = sum(
+                        len(h_key.encode("utf-8")) + (len(h_val) if h_val is not None else 0) for h_key, h_val in
+                        headers) if headers else -1
                     yield ConsumerRecord(
                         tp.topic, tp.partition, record.offset, record.timestamp,
                         record.timestamp_type, key, value, headers, record.checksum,
@@ -776,6 +778,14 @@ class Fetcher(six.Iterator):
             self._sensors.fetch_throttle_time_sensor.record(response.throttle_time_ms)
         self._sensors.fetch_latency.record((time.time() - send_time) * 1000)
 
+    def _handle_fetch_error(self, node_id, exception):
+        log.log(
+            logging.INFO if isinstance(exception, Errors.Cancelled) else logging.ERROR,
+            'Fetch to node %s failed: %s',
+            node_id,
+            exception
+        )
+
     def _parse_fetched_data(self, completed_fetch):
         tp = completed_fetch.topic_partition
         fetch_offset = completed_fetch.fetched_offset
@@ -815,8 +825,9 @@ class Fetcher(six.Iterator):
                               position)
                     unpacked = list(self._unpack_message_set(tp, records))
                     parsed_records = self.PartitionRecords(fetch_offset, tp, unpacked)
-                    last_offset = unpacked[-1].offset
-                    self._sensors.records_fetch_lag.record(highwater - last_offset)
+                    if unpacked:
+                        last_offset = unpacked[-1].offset
+                        self._sensors.records_fetch_lag.record(highwater - last_offset)
                     num_bytes = records.valid_bytes()
                     records_count = len(unpacked)
                 elif records.size_in_bytes() > 0:

kafka/consumer/group.py

@@ -167,7 +167,8 @@ class KafkaConsumer(six.Iterator):
             message iteration before raising StopIteration (i.e., ending the
             iterator). Default block forever [float('inf')].
         security_protocol (str): Protocol used to communicate with brokers.
-            Valid values are: PLAINTEXT, SSL. Default: PLAINTEXT.
+            Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
+            Default: PLAINTEXT.
         ssl_context (ssl.SSLContext): Pre-configured SSLContext for wrapping
             socket connections. If provided, all other ssl_* configurations
             will be ignored. Default: None.
@@ -243,6 +244,7 @@ class KafkaConsumer(six.Iterator):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     Note:
         Configuration parameters are described in more detail at
@@ -305,6 +307,7 @@ class KafkaConsumer(six.Iterator):
         'sasl_kerberos_domain_name': None,
         'sasl_oauth_token_provider': None,
         'legacy_iterator': False, # enable to revert to < 1.4.7 iterator
+        'kafka_client': KafkaClient,
     }
     DEFAULT_SESSION_TIMEOUT_MS_0_9 = 30000
 
@@ -352,7 +355,7 @@ class KafkaConsumer(six.Iterator):
             log.warning('use api_version=%s [tuple] -- "%s" as str is deprecated',
                         str(self.config['api_version']), str_version)
 
-        self._client = KafkaClient(metrics=self._metrics, **self.config)
+        self._client = self.config['kafka_client'](metrics=self._metrics, **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:
@@ -650,7 +653,7 @@ class KafkaConsumer(six.Iterator):
         # Poll for new data until the timeout expires
         start = time.time()
         remaining = timeout_ms
-        while True:
+        while not self._closed:
             records = self._poll_once(remaining, max_records, update_offsets=update_offsets)
             if records:
                 return records
@@ -659,7 +662,9 @@ class KafkaConsumer(six.Iterator):
             remaining = timeout_ms - elapsed_ms
 
             if remaining <= 0:
-                return {}
+                break
+
+        return {}
 
     def _poll_once(self, timeout_ms, max_records, update_offsets=True):
         """Do one round of polling. In addition to checking for new data, this does