Add topic_type to the Topic objects (#847)

Author: daniil-quix

quixstreams/models/topics/manager.py

@@ -10,7 +10,7 @@
     TopicNameLengthExceeded,
     TopicNotFoundError,
 )
-from .topic import TimestampExtractor, Topic, TopicConfig
+from .topic import TimestampExtractor, Topic, TopicConfig, TopicType
 
 logger = logging.getLogger(__name__)
 
@@ -56,24 +56,16 @@ def __init__(
         """
         self._admin = topic_admin
         self._consumer_group = consumer_group
-        self._topics: Dict[str, Topic] = {}
+        self._regular_topics: Dict[str, Topic] = {}
         self._repartition_topics: Dict[str, Topic] = {}
         self._changelog_topics: Dict[Optional[str], Dict[str, Topic]] = {}
         self._timeout = timeout
         self._create_timeout = create_timeout
         self._auto_create_topics = auto_create_topics
 
-    @property
-    def _all_topics_list(self) -> List[Topic]:
-        return (
-            list(self._topics.values())
-            + list(self._repartition_topics.values())
-            + self.changelog_topics_list
-        )
-
     @property
     def topics(self) -> Dict[str, Topic]:
-        return self._topics
+        return self._regular_topics
 
     @property
     def repartition_topics(self) -> Dict[str, Topic]:
@@ -102,7 +94,7 @@ def non_changelog_topics(self) -> Dict[str, Topic]:
         """
         Returns a dict with normal and repartition topics
         """
-        return {**self._topics, **self._repartition_topics}
+        return {**self._regular_topics, **self._repartition_topics}
 
     @property
     def all_topics(self) -> Dict[str, Topic]:
@@ -111,7 +103,12 @@ def all_topics(self) -> Dict[str, Topic]:
 
         returns: full topic dict, {topic_name: Topic}
         """
-        return {topic.name: topic for topic in self._all_topics_list}
+        all_topics_list = (
+            list(self._regular_topics.values())
+            + list(self._repartition_topics.values())
+            + self.changelog_topics_list
+        )
+        return {topic.name: topic for topic in all_topics_list}
 
     def topic_config(
         self,
@@ -175,10 +172,11 @@ def topic(
             key_deserializer=key_deserializer,
             create_config=create_config,
             timestamp_extractor=timestamp_extractor,
+            topic_type=TopicType.REGULAR,
         )
         broker_topic = self._get_or_create_broker_topic(topic)
         topic = self._configure_topic(topic, broker_topic)
-        self._topics[topic.name] = topic
+        self._regular_topics[topic.name] = topic
         return topic
 
     def register(self, topic: Topic) -> Topic:
@@ -197,7 +195,7 @@ def register(self, topic: Topic) -> Topic:
             )
         broker_topic = self._get_or_create_broker_topic(topic)
         topic = self._configure_topic(topic, broker_topic)
-        self._topics[topic.name] = topic
+        self._regular_topics[topic.name] = topic
         return topic
 
     def repartition_topic(
@@ -230,6 +228,7 @@ def repartition_topic(
             value_serializer=value_serializer,
             key_serializer=key_serializer,
             create_config=config,
+            topic_type=TopicType.REPARTITION,
         )
         broker_topic = self._get_or_create_broker_topic(topic)
         topic = self._configure_topic(topic, broker_topic)
@@ -269,6 +268,7 @@ def changelog_topic(
             key_deserializer="bytes",
             value_deserializer="bytes",
             create_config=config,
+            topic_type=TopicType.CHANGELOG,
         )
         broker_topic = self._get_or_create_broker_topic(topic)
         topic = self._configure_topic(topic, broker_topic)

quixstreams/models/topics/topic.py

@@ -1,5 +1,6 @@
 import copy
 import dataclasses
+import enum
 import logging
 from typing import Any, Callable, List, Optional, Union
 
@@ -30,7 +31,7 @@
     SuccessfulConfluentKafkaMessageProto,
 )
 
-__all__ = ("Topic", "TopicConfig", "TimestampExtractor")
+__all__ = ("Topic", "TopicConfig", "TimestampExtractor", "TopicType")
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +83,12 @@ def _get_deserializer(
     return deserializer
 
 
+class TopicType(enum.Enum):
+    REGULAR = 1
+    REPARTITION = 2
+    CHANGELOG = 3
+
+
 class Topic:
     """
     A definition of a Kafka topic.
@@ -94,6 +101,7 @@ class Topic:
     def __init__(
         self,
         name: str,
+        topic_type: TopicType = TopicType.REGULAR,
         create_config: Optional[TopicConfig] = None,
         value_deserializer: Optional[DeserializerType] = None,
         key_deserializer: Optional[DeserializerType] = BytesDeserializer(),
@@ -104,6 +112,13 @@ def __init__(
     ):
         """
         :param name: topic name
+        :param topic_type: a type of the topic, can be one of:
+            - `TopicType.REGULAR` - the regular input and output topics
+            - `TopicType.REPARTITION` - a repartition topic used for re-keying the data
+            - `TopicType.CHANGELOG` - a changelog topic to back up the state stores.
+
+            Default - `TopicType.REGULAR`.
+
         :param create_config: a `TopicConfig` to create a new topic if it does not exist
         :param value_deserializer: a deserializer type for values
         :param key_deserializer: a deserializer type for keys
@@ -123,6 +138,7 @@ def __init__(
         self._value_serializer = _get_serializer(value_serializer)
         self._key_serializer = _get_serializer(key_serializer)
         self._timestamp_extractor = timestamp_extractor
+        self._type = topic_type
 
     def __clone__(
         self,
@@ -138,6 +154,7 @@ def __clone__(
             value_serializer=self._value_serializer,
             key_serializer=self._key_serializer,
             timestamp_extractor=self._timestamp_extractor,
+            topic_type=self._type,
         )
 
     @property
@@ -166,6 +183,18 @@ def broker_config(self) -> TopicConfig:
     def broker_config(self, config: TopicConfig):
         self._broker_config = copy.deepcopy(config)
 
+    @property
+    def is_changelog(self) -> bool:
+        return self._type == TopicType.CHANGELOG
+
+    @property
+    def is_regular(self) -> bool:
+        return self._type == TopicType.REGULAR
+
+    @property
+    def is_repartition(self) -> bool:
+        return self._type == TopicType.REPARTITION
+
     def row_serialize(self, row: Row, key: Any) -> KafkaMessage:
         """
         Serialize Row to a Kafka message structure

tests/test_quixstreams/test_models/test_topics/test_manager.py

@@ -50,6 +50,7 @@ def test_topic_with_config(self, topic_manager_factory):
 
         assert topic.name == topic_name
         assert topic.quix_name == topic_name
+        assert topic.is_regular
         assert topic.create_config.num_partitions == create_config.num_partitions
         assert (
             topic.create_config.replication_factor == create_config.replication_factor
@@ -187,6 +188,7 @@ def test_changelog_topic(self, topic_manager_factory):
         assert changelog.broker_config.num_partitions == 1
         assert changelog.broker_config.replication_factor == 1
         assert changelog.broker_config.extra_config["cleanup.policy"] == "compact"
+        assert changelog.is_changelog
 
     def test_changelog_topic_partition_count_mismatch(
         self, topic_manager_factory, kafka_admin_client
@@ -260,6 +262,7 @@ def test_repartition_topic(self, topic_manager_factory):
         assert repartition.broker_config.replication_factor == 1
         assert repartition.broker_config.extra_config["retention.ms"] == "1000"
         assert repartition.broker_config.extra_config["retention.bytes"] == "1000"
+        assert repartition.is_repartition
 
     def test_changelog_nested_internal_topic_naming(self, topic_manager_factory):
         """