deps: python3-kafka sync our package with upstream [SRE-7606]

kafka/admin/client.py

@@ -147,6 +147,7 @@ class KafkaAdminClient(object):
         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
+        socks5_proxy (str): Socks5 proxy url. Default: None
 
     """
     DEFAULT_CONFIG = {
@@ -182,6 +183,7 @@ class KafkaAdminClient(object):
         'sasl_kerberos_service_name': 'kafka',
         'sasl_kerberos_domain_name': None,
         'sasl_oauth_token_provider': None,
+        'socks5_proxy': None,
 
         # metrics configs
         'metric_reporters': [],
@@ -272,7 +274,7 @@ def _refresh_controller_id(self):
         version = self._matching_api_version(MetadataRequest)
         if 1 <= version <= 6:
             request = MetadataRequest[version]()
-            future = self._send_request_to_node(self._client.least_loaded_node(), request)
+            future = self._send_request_to_least_loaded_node(request)
 
             self._wait_for_futures([future])
 
@@ -310,7 +312,7 @@ def _find_coordinator_id_send_request(self, group_id):
             raise NotImplementedError(
                 "Support for GroupCoordinatorRequest_v{} has not yet been added to KafkaAdminClient."
                 .format(version))
-        return self._send_request_to_node(self._client.least_loaded_node(), request)
+        return self._send_request_to_least_loaded_node(request)
 
     def _find_coordinator_id_process_response(self, response):
         """Process a FindCoordinatorResponse.
@@ -355,9 +357,36 @@ def _find_coordinator_ids(self, group_ids):
         }
         return groups_coordinators
 
-    def _send_request_to_node(self, node_id, request, wakeup=True):
+    def _send_request_to_least_loaded_node(self, request):
+        """Send a Kafka protocol message to the least loaded broker.
+
+        Returns a future that may be polled for status and results.
+
+        :param request: The message to send.
+        :return: A future object that may be polled for status and results.
+        :exception: The exception if the message could not be sent.
+        """
+        node_id = self._client.least_loaded_node()
+        while not self._client.ready(node_id):
+            # poll until the connection to broker is ready, otherwise send()
+            # will fail with NodeNotReadyError
+            self._client.poll()
+
+            # node_id is not part of the cluster anymore, choose a new broker
+            # to connect to
+            if self._client.cluster.broker_metadata(node_id) is None:
+                node_id = self._client.least_loaded_node()
+
+        return self._client.send(node_id, request)
+
+    def _send_request_to_node(self, node_id, request):
         """Send a Kafka protocol message to a specific broker.
 
+        .. note::
+
+            This function will enter in an infinite loop if `node_id` is
+            removed from the cluster.
+
         Returns a future that may be polled for status and results.
 
         :param node_id: The broker id to which to send the message.
@@ -383,10 +412,23 @@ def _send_request_to_controller(self, request):
         tries = 2  # in case our cached self._controller_id is outdated
         while tries:
             tries -= 1
-            future = self._send_request_to_node(self._controller_id, request)
+            future = self._client.send(self._controller_id, request)
 
             self._wait_for_futures([future])
 
+            if future.exception is not None:
+                log.error(
+                    "Sending request to controller_id %s failed with %s",
+                    self._controller_id,
+                    future.exception,
+                )
+                is_outdated_controler = (
+                    self._client.cluster.broker_metadata(self._controller_id) is None
+                )
+                if is_outdated_controler:
+                    self._refresh_controller_id()
+                continue
+
             response = future.value
             # In Java, the error field name is inconsistent:
             #  - CreateTopicsResponse / CreatePartitionsResponse uses topic_errors
@@ -507,10 +549,7 @@ def _get_cluster_metadata(self, topics=None, auto_topic_creation=False):
                 allow_auto_topic_creation=auto_topic_creation
             )
 
-        future = self._send_request_to_node(
-            self._client.least_loaded_node(),
-            request
-        )
+        future = self._send_request_to_least_loaded_node(request)
         self._wait_for_futures([future])
         return future.value
 
@@ -602,7 +641,7 @@ def describe_acls(self, acl_filter):
                     .format(version)
             )
 
-        future = self._send_request_to_node(self._client.least_loaded_node(), request)
+        future = self._send_request_to_least_loaded_node(request)
         self._wait_for_futures([future])
         response = future.value
 
@@ -693,7 +732,7 @@ def create_acls(self, acls):
                     .format(version)
             )
 
-        future = self._send_request_to_node(self._client.least_loaded_node(), request)
+        future = self._send_request_to_least_loaded_node(request)
         self._wait_for_futures([future])
         response = future.value
 
@@ -787,7 +826,7 @@ def delete_acls(self, acl_filters):
                     .format(version)
             )
 
-        future = self._send_request_to_node(self._client.least_loaded_node(), request)
+        future = self._send_request_to_least_loaded_node(request)
         self._wait_for_futures([future])
         response = future.value
 
@@ -847,8 +886,7 @@ def describe_configs(self, config_resources, include_synonyms=False):
                     ))
 
             if len(topic_resources) > 0:
-                futures.append(self._send_request_to_node(
-                    self._client.least_loaded_node(),
+                futures.append(self._send_request_to_least_loaded_node(
                     DescribeConfigsRequest[version](resources=topic_resources)
                 ))
 
@@ -868,8 +906,7 @@ def describe_configs(self, config_resources, include_synonyms=False):
                     ))
 
             if len(topic_resources) > 0:
-                futures.append(self._send_request_to_node(
-                    self._client.least_loaded_node(),
+                futures.append(self._send_request_to_least_loaded_node(
                     DescribeConfigsRequest[version](resources=topic_resources, include_synonyms=include_synonyms)
                 ))
         else:
@@ -916,7 +953,7 @@ def alter_configs(self, config_resources):
         # // a single request that may be sent to any broker.
         #
         # So this is currently broken as it always sends to the least_loaded_node()
-        future = self._send_request_to_node(self._client.least_loaded_node(), request)
+        future = self._send_request_to_least_loaded_node(request)
 
         self._wait_for_futures([future])
         response = future.value

kafka/client_async.py

@@ -154,6 +154,7 @@ class KafkaClient(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        socks5_proxy (str): Socks5 proxy URL. Default: None
     """
 
     DEFAULT_CONFIG = {
@@ -192,7 +193,8 @@ class KafkaClient(object):
         'sasl_plain_password': None,
         'sasl_kerberos_service_name': 'kafka',
         'sasl_kerberos_domain_name': None,
-        'sasl_oauth_token_provider': None
+        'sasl_oauth_token_provider': None,
+        'socks5_proxy': None,
     }
 
     def __init__(self, **configs):
