review comments

Author: ovv

quixstreams/app.py

@@ -281,15 +281,8 @@ def __init__(
             extra_config=consumer_extra_config,
             on_error=on_consumer_error,
         )
-        self._producer = RowProducer(
-            broker_address=broker_address,
-            extra_config=producer_extra_config,
-            on_error=on_producer_error,
-            flush_timeout=consumer_extra_config.get(
-                "max.poll.interval.ms", _default_max_poll_interval_ms
-            )
-            / 1000,  # convert to seconds
-            transactional=self._uses_exactly_once,
+        self._producer = self._get_rowproducer(
+            on_error=on_producer_error, transactional=self._uses_exactly_once
         )
         self._consumer_poll_timeout = consumer_poll_timeout
         self._producer_poll_timeout = producer_poll_timeout
@@ -603,7 +596,7 @@ def dataframe(
         :return: `StreamingDataFrame` object
         """
         if not source and not topic:
-            raise TypeError("one of `source` or `topic` is required")
+            raise ValueError("one of `source` or `topic` is required")
         elif source and not topic:
             topic = self._topic_manager.source_topic(source)
 
@@ -638,6 +631,26 @@ def stop(self, fail: bool = False):
         if self._state_manager.using_changelogs:
             self._state_manager.stop_recovery()
 
+    def _get_rowproducer(
+        self,
+        on_error: Optional[ProducerErrorCallback] = None,
+        transactional: bool = False,
+    ) -> RowProducer:
+        flush_timeout = (
+            self._consumer_extra_config.get(
+                "max.poll.interval.ms", _default_max_poll_interval_ms
+            )
+            / 1000
+        )
+
+        return RowProducer(
+            broker_address=self._broker_address,
+            extra_config=self._producer_extra_config,
+            flush_timeout=flush_timeout,
+            on_error=on_error,
+            transactional=transactional,
+        )
+
     def get_producer(self) -> Producer:
         """
         Create and return a pre-configured Producer instance.
@@ -721,7 +734,7 @@ def clear_state(self):
         """
         self._state_manager.clear_stores()
 
-    def source(self, source: BaseSource, topic: Optional[Topic] = None) -> Topic:
+    def add_source(self, source: BaseSource, topic: Optional[Topic] = None) -> Topic:
         """
         Add a source to the application.
 
@@ -734,15 +747,7 @@ def source(self, source: BaseSource, topic: Optional[Topic] = None) -> Topic:
         if not topic:
             topic = self._topic_manager.source_topic(source)
 
-        producer = RowProducer(
-            broker_address=self._broker_address,
-            extra_config=self._producer_extra_config,
-            flush_timeout=self._consumer_extra_config.get(
-                "max.poll.interval.ms", _default_max_poll_interval_ms
-            )
-            / 1000,  # convert to seconds
-            transactional=False,
-        )
+        producer = self._get_rowproducer(transactional=False)
         source.configure(topic, producer)
         self._source_manager.register(source)
         return topic
@@ -804,7 +809,7 @@ def _run(self, dataframe: Optional[StreamingDataFrame] = None):
         self._setup_topics()
 
         if dataframe is not None and dataframe.source:
-            self.source(dataframe.source, dataframe.topic)
+            self.add_source(dataframe.source, dataframe.topic)
 
         exit_stack = contextlib.ExitStack()
         exit_stack.enter_context(self._processing_context)
@@ -852,7 +857,7 @@ def _run_sources(self):
         while self._running:
             self._source_manager.raise_for_error()
 
-            if not self._source_manager.alives():
+            if not self._source_manager.is_alive():
                 self.stop()
 
             time.sleep(1)

quixstreams/dataframe/dataframe.py

@@ -49,7 +49,6 @@
 from quixstreams.sinks import BaseSink
 from quixstreams.state.types import State
 from quixstreams.sources import BaseSource
-from quixstreams.rowproducer import RowProducer
 from .base import BaseStreaming
 from .exceptions import InvalidOperation, GroupByLimitExceeded, DataFrameLocked
 from .series import StreamingSeries

quixstreams/models/topics/manager.py

@@ -289,7 +289,7 @@ def topic(
 
     def source_topic(self, source: "BaseSource") -> Topic:
         topic_args = source.default_topic()
-        return self.topic(**topic_args.asdict())
+        return self.topic(**topic_args.asargs())
 
     def repartition_topic(
         self,

quixstreams/sources/base.py

@@ -46,7 +46,7 @@ class SourceTopic:
     key_serializer: Optional[SerializerType] = BytesSerializer()
     timestamp_extractor: Optional[TimestampExtractor] = None
 
-    def asdict(self):
+    def asargs(self):
         return {
             field.name: getattr(self, field.name) for field in dataclasses.fields(self)
         }
@@ -257,15 +257,20 @@ def __init__(
 
         self._checkpoint_interval = checkpoint_interval
         self._checkpoint_error: Optional[BaseException] = None
+        self._checkpoint_timer: Optional[threading.Timer] = None
 
-        self.__checkpoint_timer = self.__make_timer()
+    @abstractmethod
+    def run(self):
+        self._checkpoint_timer = self._make_timer()
+        self._checkpoint_timer.start()
+        super().run()
 
-    def __make_timer(self):
+    def _make_timer(self) -> threading.Timer:
         return threading.Timer(
             interval=self._checkpoint_interval, function=self.__checkpoint
         )
 
-    def __checkpoint(self):
+    def _checkpoint(self):
         try:
             self.checkpoint()
         except BaseException as err:
@@ -274,16 +279,11 @@ def __checkpoint(self):
             self.stop()
             return
 
-        self.__checkpoint_timer = self.__make_timer()
-        self.__checkpoint_timer.start()
-
-    @abstractmethod
-    def run(self):
-        self.__checkpoint_timer.start()
-        super().run()
+        self._checkpoint_timer = self._make_timer()
+        self._checkpoint_timer.start()
 
     def cleanup(self, failed):
-        self.__checkpoint_timer.cancel()
+        self._checkpoint_timer.cancel()
         if self._checkpoint_error:
             super().cleanup(True)
             raise self._checkpoint_error
@@ -321,10 +321,12 @@ def __init__(
         super().__init__(name, shutdown_timeout)
 
         self._polling_delay = polling_delay
-        self._stopping = threading.Event()
+        self._stopping: Optional[threading.Event] = None
 
     def run(self) -> None:
         super().run()
+
+        self._stopping = threading.Event()
         while not self._stopping.is_set():
             try:
                 msg = self.poll()
@@ -350,7 +352,8 @@ def sleep(self, seconds: float):
             raise PollingSourceShutdown("shutdown")
 
     def stop(self) -> None:
-        self._stopping.set()
+        if self._stopping is not None:
+            self._stopping.set()
         super().stop()
 
     def cleanup(self, failed: bool) -> None:

quixstreams/sources/iterable.py

@@ -1,6 +1,6 @@
 import logging
 
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Optional, Tuple, Callable, Generator
 
 from quixstreams.models.messages import KafkaMessage
 
@@ -10,6 +10,7 @@
 
 
 __all__ = (
+    "GeneratorSource",
     "ValueIterableSource",
     "KeyValueIterableSource",
 )
@@ -26,7 +27,7 @@ class ValueIterableSource(PollingSource):
     from quixstreams.sources import ValueIterableSource
 
     app = Application(broker_address='localhost:9092', consumer_group='group')
-    source = ValueIterableSource(name="my_source", values=iter(range(10)))
+    source = ValueIterableSource(name="my_source", values=range(10))
 
     sdf = app.dataframe(source=source)
     sdf.print()
@@ -51,7 +52,7 @@ def __init__(
         super().__init__(name, shutdown_timeout)
 
         self._key = key
-        self._values = values
+        self._values = iter(values)
 
     def poll(self) -> KafkaMessage:
         try:
@@ -75,14 +76,9 @@ class KeyValueIterableSource(PollingSource):
     from quixstreams import Application
     from quixstreams.sources import KeyValueIterableSource
 
-    def messages():
-        yield "one", 1
-        yield "two", 2
-        yield "three", 3
-        yield "four", 4
-
+    keys = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
     app = Application(broker_address='localhost:9092', consumer_group='group')
-    source = KeyValueIterableSource(name="my_source", iterable=messages())
+    source = KeyValueIterableSource(name="my_source", iterable=zip(keys, range(10)))
 
     sdf = app.dataframe(source=source)
     sdf.print()
@@ -104,11 +100,70 @@ def __init__(
         """
         super().__init__(name, shutdown_timeout)
 
-        self._iterable = iterable
+        self._iterator = iter(iterable)
+
+    def poll(self) -> KafkaMessage:
+        try:
+            data = next(self._iterator)
+        except StopIteration:
+            raise PollingSourceShutdown()
+
+        if data is None:
+            return data
+
+        key, value = data
+        return self.serialize(key=key, value=value)
+
+
+class GeneratorSource(PollingSource):
+    """
+    PollingSource implementation that iterator over a generator of (key, value)
+
+    Example Snippet:
+
+    ```python
+    from quixstreams import Application
+    from quixstreams.sources import GeneratorSource
+
+    def messages():
+        yield "one", 1
+        yield "two", 2
+        yield "three", 3
+        yield "four", 4
+
+    app = Application(broker_address='localhost:9092', consumer_group='group')
+    source = GeneratorSource(name="my_source", generator=messages)
+
+    sdf = app.dataframe(source=source)
+    sdf.print()
+
+    app.run(sdf)
+    ```
+    """
+
+    def __init__(
+        self,
+        name: str,
+        generator: Callable[[], Generator[Optional[Tuple[any, any]], None, None]],
+        polling_delay: float = 1,
+        shutdown_timeout: float = 10,
+    ) -> None:
+        super().__init__(name, polling_delay, shutdown_timeout)
+
+        self._generator = generator
+        self._generator_instance: Optional[
+            Generator[Optional[Tuple[any, any]], None, None]
+        ] = None
+
+    def run(self):
+        self._generator_instance: Generator[Optional[Tuple[any, any]], None, None] = (
+            self._generator()
+        )
+        super().run()
 
     def poll(self) -> KafkaMessage:
         try:
-            data = next(self._iterable)
+            data = next(self._generator_instance)
         except StopIteration:
             raise PollingSourceShutdown()
 

quixstreams/sources/manager.py

@@ -211,7 +211,7 @@ def raise_for_error(self) -> None:
         for process in self.processes:
             process.raise_for_error()
 
-    def alives(self) -> bool:
+    def is_alive(self) -> bool:
         """
         Check if any process is alive
 

tests/test_quixstreams/test_app.py

@@ -2145,7 +2145,7 @@ def values():
             yield 2
 
         source = ValueIterableSource(name="foo", key="foo", values=values())
-        app.source(source, topic=app.topic(topic_name))
+        app.add_source(source, topic=app.topic(topic_name))
         app._run()
 
         results = []