Add auto promototing client class

rnpridgeon · rnpridgeon · commit c8ae724dbc25 · 2018-12-30T13:07:01.000-05:00
diff --git a/confluent_kafka/consumer.py b/confluent_kafka/consumer.py
@@ -22,6 +22,34 @@
 
 
 class Consumer(_impl):
+    """
+        Create a new Kafka Consumer instance with or without serializer support.
+
+        To avoid spontaneous calls from non-Python threads all callbacks will only be served upon
+        calling ```client.poll()``` or ```client.flush()```.
+
+        :param dict conf: Configuration properties.
+            See https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md for more information.
+        :param func key_deserializer(topic, key): Converts message key bytes to object.
+            **note** deserializers are responsible for handling NULL keys
+        :param func value_deserializer(topic, value): Converts message value bytes to object.
+            **note** deserializers are responsible for handling NULL values
+        :param func on_commit(err, [partitions]): Callback used to indicate success or failure
+            of an offset commit.
+        :param func stats_cb(json_str): Callback for statistics emitted every ``statistics.interval.ms``.
+            See https://github.com/edenhill/librdkafka/wiki/Statistics” for more information.
+        :param func throttle_cb(confluent_kafka.ThrottleEvent): Callback for throttled request reporting.
+        :param logging.handlers logger: Forwards logs from the Kafka client to the provided handler instance.
+            Log messages will only be forwarded when ``client.poll()`` or ``producer.flush()`` are called.
+        :raises TypeError: If conf is not a dict.
+    """
+    def __new__(cls, *args, **kwargs):
+        if 'key_deserializer' in kwargs or 'value_deserializer' in kwargs:
+            return super(Consumer, cls).__new__(DeserializingConsumer, *args, **kwargs)
+        return super(Consumer, cls).__new__(cls, *args, **kwargs)
+
+
+class DeserializingConsumer(Consumer):
     """
         Create a new Kafka Consumer instance.
 
@@ -78,10 +106,11 @@ def __init__(self, conf={}, key_deserializer=None, value_deserializer=None,
 
         if not logger:
             logger = conf.get('logger', None)
-        super(Consumer, self).__init__(conf, on_commit=on_commit, stats_cb=stats_cb,
-                                       throttle_cb=throttle_cb, logger=logger)
 
-    def poll(self, timeout=-1.0, key_deserializer=None, value_deserializer=None):
+        super(DeserializingConsumer, self).__init__(conf, on_commit=on_commit, stats_cb=stats_cb,
+                                                    throttle_cb=throttle_cb, logger=logger)
+
+    def poll(self, timeout=-1.0):
         """
             Consumes a message, triggers callbacks, returns an event.
 
@@ -97,25 +126,15 @@ def poll(self, timeout=-1.0, key_deserializer=None, value_deserializer=None):
             :raises RuntimeError: If called on a closed consumer.
         """
 
-        msg = super(Consumer, self).poll(timeout)
+        msg = super(DeserializingConsumer, self).poll(timeout)
 
-        if not msg or msg.error():
+        if msg is None or msg.error():
             return msg
 
         topic = msg.topic()
 
-        # parameter overrides take precedence over instance functions
-        if not key_deserializer:
-            key_deserializer = self._key_deserializer
-
-        if key_deserializer:
-            msg.set_key(key_deserializer(topic, msg.key()))
-
-        if not value_deserializer:
-            value_deserializer = self._value_deserializer
-
-        if value_deserializer:
-            msg.set_value(value_deserializer(topic, msg.value()))
+        msg.set_key(self._key_deserializer(topic, msg.key()))
+        msg.set_value(self._value_deserializer(topic, msg.value()))
 
         return msg
 
@@ -138,8 +157,12 @@ def consume(self, num_messages=1, timeout=-1):
             :raises ValueError: If num_messages > 1M.
         """
 
-        # Disable consume() method when serializers are in use.
-        if self._key_deserializer or self._value_deserializer:
-            raise(NotImplementedError, "Batch consumption does not support the use of deserializers")
+        msgset = super(DeserializingConsumer, self).consume(num_messages, timeout)
+        for msg in msgset:
+            if msg.error():
+                continue
+
+            msg.set_key(self._key_deserializer(topic, msg.key()))
+            msg.set_value(self._value_deserializer(topic, msg.value()))
 