@@ -368,18 +370,26 @@ def _maybe_connect(self, node_id):
             conn = self._conns.get(node_id)
 
             if conn is None:
-                broker = self.cluster.broker_metadata(node_id)
-                assert broker, 'Broker id %s not in current metadata' % (node_id,)
-
-                log.debug("Initiating connection to node %s at %s:%s",
-                          node_id, broker.host, broker.port)
-                host, port, afi = get_ip_port_afi(broker.host)
-                cb = WeakMethod(self._conn_state_change)
-                conn = BrokerConnection(host, broker.port, afi,
-                                        state_change_callback=cb,
-                                        node_id=node_id,
-                                        **self.config)
-                self._conns[node_id] = conn
+                broker_metadata = self.cluster.broker_metadata(node_id)
+
+                # The broker may have been removed from the cluster after the
+                # call to `maybe_connect`. At this point there is no way to
+                # recover, so just ignore the connection
+                if broker_metadata is None:
+                    log.debug("Node %s is not available anymore, discarding connection", node_id)
+                    if node_id in self._connecting:
+                        self._connecting.remove(node_id)
+                    return False
+                else:
+                    log.debug("Initiating connection to node %s at %s:%s",
+                              node_id, broker_metadata.host, broker_metadata.port)
+                    host, port, afi = get_ip_port_afi(broker_metadata.host)
+                    cb = WeakMethod(self._conn_state_change)
+                    conn = BrokerConnection(host, broker_metadata.port, afi,
+                                            state_change_callback=cb,
+                                            node_id=node_id,
+                                            **self.config)
+                    self._conns[node_id] = conn
 
             # Check if existing connection should be recreated because host/port changed
             elif self._should_recycle_connection(conn):
