diff --git a/.coveragerc b/.coveragerc index 31cb9f69d..4f0ef0cc5 100644 --- a/.coveragerc +++ b/.coveragerc @@ -23,6 +23,10 @@ omit = */faust/assignor/* */faust/transport/drivers/memory.py + # optional driver: only importable/tested with the ckafka extra installed + # (faust[ckafka]); not part of the CI test environment. + */faust/transport/drivers/confluent.py + # tested by integration */faust/tables/recovery.py diff --git a/docs/conf.py b/docs/conf.py index 9763c0c43..c563f3628 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -54,6 +54,7 @@ 'faust.cli', 'faust.models', 'faust.serializers', + 'faust.transport.drivers.confluent', 'faust.types', 'faust.types._env', 'faust.utils', diff --git a/docs/includes/settingref.txt b/docs/includes/settingref.txt index 91bec6c02..669894cff 100644 --- a/docs/includes/settingref.txt +++ b/docs/includes/settingref.txt @@ -391,6 +391,15 @@ You can also pass a list of URLs: Limitations: None +- ``confluent://`` + + Experimental transport using the :pypi:`confluent-kafka` client. + + Limitations: Does not do sticky partition assignment (not + suitable for tables), and do not create any necessary internal + topics (you have to create them manually). + + .. setting:: broker_credentials ``broker_credentials`` diff --git a/extra/bandit/baseline.json b/extra/bandit/baseline.json index 4c9cc6b2d..affc80259 100644 --- a/extra/bandit/baseline.json +++ b/extra/bandit/baseline.json @@ -1070,6 +1070,18 @@ "loc": 890, "nosec": 0 }, + "faust/transport/drivers/confluent.py": { + "CONFIDENCE.HIGH": 0.0, + "CONFIDENCE.LOW": 0.0, + "CONFIDENCE.MEDIUM": 0.0, + "CONFIDENCE.UNDEFINED": 0.0, + "SEVERITY.HIGH": 0.0, + "SEVERITY.LOW": 0.0, + "SEVERITY.MEDIUM": 0.0, + "SEVERITY.UNDEFINED": 0.0, + "loc": 486, + "nosec": 0 + }, "faust/transport/producer.py": { "CONFIDENCE.HIGH": 0.0, "CONFIDENCE.LOW": 0.0, @@ -2043,4 +2055,4 @@ "test_name": "blacklist" } ] -} \ No newline at end of file +} diff --git a/faust/transport/drivers/__init__.py b/faust/transport/drivers/__init__.py index 8deec0196..4891a75fb 100644 --- a/faust/transport/drivers/__init__.py +++ b/faust/transport/drivers/__init__.py @@ -1,22 +1,18 @@ """Transport registry.""" -from yarl import URL +from typing import Type -from .aiokafka import Transport as AIOKafkaTransport +from mode.utils.imports import FactoryMapping -__all__ = ["by_name", "by_url"] - - -DRIVERS = { - "aiokafka": AIOKafkaTransport, - "kafka": AIOKafkaTransport, -} - - -def by_name(driver_name: str): - return DRIVERS[driver_name] +from faust.types import TransportT +__all__ = ["by_name", "by_url"] -def by_url(url: URL): - scheme = url.scheme - return DRIVERS[scheme] +TRANSPORTS: FactoryMapping[Type[TransportT]] = FactoryMapping( + aiokafka="faust.transport.drivers.aiokafka:Transport", + confluent="faust.transport.drivers.confluent:Transport", + kafka="faust.transport.drivers.aiokafka:Transport", +) +TRANSPORTS.include_setuptools_namespace("faust.transports") +by_name = TRANSPORTS.by_name +by_url = TRANSPORTS.by_url diff --git a/faust/transport/drivers/confluent.py b/faust/transport/drivers/confluent.py new file mode 100644 index 000000000..2573a1545 --- /dev/null +++ b/faust/transport/drivers/confluent.py @@ -0,0 +1,695 @@ +"""Message transport using :pypi:`confluent_kafka`.""" + +import asyncio +import typing +from collections import defaultdict +from time import monotonic +from typing import ( + Any, + Awaitable, + Callable, + ClassVar, + Iterable, + List, + Mapping, + MutableMapping, + Optional, + Set, + Type, + cast, +) + +import confluent_kafka +from confluent_kafka import KafkaException, TopicPartition as _TopicPartition +from confluent_kafka.admin import AdminClient +from mode import Service, get_logger +from mode.threads import QueueServiceThread +from mode.utils.futures import notify +from mode.utils.times import Seconds, want_seconds +from yarl import URL + +from faust.exceptions import ConsumerNotStarted, ProducerSendError +from faust.transport import base +from faust.transport.consumer import ( + ConsumerThread, + RecordMap, + ThreadDelegateConsumer, + ensure_TP, + ensure_TPset, +) +from faust.types import TP, AppT, ConsumerMessage, HeadersArg, RecordMetadata +from faust.types.transports import ConsumerT, ProducerT + +if typing.TYPE_CHECKING: + from confluent_kafka import ( + Consumer as _Consumer, + Message as _Message, + Producer as _Producer, + ) +else: + + class _Consumer: ... # noqa + + class _Producer: ... # noqa + + class _Message: ... # noqa + + +__all__ = ["Consumer", "Producer", "Transport"] + + +logger = get_logger(__name__) + + +def server_list(urls: List[URL], default_port: int) -> str: + default_host = "127.0.0.1" + return ",".join( + [f"{u.host or default_host}:{u.port or default_port}" for u in urls] + ) + + +class Consumer(ThreadDelegateConsumer): + """Kafka consumer using :pypi:`confluent_kafka`.""" + + logger = logger + + def _new_consumer_thread(self) -> ConsumerThread: + return ConfluentConsumerThread(self, loop=self.loop, beacon=self.beacon) + + async def create_topic( + self, + topic: str, + partitions: int, + replication: int, + *, + config: Mapping[str, Any] = None, + timeout: Seconds = 30.0, + retention: Seconds = None, + compacting: bool = None, + deleting: bool = None, + ensure_created: bool = False, + ) -> None: + """Create topic on broker.""" + if self.app.conf.topic_allow_declare: + await self._thread.create_topic( + topic, + partitions, + replication, + config=config, + timeout=int(want_seconds(timeout) * 1000.0), + retention=( + int(want_seconds(retention) * 1000.0) + if retention is not None + else None + ), + compacting=compacting, + deleting=deleting, + ensure_created=ensure_created, + ) + else: + logger.warning(f"Topic creation disabled! Can't create topic {topic}") + + def _to_message(self, tp: TP, record: Any) -> ConsumerMessage: + # convert timestamp to seconds from int milliseconds. + timestamp_type: int + timestamp: Optional[int] + timestamp_type, timestamp = record.timestamp() + timestamp_s: float = cast(float, None) + if timestamp is not None: + timestamp_s = timestamp / 1000.0 + key = record.key() + key_size = len(key) if key is not None else 0 + value = record.value() + value_size = len(value) if value is not None else 0 + return ConsumerMessage( + record.topic(), + record.partition(), + record.offset(), + timestamp_s, + timestamp_type, + [], # headers + key, + value, + None, + key_size, + value_size, + tp, + ) + + def _new_topicpartition(self, topic: str, partition: int) -> TP: + return cast(TP, _TopicPartition(topic, partition)) + + async def on_stop(self) -> None: + """Call when consumer is stopping.""" + await super().on_stop() + + def verify_event_path(self, now: float, tp: TP) -> None: + return self._thread.verify_event_path(now, tp) + + +class AsyncConsumer: + def __init__( + self, + config, + logger=None, + callback=None, + loop=None, + on_partitions_revoked=None, + on_partitions_assigned=None, + beacon=None, + ): + """Construct a Consumer usable within asyncio. + + :param config: A configuration dict for this Consumer + :param logger: A python logger instance. + """ + self.consumer = confluent_kafka.Consumer(**config) + self.callback = callback + self.on_partitions_revoked = on_partitions_revoked + self.on_partitions_assigned = on_partitions_assigned + self.beacon = beacon + self.loop = loop + + def close(self): + self.consumer.close() + + def subscribe(self, *args, **kwargs): + self.consumer.subscribe(*args, **kwargs) + + def assign(self, *args, **kwargs): + self.consumer.assign(*args, **kwargs) + + async def poll(self, timeout=1.0): + """Consume a single message and call the registered callback. + + Delegates to the blocking :meth:`confluent_kafka.Consumer.poll`. + This is always invoked from within the dedicated consumer thread's + event loop (see :class:`ConfluentConsumerThread`), so parking on the + blocking call only affects that thread -- never the application event + loop -- mirroring how :meth:`ConfluentConsumerThread.getmany` drives + :meth:`confluent_kafka.Consumer.consume`. + + Returns ``None`` if no message arrives within ``timeout`` seconds. + """ + if not timeout or timeout <= 0: + timeout = 1.0 + msg = self.consumer.poll(timeout=timeout) + if msg is not None and self.callback is not None: + await self.callback(msg) + return msg + + def assignment(self) -> Set[TP]: + return self.consumer.assignment() + + +class ConfluentConsumerThread(ConsumerThread): + """Thread managing underlying :pypi:`confluent_kafka` consumer.""" + + _consumer: Optional[AsyncConsumer] = None + _assigned: bool = False + + # _pending_rebalancing_spans: Deque[opentracing.Span] + + tp_last_committed_at: MutableMapping[TP, float] + time_started: float + + tp_fetch_request_timeout_secs: float + tp_fetch_response_timeout_secs: float + tp_stream_timeout_secs: float + tp_commit_timeout_secs: float + + async def on_start(self) -> None: + self._consumer = self._create_consumer(loop=self.thread_loop) + self.time_started = monotonic() + # await self._consumer.start() + + def _create_consumer(self, loop: asyncio.AbstractEventLoop) -> AsyncConsumer: + transport = cast(Transport, self.transport) + if self.app.client_only: + return self._create_client_consumer(transport, loop=loop) + else: + return self._create_worker_consumer(transport, loop=loop) + + def _create_worker_consumer( + self, transport: "Transport", loop: asyncio.AbstractEventLoop + ) -> AsyncConsumer: + conf = self.app.conf + self._assignor = self.app.assignor + + # XXX parition.assignment.strategy is string + # need to write C wrapper for this + # 'partition.assignment.strategy': [self._assignor] + return AsyncConsumer( + { + "bootstrap.servers": server_list(transport.url, transport.default_port), + "group.id": conf.id, + "client.id": conf.broker_client_id, + "default.topic.config": { + "auto.offset.reset": "earliest", + }, + "enable.auto.commit": False, + "fetch.max.bytes": conf.consumer_max_fetch_size, + "request.timeout.ms": int(conf.broker_request_timeout * 1000.0), + "check.crcs": conf.broker_check_crcs, + "session.timeout.ms": int(conf.broker_session_timeout * 1000.0), + "heartbeat.interval.ms": int(conf.broker_heartbeat_interval * 1000.0), + }, + self.logger, + ) + + def _create_client_consumer( + self, transport: "Transport", loop: asyncio.AbstractEventLoop + ) -> AsyncConsumer: + conf = self.app.conf + return AsyncConsumer( + { + "bootstrap.servers": server_list(transport.url, transport.default_port), + "client.id": conf.broker_client_id, + "enable.auto.commit": True, + "default.topic.config": { + "auto.offset.reset": "earliest", + }, + }, + self.logger, + ) + + def close(self) -> None: ... + + async def subscribe(self, topics: Iterable[str]) -> None: + # XXX pattern does not work :/ + await self.cast_thread( + self._ensure_consumer().subscribe, + topics=list(topics), + on_assign=self._on_assign, + on_revoke=self._on_revoke, + # listener=self._rebalance_listener, + ) + + def _on_assign(self, consumer: _Consumer, assigned: List[_TopicPartition]) -> None: + self._assigned = True + self.thread_loop.create_task( + self.on_partitions_assigned({TP(tp.topic, tp.partition) for tp in assigned}) + ) + + def _on_revoke(self, consumer: _Consumer, revoked: List[_TopicPartition]) -> None: + self.thread_loop.create_task( + self.on_partitions_revoked({TP(tp.topic, tp.partition) for tp in revoked}) + ) + + async def seek_to_committed(self) -> Mapping[TP, int]: + return await self.call_thread(self._seek_to_committed) + + async def _seek_to_committed(self) -> Mapping[TP, int]: + consumer = self._ensure_consumer() + assignment = consumer.assignment() + committed = consumer.consumer.committed(assignment) + for tp in committed: + await consumer.consumer.seek(tp) + return {ensure_TP(tp): tp.offset for tp in committed} + + async def _committed_offsets(self, partitions: List[TP]) -> MutableMapping[TP, int]: + consumer = self._ensure_consumer() + committed = consumer.consumer.committed( + [_TopicPartition(tp[0], tp[1]) for tp in partitions] + ) + return {TP(tp.topic, tp.partition): tp.offset for tp in committed} + + async def commit(self, tps: Mapping[TP, int]) -> bool: + await self.call_thread( + self._ensure_consumer().consumer.commit, + offsets=[ + _TopicPartition(tp.topic, tp.partition, offset=offset) + for tp, offset in tps.items() + ], + asynchronous=False, + ) + return True + + async def position(self, tp: TP) -> Optional[int]: + return await self.call_thread(self._ensure_consumer().consumer.position, tp) + + async def seek_to_beginning(self, *partitions: _TopicPartition) -> None: + await self.call_thread( + self._ensure_consumer().consumer.seek_to_beginning, *partitions + ) + + async def seek_wait(self, partitions: Mapping[TP, int]) -> None: + consumer = self._ensure_consumer() + await self.call_thread(self._seek_wait, consumer, partitions) + + async def _seek_wait( + self, consumer: Consumer, partitions: Mapping[TP, int] + ) -> None: + for tp, offset in partitions.items(): + self.log.dev("SEEK %r -> %r", tp, offset) + await consumer.seek(tp, offset) + await asyncio.gather(*[consumer.position(tp) for tp in partitions]) + + def seek(self, partition: TP, offset: int) -> None: + self._ensure_consumer().consumer.seek(partition, offset) + + def assignment(self) -> Set[TP]: + return ensure_TPset(self._ensure_consumer().assignment()) + + def highwater(self, tp: TP) -> int: + _, hw = self._ensure_consumer().consumer.get_watermark_offsets( + _TopicPartition(tp.topic, tp.partition), cached=True + ) + return hw + + def topic_partitions(self, topic: str) -> Optional[int]: + # XXX NotImplemented + return None + + async def earliest_offsets(self, *partitions: TP) -> MutableMapping[TP, int]: + if not partitions: + return {} + return await self.call_thread(self._earliest_offsets, partitions) + + async def _earliest_offsets(self, partitions: List[TP]) -> MutableMapping[TP, int]: + consumer = self._ensure_consumer() + return { + tp: consumer.consumer.get_watermark_offsets(_TopicPartition(tp[0], tp[1]))[ + 0 + ] + for tp in partitions + } + + async def highwaters(self, *partitions: TP) -> MutableMapping[TP, int]: + if not partitions: + return {} + return await self.call_thread(self._highwaters, partitions) + + async def _highwaters(self, partitions: List[TP]) -> MutableMapping[TP, int]: + consumer = self._ensure_consumer() + return { + tp: consumer.consumer.get_watermark_offsets(_TopicPartition(tp[0], tp[1]))[ + 1 + ] + for tp in partitions + } + + def _ensure_consumer(self) -> AsyncConsumer: + if self._consumer is None: + raise ConsumerNotStarted("Consumer thread not yet started") + return self._consumer + + async def getmany( + self, active_partitions: Optional[Set[TP]], timeout: float + ) -> RecordMap: + # Implementation for the Fetcher service. + _consumer = self._ensure_consumer() + messages = await self.call_thread( + _consumer.consumer.consume, + num_messages=10000, + timeout=timeout, + ) + records: RecordMap = defaultdict(list) + for message in messages: + tp = TP(message.topic(), message.partition()) + records[tp].append(message) + return records + + async def create_topic( + self, + topic: str, + partitions: int, + replication: int, + *, + config: Mapping[str, Any] = None, + timeout: Seconds = 30.0, + retention: Seconds = None, + compacting: bool = None, + deleting: bool = None, + ensure_created: bool = False, + ) -> None: + return # XXX + + def key_partition( + self, topic: str, key: Optional[bytes], partition: int = None + ) -> Optional[int]: + metadata = self._consumer.consumer.list_topics(topic) + partition_count = len(metadata.topics[topic]["partitions"]) + + # Calculate the partition number based on the key hash + key_bytes = str(key).encode("utf-8") + return abs(hash(key_bytes)) % partition_count + + +class ProducerProduceFuture(asyncio.Future): + def set_from_on_delivery(self, err: Optional[BaseException], msg: _Message) -> None: + if err: + # XXX Not sure what err' is here, hopefully it's an exception + # object and not a string [ask]. + self.set_exception(err) + else: + metadata: RecordMetadata = self.message_to_metadata(msg) + self.set_result(metadata) + + def message_to_metadata(self, message: _Message) -> RecordMetadata: + topic, partition = tp = TP(message.topic(), message.partition()) + return RecordMetadata(topic, partition, tp, message.offset()) + + +class ProducerThread(QueueServiceThread): + """Thread managing underlying :pypi:`confluent_kafka` producer.""" + + app: AppT + producer: "Producer" + transport: "Transport" + _producer: Optional[_Producer] = None + _flush_soon: Optional[asyncio.Future] = None + + def __init__(self, producer: "Producer", **kwargs: Any) -> None: + self.producer = producer + self.transport = cast(Transport, self.producer.transport) + self.app = self.transport.app + super().__init__(**kwargs) + + async def on_start(self) -> None: + self._producer = confluent_kafka.Producer( + { + "bootstrap.servers": server_list( + self.transport.url, self.transport.default_port + ), + "client.id": self.app.conf.broker_client_id, + "max.in.flight.requests.per.connection": 1, + } + ) + + async def flush(self) -> None: + if self._producer is not None: + self._producer.flush() + + async def on_thread_stop(self) -> None: + if self._producer is not None: + self._producer.flush() + + def produce( + self, + topic: str, + key: bytes, + value: bytes, + partition: int, + on_delivery: Callable, + ) -> None: + if self._producer is None: + raise RuntimeError("Producer not started") + if partition is not None: + self._producer.produce( + topic, + key, + value, + partition, + on_delivery=on_delivery, + ) + else: + self._producer.produce( + topic, + key, + value, + on_delivery=on_delivery, + ) + notify(self._flush_soon) + + @Service.task + async def _background_flush(self) -> None: + producer = cast(_Producer, self._producer) + _size = producer.__len__ + _flush = producer.flush + _poll = producer.poll + _sleep = self.sleep + _create_future = self.loop.create_future + while not self.should_stop: + if not _size(): + flush_soon = self._flush_soon + if flush_soon is None: + flush_soon = self._flush_soon = _create_future() + stopped: bool = False + try: + stopped = await self.wait_for_stopped(flush_soon, timeout=1.0) + finally: + self._flush_soon = None + if not stopped: + _flush(timeout=100) + _poll(timeout=1) + await _sleep(0) + + +class Producer(base.Producer): + """Kafka producer using :pypi:`confluent_kafka`.""" + + logger = logger + + _producer_thread: ProducerThread + _admin: AdminClient + _quick_produce: Any = None + + def __post_init__(self) -> None: + self._producer_thread = ProducerThread(self, loop=self.loop, beacon=self.beacon) + self._quick_produce = self._producer_thread.produce + + async def _on_irrecoverable_error(self, exc: BaseException) -> None: + consumer = self.transport.app.consumer + if consumer is not None: + await consumer.crash(exc) + await self.crash(exc) + + async def on_restart(self) -> None: + """Call when producer is restarting.""" + self.on_init() + + async def create_topic( + self, + topic: str, + partitions: int, + replication: int, + *, + config: Mapping[str, Any] = None, + timeout: Seconds = 20.0, + retention: Seconds = None, + compacting: bool = None, + deleting: bool = None, + ensure_created: bool = False, + ) -> None: + """Create topic on broker.""" + return # XXX + _retention = int(want_seconds(retention) * 1000.0) if retention else None + await cast(Transport, self.transport)._create_topic( + self, + self._producer.client, + topic, + partitions, + replication, + config=config, + timeout=int(want_seconds(timeout) * 1000.0), + retention=_retention, + compacting=compacting, + deleting=deleting, + ensure_created=ensure_created, + ) + + async def on_start(self) -> None: + """Call when producer is starting.""" + await self._producer_thread.start() + await self.sleep(0.5) # cannot remember why, necessary? [ask] + + async def on_stop(self) -> None: + """Call when producer is stopping.""" + await self._producer_thread.stop() + + async def send( + self, + topic: str, + key: Optional[bytes], + value: Optional[bytes], + partition: Optional[int], + timestamp: Optional[float], + headers: Optional[HeadersArg], + *, + transactional_id: str = None, + ) -> Awaitable[RecordMetadata]: + """Send message for future delivery.""" + fut = ProducerProduceFuture(loop=self.loop) + self._quick_produce( + topic, + value, + key, + partition, + on_delivery=fut.set_from_on_delivery, + ) + return cast(Awaitable[RecordMetadata], fut) + try: + return cast( + Awaitable[RecordMetadata], + await self._producer.send(topic, value, key=key, partition=partition), + ) + except KafkaException as exc: + raise ProducerSendError(f"Error while sending: {exc!r}") from exc + + async def send_and_wait( + self, + topic: str, + key: Optional[bytes], + value: Optional[bytes], + partition: Optional[int], + timestamp: Optional[float], + headers: Optional[HeadersArg], + *, + transactional_id: str = None, + ) -> RecordMetadata: + """Send message and wait for it to be delivered to broker(s).""" + fut = await self.send( + topic, + key, + value, + partition, + timestamp, + headers, + ) + return await fut + + async def flush(self) -> None: + """Flush producer buffer. + + This will wait until the producer has written + all buffered up messages to any connected brokers. + """ + await self._producer_thread.flush() + + def key_partition(self, topic: str, key: bytes) -> TP: + """Return topic and partition destination for key.""" + # Get the partition count for the topic + metadata = self._producer_thread.producer.list_topics(topic) + partition_count = len(metadata.topics[topic].partitions) + + # Calculate the partition number based on the key hash + key_bytes = str(key).encode("utf-8") + partition = abs(hash(key_bytes)) % partition_count + + return TP(topic, partition) + + +class Transport(base.Transport): + """Kafka transport using :pypi:`confluent_kafka`.""" + + Consumer: ClassVar[Type[ConsumerT]] = Consumer + Producer: ClassVar[Type[ProducerT]] = Producer + + default_port = 9092 + driver_version = f"confluent_kafka={confluent_kafka.__version__}" + + def _topic_config( + self, retention: int = None, compacting: bool = None, deleting: bool = None + ) -> MutableMapping[str, Any]: + config: MutableMapping[str, Any] = {} + cleanup_flags: Set[str] = set() + if compacting: + cleanup_flags |= {"compact"} + if deleting: + cleanup_flags |= {"delete"} + if cleanup_flags: + config["cleanup.policy"] = ",".join(sorted(cleanup_flags)) + if retention: + config["retention.ms"] = retention + return config diff --git a/faust/types/settings/settings.py b/faust/types/settings/settings.py index 2fbc84d17..9a259b8dd 100644 --- a/faust/types/settings/settings.py +++ b/faust/types/settings/settings.py @@ -703,6 +703,15 @@ def broker(self) -> List[URL]: The recommended transport using the :pypi:`aiokafka` client. Limitations: None + + + - ``confluent://`` + + Experimental transport using the :pypi:`confluent-kafka` client. + + Limitations: Does not do sticky partition assignment (not + suitable for tables), and do not create any necessary internal + topics (you have to create them manually). """ @broker.on_set_default # type: ignore diff --git a/requirements/extras/ckafka.txt b/requirements/extras/ckafka.txt new file mode 100644 index 000000000..42edc695d --- /dev/null +++ b/requirements/extras/ckafka.txt @@ -0,0 +1 @@ +confluent-kafka>=2.0.0 diff --git a/setup.py b/setup.py index e8047d085..6608fba62 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,7 @@ "aiomonitor", "cchardet", "ciso8601", + "ckafka", "cython", "datadog", "debug", diff --git a/tests/unit/transport/drivers/test_confluent.py b/tests/unit/transport/drivers/test_confluent.py new file mode 100644 index 000000000..2b58d0d87 --- /dev/null +++ b/tests/unit/transport/drivers/test_confluent.py @@ -0,0 +1,564 @@ +"""Unit tests for the :pypi:`confluent_kafka` transport driver. + +These tests exercise the driver's real surface -- the thin async wrappers +around :pypi:`confluent_kafka` -- with the underlying ``confluent_kafka`` +``Consumer``/``Producer`` mocked out, so no broker is required. +""" + +from unittest.mock import Mock, patch + +import pytest + +# confluent-kafka is an optional extra (faust[ckafka]); skip the whole module +# when it -- and therefore the confluent driver -- is not installed. +pytest.importorskip("confluent_kafka") + +from confluent_kafka import TopicPartition # noqa: E402 +from mode.utils.futures import done_future # noqa: E402 + +from faust.transport.drivers import confluent as mod # noqa: E402 +from faust.transport.drivers.confluent import ( # noqa: E402 + AsyncConsumer, + ConfluentConsumerThread, + Consumer, + ConsumerNotStarted, + Producer, + ProducerProduceFuture, + ProducerThread, + Transport, + server_list, +) +from faust.types import TP # noqa: E402 +from faust.types.tuples import RecordMetadata # noqa: E402 +from tests.helpers import AsyncMock # noqa: E402 + +TP1 = TP("topic", 0) +TP2 = TP("topic", 3) +TP3 = TP("topix", 1) + +TESTED_MODULE = "faust.transport.drivers.confluent" + + +async def _passthrough_call_thread(fn, *args, **kwargs): + """Stand-in for ConsumerThread.call_thread that runs inline. + + The real implementation dispatches ``fn`` onto the consumer's dedicated + thread; for unit tests we just await/run it in place. + """ + result = fn(*args, **kwargs) + if hasattr(result, "__await__"): + return await result + return result + + +@pytest.fixture() +def callback(): + return AsyncMock(name="callback") + + +@pytest.fixture() +def on_partitions_revoked(): + return AsyncMock(name="on_partitions_revoked") + + +@pytest.fixture() +def on_partitions_assigned(): + return AsyncMock(name="on_partitions_assigned") + + +@pytest.fixture() +def consumer(app, callback, on_partitions_revoked, on_partitions_assigned): + consumer = Consumer( + app.transport, + callback=callback, + on_partitions_revoked=on_partitions_revoked, + on_partitions_assigned=on_partitions_assigned, + ) + consumer._thread = Mock( + name="thread", + create_topic=AsyncMock(), + verify_event_path=Mock(), + ) + return consumer + + +@pytest.fixture() +def cthread(consumer): + return ConfluentConsumerThread(consumer) + + +@pytest.fixture() +def underlying(): + """A mock of the low-level confluent_kafka.Consumer.""" + return Mock(name="confluent_kafka.Consumer") + + +@pytest.fixture() +def _consumer(underlying): + """A mock AsyncConsumer as held by ConfluentConsumerThread._consumer.""" + _consumer = Mock(name="AsyncConsumer") + _consumer.consumer = underlying + return _consumer + + +class Test_server_list: + def test_single(self): + from yarl import URL + + assert server_list([URL("kafka://localhost:9092")], 9092) == "localhost:9092" + + def test_multiple_and_default_port(self): + from yarl import URL + + urls = [URL("kafka://h1:9092"), URL("kafka://h2")] + assert server_list(urls, 9092) == "h1:9092,h2:9092" + + +class TestConsumer: + def test__new_consumer_thread(self, *, consumer): + thread = Consumer._new_consumer_thread(consumer) + assert isinstance(thread, ConfluentConsumerThread) + + @pytest.mark.asyncio + async def test_create_topic__allowed(self, *, consumer): + consumer.app.conf.topic_allow_declare = True + await consumer.create_topic( + "topic", 3, 1, retention=3000.0, compacting=True, deleting=False + ) + consumer._thread.create_topic.assert_called_once_with( + "topic", + 3, + 1, + config=None, + timeout=30000, + retention=3000000, + compacting=True, + deleting=False, + ensure_created=False, + ) + + @pytest.mark.asyncio + async def test_create_topic__default_retention(self, *, consumer): + # retention defaults to None; must not crash on int(None * 1000). + consumer.app.conf.topic_allow_declare = True + await consumer.create_topic("topic", 3, 1) + _, kwargs = consumer._thread.create_topic.call_args + assert kwargs["retention"] is None + + @pytest.mark.asyncio + async def test_create_topic__not_allowed(self, *, consumer): + consumer.app.conf.topic_allow_declare = False + with patch(TESTED_MODULE + ".logger") as logger: + await consumer.create_topic("topic", 3, 1, retention=3000.0) + consumer._thread.create_topic.assert_not_called() + logger.warning.assert_called_once() + + def test__to_message(self, *, consumer): + record = Mock(name="record") + record.timestamp.return_value = (1, 1000) + record.key.return_value = b"key" + record.value.return_value = b"value" + record.topic.return_value = "topic" + record.partition.return_value = 0 + record.offset.return_value = 7 + + msg = consumer._to_message(TP1, record) + + assert msg.topic == "topic" + assert msg.partition == 0 + assert msg.offset == 7 + assert msg.key == b"key" + assert msg.value == b"value" + assert msg.timestamp == 1.0 + assert msg.timestamp_type == 1 + assert msg.tp == TP1 + + def test__to_message__no_timestamp_no_key(self, *, consumer): + record = Mock(name="record") + record.timestamp.return_value = (0, None) + record.key.return_value = None + record.value.return_value = None + record.topic.return_value = "topic" + record.partition.return_value = 0 + record.offset.return_value = 0 + + msg = consumer._to_message(TP1, record) + + assert msg.key is None + assert msg.value is None + + def test__new_topicpartition(self, *, consumer): + tp = consumer._new_topicpartition("topic", 3) + assert tp.topic == "topic" + assert tp.partition == 3 + + def test_verify_event_path(self, *, consumer): + consumer.verify_event_path(303.3, TP1) + consumer._thread.verify_event_path.assert_called_once_with(303.3, TP1) + + +class TestAsyncConsumer: + @pytest.fixture() + def async_consumer(self, callback): + with patch.object(mod.confluent_kafka, "Consumer") as MockConsumer: + ac = AsyncConsumer({"group.id": "g"}, callback=callback) + assert ac.consumer is MockConsumer.return_value + yield ac + + def test_construct_does_not_call_io_event_enable(self): + with patch.object(mod.confluent_kafka, "Consumer") as MockConsumer: + AsyncConsumer({"group.id": "g"}) + # io_event_enable was the abandoned-fork-only API; it must not be used. + assert not MockConsumer.return_value.io_event_enable.called + + def test_close(self, *, async_consumer): + async_consumer.close() + async_consumer.consumer.close.assert_called_once_with() + + def test_subscribe(self, *, async_consumer): + async_consumer.subscribe(["topic"], on_assign=1) + async_consumer.consumer.subscribe.assert_called_once_with( + ["topic"], on_assign=1 + ) + + def test_assign(self, *, async_consumer): + async_consumer.assign([TP1]) + async_consumer.consumer.assign.assert_called_once_with([TP1]) + + def test_assignment(self, *, async_consumer): + async_consumer.consumer.assignment.return_value = {TP1} + assert async_consumer.assignment() == {TP1} + + @pytest.mark.asyncio + async def test_poll__message(self, *, async_consumer, callback): + message = Mock(name="message") + async_consumer.consumer.poll.return_value = message + result = await async_consumer.poll(timeout=5.0) + assert result is message + async_consumer.consumer.poll.assert_called_once_with(timeout=5.0) + callback.assert_called_once_with(message) + + @pytest.mark.asyncio + async def test_poll__no_message(self, *, async_consumer, callback): + async_consumer.consumer.poll.return_value = None + result = await async_consumer.poll(timeout=5.0) + assert result is None + callback.assert_not_called() + + @pytest.mark.asyncio + async def test_poll__coerces_nonpositive_timeout(self, *, async_consumer): + async_consumer.consumer.poll.return_value = None + await async_consumer.poll(timeout=0) + async_consumer.consumer.poll.assert_called_once_with(timeout=1.0) + + +class TestConfluentConsumerThread: + def test__ensure_consumer__not_started(self, *, cthread): + cthread._consumer = None + with pytest.raises(ConsumerNotStarted): + cthread._ensure_consumer() + + def test__ensure_consumer(self, *, cthread, _consumer): + cthread._consumer = _consumer + assert cthread._ensure_consumer() is _consumer + + def test__create_worker_consumer(self, *, cthread): + transport = cthread.transport + with patch(TESTED_MODULE + ".AsyncConsumer") as AC: + result = cthread._create_worker_consumer(transport, loop=None) + assert result is AC.return_value + config = AC.call_args[0][0] + assert config["group.id"] == cthread.app.conf.id + assert config["enable.auto.commit"] is False + assert "bootstrap.servers" in config + assert cthread._assignor is cthread.app.assignor + + def test__create_client_consumer(self, *, cthread): + transport = cthread.transport + with patch(TESTED_MODULE + ".AsyncConsumer") as AC: + result = cthread._create_client_consumer(transport, loop=None) + assert result is AC.return_value + config = AC.call_args[0][0] + assert config["enable.auto.commit"] is True + assert "bootstrap.servers" in config + + @pytest.mark.asyncio + async def test_getmany(self, *, cthread, _consumer): + cthread._consumer = _consumer + m1, m2, m3 = (self._message(TP1), self._message(TP1), self._message(TP3)) + cthread.call_thread = AsyncMock(return_value=[m1, m2, m3]) + + records = await cthread.getmany({TP1}, timeout=1.0) + + cthread.call_thread.assert_called_once_with( + _consumer.consumer.consume, num_messages=10000, timeout=1.0 + ) + assert records[TP1] == [m1, m2] + assert records[TP3] == [m3] + + @staticmethod + def _message(tp): + message = Mock(name="message") + message.topic.return_value = tp.topic + message.partition.return_value = tp.partition + return message + + @pytest.mark.asyncio + async def test_commit(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.call_thread = AsyncMock() + assert await cthread.commit({TP1: 100}) is True + cthread.call_thread.assert_called_once() + + @pytest.mark.asyncio + async def test_position(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.call_thread = AsyncMock(return_value=42) + assert await cthread.position(TP1) == 42 + cthread.call_thread.assert_called_once_with(_consumer.consumer.position, TP1) + + def test_assignment(self, *, cthread, _consumer): + cthread._consumer = _consumer + _consumer.assignment.return_value = {TopicPartition(TP1.topic, TP1.partition)} + assert cthread.assignment() == {TP1} + + def test_highwater(self, *, cthread, _consumer): + cthread._consumer = _consumer + _consumer.consumer.get_watermark_offsets.return_value = (5, 13) + assert cthread.highwater(TP1) == 13 + + def test_seek(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.seek(TP1, 300) + _consumer.consumer.seek.assert_called_once_with(TP1, 300) + + @pytest.mark.asyncio + async def test_seek_to_beginning(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.call_thread = AsyncMock() + part = TopicPartition(TP1.topic, TP1.partition) + await cthread.seek_to_beginning(part) + cthread.call_thread.assert_called_once_with( + _consumer.consumer.seek_to_beginning, part + ) + + @pytest.mark.asyncio + async def test_earliest_offsets(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.call_thread = _passthrough_call_thread + _consumer.consumer.get_watermark_offsets.return_value = (2, 20) + offsets = await cthread.earliest_offsets(TP1) + assert offsets[TP1] == 2 + + @pytest.mark.asyncio + async def test_highwaters(self, *, cthread, _consumer): + cthread._consumer = _consumer + cthread.call_thread = _passthrough_call_thread + _consumer.consumer.get_watermark_offsets.return_value = (2, 20) + offsets = await cthread.highwaters(TP1) + assert offsets[TP1] == 20 + + def test__on_assign(self, *, cthread): + cthread.thread_loop = Mock(name="thread_loop") + cthread.on_partitions_assigned = Mock(name="on_partitions_assigned") + assigned = [TopicPartition(TP1.topic, TP1.partition)] + cthread._on_assign(Mock(name="consumer"), assigned) + assert cthread._assigned is True + cthread.on_partitions_assigned.assert_called_once_with({TP1}) + cthread.thread_loop.create_task.assert_called_once() + + def test__on_revoke(self, *, cthread): + cthread.thread_loop = Mock(name="thread_loop") + cthread.on_partitions_revoked = Mock(name="on_partitions_revoked") + revoked = [TopicPartition(TP1.topic, TP1.partition)] + cthread._on_revoke(Mock(name="consumer"), revoked) + cthread.on_partitions_revoked.assert_called_once_with({TP1}) + cthread.thread_loop.create_task.assert_called_once() + + def test_key_partition(self, *, cthread, _consumer): + cthread._consumer = _consumer + metadata = Mock(name="metadata") + metadata.topics = {"topic": {"partitions": {0: 1, 1: 1, 2: 1}}} + _consumer.consumer.list_topics.return_value = metadata + partition = cthread.key_partition("topic", b"key") + assert 0 <= partition < 3 + + def test_topic_partitions(self, *, cthread): + assert cthread.topic_partitions("topic") is None + + +class TestProducer: + @pytest.fixture() + def producer(self, app): + producer = Producer(app.transport) + producer._producer_thread = Mock( + name="producer_thread", + start=AsyncMock(), + stop=AsyncMock(), + flush=AsyncMock(), + ) + producer.sleep = AsyncMock() + return producer + + def test___post_init__(self, *, producer): + assert producer._producer_thread is not None + assert producer._quick_produce is not None + + @pytest.mark.asyncio + async def test_on_start(self, *, producer): + await producer.on_start() + producer._producer_thread.start.assert_called_once_with() + + @pytest.mark.asyncio + async def test_on_stop(self, *, producer): + await producer.on_stop() + producer._producer_thread.stop.assert_called_once_with() + + @pytest.mark.asyncio + async def test_send(self, *, producer): + producer._quick_produce = Mock(name="quick_produce") + fut = await producer.send("topic", b"key", b"value", 0, None, None) + assert isinstance(fut, ProducerProduceFuture) + producer._quick_produce.assert_called_once() + args, kwargs = producer._quick_produce.call_args + assert args[0] == "topic" + assert "on_delivery" in kwargs + + @pytest.mark.asyncio + async def test_send_and_wait(self, *, producer, event_loop): + expected = Mock(name="metadata") + producer.send = AsyncMock(return_value=done_future(expected, loop=event_loop)) + result = await producer.send_and_wait("topic", b"k", b"v", 0, None, None) + assert result is expected + + @pytest.mark.asyncio + async def test_flush(self, *, producer): + await producer.flush() + producer._producer_thread.flush.assert_called_once_with() + + @pytest.mark.asyncio + async def test_create_topic__is_noop(self, *, producer): + # create_topic short-circuits (XXX) -- must not raise. + assert await producer.create_topic("topic", 3, 1) is None + + def test_key_partition(self, *, producer): + metadata = Mock(name="metadata") + topic_meta = Mock() + topic_meta.partitions = {0: 1, 1: 1} + metadata.topics = {"topic": topic_meta} + producer._producer_thread.producer = Mock() + producer._producer_thread.producer.list_topics.return_value = metadata + tp = producer.key_partition("topic", b"key") + assert tp.topic == "topic" + assert 0 <= tp.partition < 2 + + +class TestProducerThread: + @pytest.fixture() + def producer_thread(self, app): + producer = Mock(name="producer") + producer.transport = app.transport + producer.loop = app.loop + producer.beacon = Mock(name="beacon") + return ProducerThread(producer, loop=app.loop, beacon=producer.beacon) + + @pytest.mark.asyncio + async def test_on_start(self, *, producer_thread): + with patch.object(mod.confluent_kafka, "Producer") as P: + await producer_thread.on_start() + assert producer_thread._producer is P.return_value + + @pytest.mark.asyncio + async def test_flush(self, *, producer_thread): + producer_thread._producer = Mock(name="_producer") + await producer_thread.flush() + producer_thread._producer.flush.assert_called_once_with() + + @pytest.mark.asyncio + async def test_flush__no_producer(self, *, producer_thread): + producer_thread._producer = None + await producer_thread.flush() # must not raise + + @pytest.mark.asyncio + async def test_on_thread_stop(self, *, producer_thread): + producer_thread._producer = Mock(name="_producer") + await producer_thread.on_thread_stop() + producer_thread._producer.flush.assert_called_once_with() + + def test_produce__with_partition(self, *, producer_thread): + producer_thread._producer = Mock(name="_producer") + on_delivery = Mock() + producer_thread.produce("topic", b"key", b"value", 0, on_delivery) + producer_thread._producer.produce.assert_called_once_with( + "topic", b"key", b"value", 0, on_delivery=on_delivery + ) + + def test_produce__without_partition(self, *, producer_thread): + producer_thread._producer = Mock(name="_producer") + on_delivery = Mock() + producer_thread.produce("topic", b"key", b"value", None, on_delivery) + producer_thread._producer.produce.assert_called_once_with( + "topic", b"key", b"value", on_delivery=on_delivery + ) + + def test_produce__not_started(self, *, producer_thread): + producer_thread._producer = None + with pytest.raises(RuntimeError): + producer_thread.produce("topic", b"key", b"value", 0, Mock()) + + +class TestProducerProduceFuture: + def test_set_from_on_delivery__success(self, event_loop): + fut = ProducerProduceFuture(loop=event_loop) + message = Mock(name="message") + message.topic.return_value = "topic" + message.partition.return_value = 3 + message.offset.return_value = 9 + fut.set_from_on_delivery(None, message) + assert fut.done() + metadata = fut.result() + assert isinstance(metadata, RecordMetadata) + assert metadata.topic == "topic" + assert metadata.partition == 3 + assert metadata.offset == 9 + + def test_set_from_on_delivery__error(self, event_loop): + fut = ProducerProduceFuture(loop=event_loop) + exc = KeyError("boom") + fut.set_from_on_delivery(exc, None) + assert fut.done() + with pytest.raises(KeyError): + fut.result() + + def test_message_to_metadata(self, event_loop): + fut = ProducerProduceFuture(loop=event_loop) + message = Mock(name="message") + message.topic.return_value = "topic" + message.partition.return_value = 1 + message.offset.return_value = 5 + metadata = fut.message_to_metadata(message) + assert metadata.topic == "topic" + assert metadata.partition == 1 + assert metadata.offset == 5 + + +class TestTransport: + @pytest.fixture() + def transport(self, app): + return Transport(url=["kafka://localhost:9092"], app=app) + + def test_attributes(self, *, transport): + assert transport.Consumer is Consumer + assert transport.Producer is Producer + assert transport.default_port == 9092 + assert "confluent_kafka" in transport.driver_version + + def test__topic_config__empty(self, *, transport): + assert transport._topic_config() == {} + + def test__topic_config__retention(self, *, transport): + assert transport._topic_config(retention=3000)["retention.ms"] == 3000 + + def test__topic_config__compacting_and_deleting(self, *, transport): + config = transport._topic_config(compacting=True, deleting=True) + assert config["cleanup.policy"] == "compact,delete"