Skip to content

Commit

Permalink
fix: ensure aiokafka commit is called with kafka.structs.TopicPartiti…
Browse files Browse the repository at this point in the history
…on (#539)
  • Loading branch information
dada-engineer committed Aug 14, 2023
1 parent 87a80a9 commit 3275ea8
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 4 deletions.
24 changes: 20 additions & 4 deletions faust/transport/drivers/aiokafka.py
Original file line number Diff line number Diff line change
Expand Up @@ -711,13 +711,15 @@ async def commit(self, offsets: Mapping[TP, int]) -> bool:
async def _commit(self, offsets: Mapping[TP, int]) -> bool:
consumer = self._ensure_consumer()
now = monotonic()
commitable_offsets = {
tp: offset for tp, offset in offsets.items() if tp in self.assignment()
}
try:
aiokafka_offsets = {
tp: OffsetAndMetadata(offset, "")
for tp, offset in offsets.items()
if tp in self.assignment()
ensure_aiokafka_TP(tp): OffsetAndMetadata(offset, "")
for tp, offset in commitable_offsets.items()
}
self.tp_last_committed_at.update({tp: now for tp in aiokafka_offsets})
self.tp_last_committed_at.update({tp: now for tp in commitable_offsets})
await consumer.commit(aiokafka_offsets)
except CommitFailedError as exc:
if "already rebalanced" in str(exc):
Expand Down Expand Up @@ -1621,3 +1623,17 @@ def credentials_to_aiokafka_auth(
}
else:
return {"security_protocol": "PLAINTEXT"}


def ensure_aiokafka_TP(tp: TP) -> _TopicPartition:
"""Convert Faust ``TP`` to aiokafka ``TopicPartition``."""
return (
tp
if isinstance(tp, _TopicPartition)
else _TopicPartition(tp.topic, tp.partition)
)


def ensure_aiokafka_TPset(tps: Iterable[TP]) -> Set[_TopicPartition]:
"""Convert set of Faust ``TP`` to aiokafka ``TopicPartition``."""
return {ensure_aiokafka_TP(tp) for tp in tps}
6 changes: 6 additions & 0 deletions tests/unit/transport/drivers/test_aiokafka.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
Transport,
credentials_to_aiokafka_auth,
server_list,
ensure_aiokafka_TPset,
)
from faust.types import TP
from faust.types.tuples import FutureMessage, PendingMessage
Expand Down Expand Up @@ -2038,3 +2039,8 @@ def test_credentials_to_aiokafka(credentials, ssl_context, expected):
def test_credentials_to_aiokafka__invalid():
with pytest.raises(ImproperlyConfigured):
credentials_to_aiokafka_auth(object())

def test_ensure_aiokafka_TPset():
actual = ensure_aiokafka_TPset({TP(topic="foo", partition=0)})
assert actual == {TopicPartition("foo", 0)}
assert all(isinstance(tp, TopicPartition) for tp in actual)

0 comments on commit 3275ea8

Please sign in to comment.