@@ -637,6 +647,9 @@ def _poll(self, timeout):
             self._sensors.select_time.record((end_select - start_select) * 1000000000)
 
         for key, events in ready:
+            if key.fileobj.fileno() < 0:
+                time.sleep(0.1)
+
             if key.fileobj is self._wake_r:
                 self._clear_wake_fd()
                 continue

kafka/conn.py

@@ -33,6 +33,7 @@
 from kafka.protocol.parser import KafkaProtocol
 from kafka.protocol.types import Int32, Int8
 from kafka.scram import ScramClient
+from kafka.socks5_wrapper import Socks5Wrapper
 from kafka.version import __version__
 
 
@@ -191,6 +192,7 @@ class BrokerConnection(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        socks5_proxy (str): Socks5 proxy url. Default: None
     """
 
     DEFAULT_CONFIG = {
@@ -224,7 +226,8 @@ class BrokerConnection(object):
         'sasl_plain_password': None,
         'sasl_kerberos_service_name': 'kafka',
         'sasl_kerberos_domain_name': None,
-        'sasl_oauth_token_provider': None
+        'sasl_oauth_token_provider': None,
+        'socks5_proxy': None,
     }
     SECURITY_PROTOCOLS = ('PLAINTEXT', 'SSL', 'SASL_PLAINTEXT', 'SASL_SSL')
     SASL_MECHANISMS = ('PLAIN', 'GSSAPI', 'OAUTHBEARER', "SCRAM-SHA-256", "SCRAM-SHA-512")
@@ -236,6 +239,7 @@ def __init__(self, host, port, afi, **configs):
         self._sock_afi = afi
         self._sock_addr = None
         self._api_versions = None
+        self._socks5_proxy = None
 
         self.config = copy.copy(self.DEFAULT_CONFIG)
         for key in self.config:
@@ -368,7 +372,11 @@ def connect(self):
                 log.debug('%s: creating new socket', self)
                 assert self._sock is None
                 self._sock_afi, self._sock_addr = next_lookup
-                self._sock = socket.socket(self._sock_afi, socket.SOCK_STREAM)
+                if self.config["socks5_proxy"] is not None:
+                    self._socks5_proxy = Socks5Wrapper(self.config["socks5_proxy"], self.afi)
+                    self._sock = self._socks5_proxy.socket(self._sock_afi, socket.SOCK_STREAM)
+                else:
+                    self._sock = socket.socket(self._sock_afi, socket.SOCK_STREAM)
 
             for option in self.config['socket_options']:
                 log.debug('%s: setting socket option %s', self, option)
@@ -385,7 +393,10 @@ def connect(self):
             # to check connection status
             ret = None
             try:
-                ret = self._sock.connect_ex(self._sock_addr)
+                if self._socks5_proxy:
+                    ret = self._socks5_proxy.connect_ex(self._sock_addr)
+                else:
+                    ret = self._sock.connect_ex(self._sock_addr)
             except socket.error as err:
                 ret = err.errno
 

kafka/consumer/group.py

@@ -245,6 +245,7 @@ class KafkaConsumer(six.Iterator):
         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
+        socks5_proxy (str): Socks5 proxy URL. Default: None
 
     Note:
         Configuration parameters are described in more detail at
@@ -308,6 +309,7 @@ class KafkaConsumer(six.Iterator):
         'sasl_oauth_token_provider': None,
         'legacy_iterator': False, # enable to revert to < 1.4.7 iterator
         'kafka_client': KafkaClient,
+        'socks5_proxy': None,
     }
     DEFAULT_SESSION_TIMEOUT_MS_0_9 = 30000
 

kafka/producer/kafka.py

@@ -281,6 +281,7 @@ class KafkaProducer(object):
         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
+        socks5_proxy (str): Socks5 proxy URL. Default: None
 
     Note:
         Configuration parameters are described in more detail at
@@ -335,6 +336,7 @@ class KafkaProducer(object):
         'sasl_kerberos_domain_name': None,
         'sasl_oauth_token_provider': None,
         'kafka_client': KafkaClient,
+        'socks5_proxy': None,
     }
 
     _COMPRESSORS = {