class AbstractConsumer(ABC):
def __init__(self, group_id: str, topic_name: str, last: bool):
offset_reset = "earliest"
if last:
offset_reset = "latest"
self._config = Config().create_confluent_config()
self._config.update({
"group.id": group_id,
"error_cb": raise_for_kafka_error,
# We need to commit offsets manually once we"re sure it got saved
# to the sink
"enable.auto.commit": True,
"enable.partition.eof": False,
# We need this to start at the last committed offset instead of the
# latest when subscribing for the first time
"default.topic.config": {
"auto.offset.reset": offset_reset
},
})
self._consumer = confluent_kafka.Consumer(self._config)
self._subscribe(topic_name)
def _subscribe(self, topic: str) -> None:
self._consumer.subscribe([topic])
@abstractmethod
def consume(self, amount: int) -> int:
pass
def _consume_single_message(self, timeout=30) -> Optional[Message]:
message = self._consumer.poll(timeout=timeout)
raise_for_message(message)
return message
class Producer(ABC):
def __init__(self):
self.queue_length = 100000
self.internal_queue_length_limit = self.queue_length / 0.5
self._config = Config().create_confluent_config()
self._config.update({
"on_delivery": delivery_callback,
"error_cb": raise_for_kafka_error,
"queue.buffering.max.messages": self.queue_length,
})
@abstractmethod
def produce(self, topic_name: str) -> int:
pass
def consumer(topic_object: Topic, consumer_group):
_config = Config().create_confluent_config()
_config.update({
"group.id": consumer_group,
"error_cb": raise_for_kafka_error,
# We need to commit offsets manually once we"re sure it got saved
# to the sink
"enable.auto.commit": False,
"enable.partition.eof": False,
# We need this to start at the last committed offset instead of the
# latest when subscribing for the first time
"default.topic.config": {
"auto.offset.reset": "latest"
},
})
_consumer = confluent_kafka.Consumer(_config)
_consumer.assign(
[TopicPartition(topic=topic_object.name, partition=0, offset=0)])
yield _consumer