-        return super(Consumer, self).consume(num_messages, timeout)
+        return msgset
diff --git a/confluent_kafka/producer.py b/confluent_kafka/producer.py
@@ -23,7 +23,35 @@
 
 class Producer(_impl):
     """
-        Create a new Kafka Producer instance.
+        Create a new Kafka Producer instance with or without serializer support.
+
+        To avoid spontaneous calls from non-Python threads all callbacks will only be served upon
+            calling ```client.poll()``` or ```client.flush()```.
+
+        :param dict conf: Configuration properties. At a minimum ``bootstrap.servers`` **should** be set.
+            See https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md for more information.
+        :param func key_serializer(topic, key): Converts key to bytes.
+            **note** serializers are responsible for handling NULL keys
+        :param func value_serializer(topic, value): Converts value to bytes.
+            **note** serializers are responsible for handling NULL keys
+        :param func error_cb(kafka.KafkaError): Callback for generic/global error events.
+        :param func stats_cb(json_str): Callback for statistics emitted every ``statistics.interval.ms``.
+            See https://github.com/edenhill/librdkafka/wiki/Statistics” for more information.
+        :param func throttle_cb(confluent_kafka.ThrottleEvent): Callback for throttled request reporting.
+        :param logging.handlers logger: Forwards logs from the Kafka client to the provided handler instance.
+            Log messages will only be forwarded when ``client.poll()`` or ``producer.flush()`` are called.
+        :raises TypeError: If conf is not a dict object.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        if 'key_serializer' in kwargs or 'value_serializer' in kwargs:
+            return super(Producer, cls).__new__(SerializingProducer, *args, **kwargs)
+        return super(Producer, cls).__new__(cls, *args, **kwargs)
+
+
+class SerializingProducer(Producer):
+    """
+        Create a new Kafka Producer instance with custom serialization support.
 
         To avoid spontaneous calls from non-Python threads all callbacks will only be served upon
             calling ```client.poll()``` or ```client.flush()```.
@@ -78,12 +106,11 @@ def __init__(self, conf={}, key_serializer=None, value_serializer=None,
         if not logger:
             logger = conf.get('logger', None)
 
-        super(Producer, self).__init__(conf, error_cb=error_cb, stats_cb=stats_cb,
-                                       throttle_cb=throttle_cb, logger=logger)
+        super(SerializingProducer, self).__init__(conf, error_cb=error_cb, stats_cb=stats_cb,
+                                                  throttle_cb=throttle_cb, logger=logger)
 
     def produce(self, topic, value=None, key=None, partition=PARTITION_UA,
-                on_delivery=None, callback=None, timestamp=0, headers=None,
-                key_serializer=None, value_serializer=None):
+                on_delivery=None, callback=None, timestamp=0, headers=None):
         """
             Produces message to Kafka
 
@@ -109,22 +136,8 @@ def produce(self, topic, value=None, key=None, partition=PARTITION_UA,
 
         """
 
-        # on_delivery is an alias for callback and take precedence
-        if callback and not on_delivery:
-            on_delivery = callback
-
-        # parameter overrides take precedence over instance functions
-        if not key_serializer:
-            key_serializer = self._key_serializer
-
-        if key_serializer:
-            key = key_serializer(topic, key)
-
-        if not value_serializer:
-            value_serializer = self._value_serializer
-
-        if value_serializer:
-            value = value_serializer(topic, value)
+        key = self._key_serializer(topic, key)
+        value = self._value_serializer(topic, value)
 
-        super(Producer, self).produce(topic, value, key, partition, on_delivery=on_delivery,
-                                      timestamp=timestamp, headers=headers)
+        super(SerializingProducer, self).produce(topic, value, key, partition, on_delivery=on_delivery,
+                                                 callback=callback, timestamp=timestamp, headers=headers)
diff --git a/tests/test_Consumer.py b/tests/test_Consumer.py
@@ -3,6 +3,7 @@
 from confluent_kafka import (Consumer, TopicPartition, KafkaError,
                              KafkaException, TIMESTAMP_NOT_AVAILABLE,
                              OFFSET_INVALID, libversion)
+from confluent_kafka.consumer import DeserializingConsumer
 import pytest
 
 
@@ -321,3 +322,19 @@ def test_consumer_withot_groupid():
     with pytest.raises(ValueError) as ex:
         Consumer({'bootstrap.servers': "mybroker:9092"})
     assert 'group.id must be set' in str(ex)
+
+
+@pytest.mark.parametrize("test_param", [
+    ({"key_deserializer": lambda key: key}, DeserializingConsumer),
+    ({"value_deserializer": lambda value: value}, DeserializingConsumer),
+    ({"key_deserializer": lambda key: key,
+      "value_deserializer": lambda value: value}, DeserializingConsumer),
+    ({}, Consumer)
+])
+def test_producer_factory(test_param):
+    """
+    Assert Consumer factory returns the correct class
+    """
+
+    c = Consumer({'group.id': 'wrapper'}, **test_param[0])
+    assert(type(c) is test_param[1])
diff --git a/tests/test_Producer.py b/tests/test_Producer.py
@@ -3,6 +3,7 @@
 import pytest
 
 from confluent_kafka import Producer, KafkaError, KafkaException, libversion
+from confluent_kafka.producer import SerializingProducer
 from struct import pack
 
 
@@ -191,3 +192,19 @@ def test_set_invalid_partitioner_murmur():
     with pytest.raises(KafkaException) as e:
         Producer({'partitioner': 'murmur'})
     assert 'Invalid value for configuration property "partitioner": murmur' in str(e)
+
+
+@pytest.mark.parametrize("test_param", [
+    ({"key_serializer": lambda key: key}, SerializingProducer),
+    ({"value_serializer": lambda value: value}, SerializingProducer),
+    ({"key_serializer": lambda key: key,
+      "value_serializer": lambda value: value}, SerializingProducer),
+    ({}, Producer)
+])
+def test_producer_factory(test_param):
+    """
+    Assert Producer factory returns the correct class
+    """
+
+    p = Producer(**test_param[0])
+    assert(type(p) is test_param[1])
