def poll(self,
timeout: Optional[float] = None) -> Optional[Message[TPayload]]:
with self.__lock:
if self.__closed:
raise RuntimeError("consumer is closed")
while self.__pending_callbacks:
callback = self.__pending_callbacks.popleft()
callback()
for partition, offset in sorted(self.__offsets.items()):
if partition in self.__paused:
continue # skip paused partitions
try:
message = self.__broker.consume(partition, offset)
except ConsumerError:
raise
except Exception as e:
raise ConsumerError("error consuming mesage") from e
if message is None:
if self.__enable_end_of_partition and (
partition not in self.__last_eof_at
or offset > self.__last_eof_at[partition]):
self.__last_eof_at[partition] = offset
raise EndOfPartition(partition, offset)
else:
self.__offsets[partition] = message.next_offset
return message
return None
def pause(self, partitions: Sequence[Partition]) -> None:
"""
Pause the consumption of messages for the provided partitions.
Raises an ``InvalidState`` if called on a closed consumer.
"""
if self.__state in {
KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR
}:
raise InvalidState(self.__state)
if set(partitions) - self.__offsets.keys():
raise ConsumerError("cannot pause unassigned partitions")
self.__consumer.pause([
ConfluentTopicPartition(partition.topic.name, partition.index)
for partition in partitions
])
self.__paused.update(partitions)
# XXX: Seeking to a specific partition offset and immediately pausing
# that partition causes the seek to be ignored for some reason.
self.seek({
partition: offset
for partition, offset in self.__offsets.items()
if partition in partitions
})
def __validate_offsets(self, offsets: Mapping[Partition, int]) -> None:
invalid_offsets: Mapping[Partition, int] = {
partition: offset for partition, offset in offsets.items() if offset < 0
}
if invalid_offsets:
raise ConsumerError(f"invalid offsets: {invalid_offsets!r}")
def pause(self, partitions: Sequence[Partition]) -> None:
with self.__lock:
if self.__closed:
raise RuntimeError("consumer is closed")
if set(partitions) - self.__offsets.keys():
raise ConsumerError("cannot pause unassigned partitions")
self.__paused.update(partitions)
def resume(self, partitions: Sequence[Partition]) -> None:
with self.__lock:
if self.__closed:
raise RuntimeError("consumer is closed")
if set(partitions) - self.__offsets.keys():
raise ConsumerError("cannot resume unassigned partitions")
for partition in partitions:
self.__paused.discard(partition)
def seek(self, offsets: Mapping[Partition, int]) -> None:
with self.__lock:
if self.__closed:
raise RuntimeError("consumer is closed")
if offsets.keys() - self.__offsets.keys():
raise ConsumerError("cannot seek on unassigned partitions")
self.__validate_offsets(offsets)
self.__offsets.update(offsets)
def stage_offsets(self, offsets: Mapping[Partition, int]) -> None:
if self.__state in {KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR}:
raise InvalidState(self.__state)
if offsets.keys() - self.__offsets.keys():
raise ConsumerError("cannot stage offsets for unassigned partitions")
self.__validate_offsets(offsets)
# TODO: Maybe log a warning if these offsets exceed the current
# offsets, since that's probably a side effect of an incorrect usage
# pattern?
self.__staged_offsets.update(offsets)
def seek(self, offsets: Mapping[Partition, int]) -> None:
"""
Change the read offsets for the provided partitions.
Raises an ``InvalidState`` if called on a closed consumer.
"""
if self.__state in {KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR}:
raise InvalidState(self.__state)
if offsets.keys() - self.__offsets.keys():
raise ConsumerError("cannot seek on unassigned partitions")
self.__seek(offsets)
def resume(self, partitions: Sequence[Partition]) -> None:
"""
Resume the consumption of messages for the provided partitions.
Raises an ``InvalidState`` if called on a closed consumer.
"""
if self.__state in {
KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR
}:
raise InvalidState(self.__state)
if set(partitions) - self.__offsets.keys():
raise ConsumerError("cannot resume unassigned partitions")
self.__consumer.resume([
ConfluentTopicPartition(partition.topic.name, partition.index)
for partition in partitions
])
for partition in partitions:
self.__paused.discard(partition)
def poll(
self,
timeout: Optional[float] = None
) -> Optional[Message[KafkaPayload]]:
"""
Return the next message available to be consumed, if one is
available. If no message is available, this method will block up to
the ``timeout`` value before returning ``None``. A timeout of
``0.0`` represents "do not block", while a timeout of ``None``
represents "block until a message is available (or forever)".
Calling this method may also invoke subscription state change
callbacks.
This method may also raise an ``EndOfPartition`` error (a subtype of
``ConsumerError``) when the consumer has reached the end of a
partition that it is subscribed to and no additional messages are
available. The ``partition`` attribute of the raised exception
specifies the end which partition has been reached. (Since this
consumer is multiplexing a set of partitions, this exception does not
mean that *all* of the partitions that the consumer is subscribed to
do not have any messages, just that it has reached the end of one of
them. This also does not mean that additional messages won't be
available in future poll calls.) Not every backend implementation
supports this feature or is configured to raise in this scenario.
Raises an ``InvalidState`` exception if called on a closed consumer.
Raises a ``TransportError`` for various other consumption-related
errors.
"""
if self.__state is not KafkaConsumerState.CONSUMING:
raise InvalidState(self.__state)
message: Optional[ConfluentMessage] = self.__consumer.poll(
*[timeout] if timeout is not None else [])
if message is None:
return None
error: Optional[KafkaError] = message.error()
if error is not None:
code = error.code()
if code == KafkaError._PARTITION_EOF:
raise EndOfPartition(
Partition(Topic(message.topic()), message.partition()),
message.offset(),
)
elif code == KafkaError._TRANSPORT:
raise TransportError(str(error))
elif code == KafkaError.OFFSET_OUT_OF_RANGE:
raise OffsetOutOfRange(str(error))
else:
raise ConsumerError(str(error))
headers: Optional[Headers] = message.headers()
result = Message(
Partition(Topic(message.topic()), message.partition()),
message.offset(),
KafkaPayload(
message.key(),
message.value(),
headers if headers is not None else [],
),
datetime.utcfromtimestamp(message.timestamp()[1] / 1000.0),
)
self.__offsets[result.partition] = result.next_offset
return result
def __resolve_partition_offset_error(
self,
partition: ConfluentTopicPartition) -> ConfluentTopicPartition:
raise ConsumerError("unable to resolve partition offsets")