Feature: Sources API (#420)

Sources API

Author: ovv

quixstreams/app.py

@@ -2,6 +2,7 @@
 import functools
 import logging
 import os
+import time
 import signal
 import warnings
 from typing import Optional, List, Callable, Union, Literal, Tuple, Type
@@ -44,6 +45,7 @@
 from .state import StateStoreManager
 from .state.recovery import RecoveryManager
 from .state.rocksdb import RocksDBOptionsType
+from .sources.manager import SourceManager, BaseSource, SourceException
 from .utils.settings import BaseSettings
 
 __all__ = ("Application", "ApplicationConfig")
@@ -114,7 +116,7 @@ def __init__(
         on_message_processed: Optional[MessageProcessedCallback] = None,
         consumer_poll_timeout: float = 1.0,
         producer_poll_timeout: float = 0.0,
-        loglevel: Optional[LogLevel] = "INFO",
+        loglevel: Optional[Union[int, LogLevel]] = "INFO",
         auto_create_topics: bool = True,
         use_changelog_topics: bool = True,
         quix_config_builder: Optional[QuixKafkaConfigsBuilder] = None,
@@ -289,13 +291,7 @@ def __init__(
             extra_config=self._config.consumer_extra_config,
             on_error=on_consumer_error,
         )
-        self._producer = RowProducer(
-            broker_address=self._config.broker_address,
-            extra_config=self._config.producer_extra_config,
-            on_error=on_producer_error,
-            flush_timeout=self._config.flush_timeout,
-            transactional=self._config.exactly_once,
-        )
+        self._producer = self._get_rowproducer(on_error=on_producer_error)
         self._running = False
         self._failed = False
 
@@ -328,6 +324,7 @@ def __init__(
             recovery_manager=recovery_manager,
         )
 
+        self._source_manager = SourceManager()
         self._sink_manager = SinkManager()
         self._pausing_manager = PausingManager(consumer=self._consumer)
         self._processing_context = ProcessingContext(
@@ -364,7 +361,7 @@ def Quix(
         on_message_processed: Optional[MessageProcessedCallback] = None,
         consumer_poll_timeout: float = 1.0,
         producer_poll_timeout: float = 0.0,
-        loglevel: Optional[LogLevel] = "INFO",
+        loglevel: Optional[Union[int, LogLevel]] = "INFO",
         quix_config_builder: Optional[QuixKafkaConfigsBuilder] = None,
         auto_create_topics: bool = True,
         use_changelog_topics: bool = True,
@@ -573,7 +570,8 @@ def custom_ts_extractor(
 
     def dataframe(
         self,
-        topic: Topic,
+        topic: Optional[Topic] = None,
+        source: Optional[BaseSource] = None,
     ) -> StreamingDataFrame:
         """
         A simple helper method that generates a `StreamingDataFrame`, which is used
@@ -603,6 +601,12 @@ def dataframe(
             to be used as an input topic.
         :return: `StreamingDataFrame` object
         """
+        if not source and not topic:
+            raise ValueError("one of `source` or `topic` is required")
+
+        if source:
+            topic = self.add_source(source, topic)
+
         sdf = StreamingDataFrame(
             topic=topic,
             topic_manager=self._topic_manager,
@@ -633,6 +637,28 @@ 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: Optional[bool] = None,
+    ) -> RowProducer:
+        """
+        Create a RowProducer using the application config
+
+        Used to create the application producer as well as the sources producers
+        """
+
+        if transactional is None:
+            transactional = self._config.exactly_once
+
+        return RowProducer(
+            broker_address=self._config.broker_address,
+            extra_config=self._config.producer_extra_config,
+            flush_timeout=self._config.flush_timeout,
+            on_error=on_error,
+            transactional=transactional,
+        )
+
     def get_producer(self) -> Producer:
         """
         Create and return a pre-configured Producer instance.
@@ -716,6 +742,25 @@ def clear_state(self):
         """
         self._state_manager.clear_stores()
 
+    def add_source(self, source: BaseSource, topic: Optional[Topic] = None) -> Topic:
+        """
+        Add a source to the application.
+
+        See :class:`quixstreams.sources.base.BaseSource` for more details.
+
+        :param source: a :class:`quixstreams.sources.BaseSource` instance
+        :param topic: the :class:`quixstreams.models.Topic` instance the source will produce to
+            Default: the source default
+        """
+        if not topic:
+            topic = source.default_topic()
+            self._topic_manager.register(topic)
+
+        producer = self._get_rowproducer(transactional=False)
+        source.configure(topic, producer)
+        self._source_manager.register(source)
+        return topic
+
     def run(
         self,
         dataframe: StreamingDataFrame,
@@ -745,6 +790,19 @@ def run(
 
         :param dataframe: instance of `StreamingDataFrame`
         """
+        self._run(dataframe)
+
+    def _exception_handler(self, exc_type, exc_val, exc_tb):
+        fail = False
+
+        # Sources and the application are independent.
+        # If a source fails, the application can shutdown gracefully.
+        if exc_val is not None and exc_type is not SourceException:
+            fail = True
+
+        self.stop(fail=fail)
+
+    def _run(self, dataframe: Optional[StreamingDataFrame] = None):
         self._setup_signal_handlers()
 
         logger.info(
@@ -765,36 +823,52 @@ def run(
         exit_stack.enter_context(self._processing_context)
         exit_stack.enter_context(self._state_manager)
         exit_stack.enter_context(self._consumer)
-        exit_stack.push(
-            lambda exc_type, exc_val, exc_tb: self.stop(fail=exc_val is not None)
-        )
+        exit_stack.enter_context(self._source_manager)
+        exit_stack.push(self._exception_handler)
 
         with exit_stack:
             # Subscribe to topics in Kafka and start polling
-            self._consumer.subscribe(
-                dataframe.topics_to_subscribe,
-                on_assign=self._on_assign,
-                on_revoke=self._on_revoke,
-                on_lost=self._on_lost,
-            )
-            logger.info("Waiting for incoming messages")
-            # Start polling Kafka for messages and callbacks
-            self._running = True
+            if dataframe is not None:
+                self._run_dataframe(dataframe)
+            else:
+                self._run_sources()
+
+    def _run_dataframe(self, dataframe):
+        self._consumer.subscribe(
+            dataframe.topics_to_subscribe,
+            on_assign=self._on_assign,
+            on_revoke=self._on_revoke,
+            on_lost=self._on_lost,
+        )
+        logger.info("Waiting for incoming messages")
+        # Start polling Kafka for messages and callbacks
+        self._running = True
+
+        # Initialize the checkpoint
+        self._processing_context.init_checkpoint()
+
+        dataframe_composed = dataframe.compose()
+
+        while self._running:
+            if self._state_manager.recovery_required:
+                self._state_manager.do_recovery()
+            else:
+                self._process_message(dataframe_composed)
+                self._processing_context.commit_checkpoint()
+                self._processing_context.resume_ready_partitions()
+                self._source_manager.raise_for_error()
 
-            # Initialize the checkpoint
-            self._processing_context.init_checkpoint()
+        logger.info("Stop processing of StreamingDataFrame")
 
-            dataframe_composed = dataframe.compose()
+    def _run_sources(self):
+        self._running = True
+        while self._running:
+            self._source_manager.raise_for_error()
 
-            while self._running:
-                if self._state_manager.recovery_required:
-                    self._state_manager.do_recovery()
-                else:
-                    self._process_message(dataframe_composed)
-                    self._processing_context.commit_checkpoint()
-                    self._processing_context.resume_ready_partitions()
+            if not self._source_manager.is_alive():
+                self.stop()
 
-            logger.info("Stop processing of StreamingDataFrame")
+            time.sleep(1)
 
     def _quix_runtime_init(self):
         """

quixstreams/dataframe/dataframe.py

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

quixstreams/logging.py

@@ -1,6 +1,7 @@
+import os
 import logging
 import sys
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
 __all__ = ("configure_logging", "LogLevel")
 
@@ -13,18 +14,17 @@
     "NOTSET",
 ]
 
-_LOGGER_NAME = "quixstreams"
+LOGGER_NAME = "quixstreams"
 
-_DEFAULT_FORMATTER = logging.Formatter(
-    f"[%(asctime)s] [%(levelname)s] [{_LOGGER_NAME}] : %(message)s"
-)
 _DEFAULT_HANDLER = logging.StreamHandler(stream=sys.stderr)
 
-logger = logging.getLogger(_LOGGER_NAME)
+logger = logging.getLogger(LOGGER_NAME)
 
 
 def configure_logging(
-    loglevel: Optional[LogLevel],
+    loglevel: Optional[Union[int, LogLevel]],
+    name: str = LOGGER_NAME,
+    pid: bool = False,
 ) -> bool:
     """
     Configure "quixstreams" logger.
@@ -35,20 +35,36 @@ def configure_logging(
 
     :param loglevel: a valid log level as a string or None.
         If None passed, this function is no-op and no logging will be configured.
+    :param name: the log name included in the output
+    :param pid: if True include the process PID in the logs
     :return: True if logging config has been updated, otherwise False.
     """
     if loglevel is None:
         # Skipping logging configuration
         return False
 
+    if pid:
+        formatter = (
+            f"[%(asctime)s] [%(levelname)s] [{name}] [{os.getpid()}] : %(message)s"
+        )
+    else:
+        formatter = f"[%(asctime)s] [%(levelname)s] [{name}] : %(message)s"
+
     if logger.handlers:
-        # There's a pre-configured handler for "quixstreams", leave it as it is
-        return False
+        if len(logger.handlers) != 1 or logger.handlers[0] is not _DEFAULT_HANDLER:
+            # There's a pre-configured handler for "quixstreams", leave it as it is
+            return False
+        else:
+            # The pre-configured handler for "quixstreams" is the default handler.
+            # Reconfigure the formatter in case we are in a subprocess and the logger
+            # was configured by mistake by the Application.
+            _DEFAULT_HANDLER.setFormatter(logging.Formatter(formatter))
+            return True
 
     # Configuring logger
     logger.setLevel(loglevel)
     logger.propagate = False
-    _DEFAULT_HANDLER.setFormatter(_DEFAULT_FORMATTER)
+    _DEFAULT_HANDLER.setFormatter(logging.Formatter(formatter))
     _DEFAULT_HANDLER.setLevel(loglevel)
     logger.addHandler(_DEFAULT_HANDLER)
     return True

quixstreams/models/topics/manager.py

@@ -283,6 +283,23 @@ def topic(
         self._topics[name] = topic
         return topic
 
+    def register(self, topic: Topic):
+        """
+        Register an already generated :class:`quixstreams.models.topics.Topic` to the topic manager.
+
+        The topic name and config can be updated by the topic manager.
+
+        :param topic: The topic to register
+        """
+        topic.name = self._resolve_topic_name(topic.name)
+        if topic.config is None:
+            topic.config = TopicConfig(
+                num_partitions=self.default_num_partitions,
+                replication_factor=self.default_replication_factor,
+                extra_config=self.default_extra_config,
+            )
+        self._topics[topic.name] = topic
+
     def repartition_topic(
         self,
         operation: str,

quixstreams/models/topics/topic.py

@@ -93,7 +93,7 @@ class Topic:
     def __init__(
         self,
         name: str,
-        config: TopicConfig,
+        config: Optional[TopicConfig] = None,
         value_deserializer: Optional[DeserializerType] = None,
         key_deserializer: Optional[DeserializerType] = BytesDeserializer(),
         value_serializer: Optional[SerializerType] = None,
@@ -110,25 +110,14 @@ def __init__(
         :param timestamp_extractor: a callable that returns a timestamp in
             milliseconds from a deserialized message.
         """
-        self._name = name
-        self._config = config
+        self.name = name
+        self.config = config
         self._key_serializer = _get_serializer(key_serializer)
         self._key_deserializer = _get_deserializer(key_deserializer)
         self._value_serializer = _get_serializer(value_serializer)
         self._value_deserializer = _get_deserializer(value_deserializer)
         self._timestamp_extractor = timestamp_extractor
 
-    @property
-    def name(self) -> str:
-        """
-        Topic name
-        """
-        return self._name
-
-    @property
-    def config(self) -> TopicConfig:
-        return self._config
-
     def row_serialize(self, row: Row, key: Any) -> KafkaMessage:
         """
         Serialize Row to a Kafka message structure

quixstreams/sources/__init__.py

@@ -0,0 +1,3 @@
+from .base import *
+from .manager import SourceException
+from .multiprocessing import multiprocessing