feat: add offset management integration to consumer stream (#541)

* feat: add offset management integration to consumer stream

* chore: update fluvio version

* feat: add delete_consumer_offset

* test: more commit consumer tests

Author: fraidev

Cargo.toml

@@ -38,11 +38,11 @@ url = "2.5.0"
 webbrowser = "1.0.0"
 
 fluvio-future = { version = "0.7.0", features = ["task", "io", "native_tls", "subscriber"] }
-fluvio = { features = ["admin", "rustls"], git = "https://github.com/infinyon/fluvio.git", tag = "v0.14.0" }
-fluvio-protocol = { git = "https://github.com/infinyon/fluvio.git", tag = "v0.14.0" }
-fluvio-types = { git = "https://github.com/infinyon/fluvio.git", tag = "v0.14.0" }
-fluvio-sc-schema = { git = "https://github.com/infinyon/fluvio.git", tag = "v0.14.0" }
-fluvio-controlplane-metadata = { git = "https://github.com/infinyon/fluvio.git", tag = "v0.14.0" }
+fluvio = { features = ["admin", "rustls"], git = "https://github.com/infinyon/fluvio.git", brach ="offset_flush_box" }
+fluvio-protocol = { git = "https://github.com/infinyon/fluvio.git", brach ="offset_flush_box" }
+fluvio-types = { git = "https://github.com/infinyon/fluvio.git", brach ="offset_flush_box" }
+fluvio-sc-schema = { git = "https://github.com/infinyon/fluvio.git", brach ="offset_flush_box" }
+fluvio-controlplane-metadata = { git = "https://github.com/infinyon/fluvio.git", brach ="offset_flush_box" }
 
 # transitive version selection
 parking_lot = "0.12.3"

fluvio/__init__.py

@@ -58,6 +58,28 @@
 records = [bytearray(next(stream).value()).decode() for _ in range(num_items)]
 ```
 
+Also you can consume usign offset management:
+
+```python
+import fluvio
+
+fluvio = Fluvio.connect()
+
+topic = "a_topic"
+builder = ConsumerConfigExtBuilder(topic)
+builder.offset_start(Offset.beginning())
+builder.offset_strategy(OffsetManagementStrategy.MANUAL)
+builder.offset_consumer("a-consumer")
+config = builder.build()
+stream = fluvio.consumer_with_config(config)
+
+num_items = 2
+records = [bytearray(next(stream).value()).decode() for _ in range(num_items)]
+
+stream.offset_commit()
+stream.offset_flush()
+```
+
 For more examples see the integration tests in the fluvio-python repository.
 
 [test_produce.py](https://github.com/infinyon/fluvio-client-python/blob/main/integration-tests/test_produce.py)
@@ -82,6 +104,8 @@
     PartitionSelectionStrategy as _PartitionSelectionStrategy,
     PartitionConsumerStream as _PartitionConsumerStream,
     AsyncPartitionConsumerStream as _AsyncPartitionConsumerStream,
+    OffsetManagementStrategy,
+    ConsumerOffset as _ConsumerOffset,
     # producer types
     TopicProducer as _TopicProducer,
     ProduceOutput as _ProduceOutput,
@@ -137,6 +161,8 @@
     "ConsumerConfig",
     "PartitionConsumer",
     "MultiplePartitionConsumer",
+    "OffsetManagementStrategy",
+    "ConsumerOffset",
     # specs
     "CommonCreateRequest",
     "SmartModuleSpec",
@@ -262,6 +288,26 @@ def smartmodule(
         )
 
 
+class ConsumerIterator:
+    def __init__(self, stream: _PartitionConsumerStream):
+        self.stream = stream
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        item = self.stream.next()
+        if item is None:
+            raise StopIteration
+        return Record(item)
+
+    def offset_commit(self):
+        self.stream.offset_commit()
+
+    def offset_flush(self):
+        self.stream.offset_flush()
+
+
 class PartitionConsumer:
     """
     An interface for consuming events from a particular partition
@@ -614,6 +660,30 @@ def with_multiple(cls, topic: typing.List[typing.Tuple[str, int]]):
         return cls(_PartitionSelectionStrategy.with_multiple(topic))
 
 
+class ConsumerOffset:
+    """Consumer offset"""
+
+    _inner: _ConsumerOffset
+
+    def __init__(self, inner: _ConsumerOffset):
+        self._inner = inner
+
+    def consumer_id(self) -> str:
+        return self._inner.consumer_id()
+
+    def topic(self) -> str:
+        return self._inner.topic()
+
+    def partition(self) -> int:
+        return self._inner.partition()
+
+    def offset(self) -> int:
+        return self._inner.offset()
+
+    def modified_time(self) -> int:
+        return self._inner.modified_time()
+
+
 class FluvioConfig:
     """Configuration for Fluvio client"""
 
@@ -687,14 +757,12 @@ def connect_with_config(cls, config: FluvioConfig):
         """Creates a new Fluvio client using the given configuration"""
         return cls(_Fluvio.connect_with_config(config._inner))
 
-    def consumer_with_config(
-        self, config: ConsumerConfigExt
-    ) -> typing.Iterator[Record]:
+    def consumer_with_config(self, config: ConsumerConfigExt) -> ConsumerIterator:
         """Creates consumer with settings defined in config
 
         This is the recommended way to create a consume records.
         """
-        return self._generator(self._inner.consumer_with_config(config))
+        return ConsumerIterator(self._inner.consumer_with_config(config))
 
     def topic_producer(self, topic: str) -> TopicProducer:
         """
@@ -740,6 +808,14 @@ def multi_topic_partition_consumer(
             self._inner.multi_partition_consumer(strategy._inner)
         )
 
+    def consumer_offsets(self) -> typing.List[ConsumerOffset]:
+        """Fetch the current offsets of the consumer"""
+        return self._inner.consumer_offsets()
+
+    def delete_consumer_offset(self, consumer: str, topic: str, partition: int):
+        """Delete the consumer offset"""
+        return self._inner.delete_consumer_offset(consumer, topic, partition)
+
     def _generator(self, stream: _PartitionConsumerStream) -> typing.Iterator[Record]:
         item = stream.next()
         while item is not None:

integration-tests/test_base.py

@@ -2,7 +2,7 @@
 import uuid
 import time
 
-from fluvio import FluvioAdmin
+from fluvio import FluvioAdmin, Fluvio
 
 
 def create_smartmodule(sm_name, sm_path):
@@ -22,7 +22,9 @@ class CommonFluvioSetup(unittest.TestCase):
     def common_setup(self, sm_path=None):
         """Optionally create a smartmodule if sm_path is provided"""
         self.admin = FluvioAdmin.connect()
+        self.fluvio = Fluvio.connect()
         self.topic = str(uuid.uuid4())
+        self.consumer_id = str(uuid.uuid4())
         self.sm_name = str(uuid.uuid4())
         self.sm_path = sm_path
 
@@ -50,6 +52,9 @@ def setUp(self):
 
     def tearDown(self):
         self.admin.delete_topic(self.topic)
+        # TODO: we can remove this delete_consumer_offset after to fix
+        # this bug: https://github.com/infinyon/fluvio/issues/4308
+        self.fluvio.delete_consumer_offset(self.consumer_id, self.topic, 0)
         time.sleep(1)
         if self.sm_path is not None:
             self.admin.delete_smartmodule(self.sm_name)

integration-tests/test_consume.py

@@ -1,6 +1,7 @@
 from string import ascii_lowercase
+import time
 
-from fluvio import ConsumerConfig, Fluvio, Offset
+from fluvio import ConsumerConfig, Fluvio, Offset, OffsetManagementStrategy
 from fluvio import ConsumerConfigExtBuilder
 from test_base import CommonFluvioSetup
 
@@ -59,6 +60,140 @@ def test_consume_at_offset_begin(self):
         self.assertEqual(records[0], "record-z")
         self.assertEqual(records[1], "record-y")
 
+    def test_consume_offset_managed_auto_manual(self):
+        """
+        Manual offset management test
+        """
+        fluvio = Fluvio.connect()
+
+        builder = ConsumerConfigExtBuilder(self.topic)
+        builder.offset_start(Offset.beginning())
+        builder.offset_strategy(OffsetManagementStrategy.MANUAL)
+        builder.offset_consumer(self.consumer_id)
+
+        config = builder.build()
+        stream = fluvio.consumer_with_config(config)
+
+        num_items = 2
+        records = [bytearray(next(stream).value()).decode() for _ in range(num_items)]
+
+        time.sleep(1)
+
+        stream.offset_commit()
+        stream.offset_flush()
+
+        self.assertEqual(len(records), 2)
+        self.assertEqual(records[0], "record-z")
+        self.assertEqual(records[1], "record-y")
+
+        consumers = fluvio.consumer_offsets()
+        self.assertEqual(len(consumers), 1)
+        self.assertEqual(consumers[0].consumer_id, self.consumer_id)
+        self.assertEqual(consumers[0].topic, self.topic)
+        self.assertEqual(consumers[0].partition, 0)
+        self.assertEqual(consumers[0].offset, 1)
+
+    def test_consume_offset_managed_auto_commit_and_flush(self):
+        """
+        Automatic offset management test
+        """
+        fluvio = Fluvio.connect()
+
+        builder = ConsumerConfigExtBuilder(self.topic)
+        builder.offset_start(Offset.beginning())
+        builder.offset_strategy(OffsetManagementStrategy.AUTO)
+        builder.offset_consumer(self.consumer_id)
+
+        config = builder.build()
+        stream = fluvio.consumer_with_config(config)
+
+        num_items = 2
+        records = [bytearray(next(stream).value()).decode() for _ in range(num_items)]
+
+        stream.offset_commit()
+        stream.offset_flush()
+
+        self.assertEqual(len(records), 2)
+        self.assertEqual(records[0], "record-z")
+        self.assertEqual(records[1], "record-y")
+
+        consumers = fluvio.consumer_offsets()
+        self.assertEqual(len(consumers), 1)
+        self.assertEqual(consumers[0].consumer_id, self.consumer_id)
+        self.assertEqual(consumers[0].topic, self.topic)
+        self.assertEqual(consumers[0].partition, 0)
+        self.assertEqual(consumers[0].offset, 1)
+
+    def test_consume_offset_managed_auto_flush(self):
+        """
+         Consume with an offset_consume id defined
+
+         This will case the cluster to track that offsets
+         that consumer has consumed, stoping and resuming
+         from last offset
+
+         The equivalent of this test with the CLI is:
+        `fluvio consume test-topic --offset-beginning --offset-consumer --end 1 foo`
+         record-z
+
+
+        `fluvio consume test-topic --offset-beginning --offset-consumer --end 2 foo`
+         record-y
+
+        """
+        fluvio = Fluvio.connect()
+
+        # configure consumer
+        builder = ConsumerConfigExtBuilder(self.topic)
+        builder.disable_continuous()
+        builder.offset_consumer(self.consumer_id)
+        builder.offset_start(Offset.beginning())
+        builder.offset_strategy(OffsetManagementStrategy.AUTO)
+
+        config = builder.build()
+        stream = fluvio.consumer_with_config(config)
+
+        num_items = 2
+        records = []
+        for _ in range(num_items):
+            records.append(bytearray(next(stream).value()).decode())
+        self.assertEqual(len(records), 2)
+        self.assertEqual(records[0], "record-z")
+        # ensure offset is flushed
+        # with Auto, this normally happens with advancement of the stream and/or
+        # closure of the client
+        # but that isn't at a guaranteed point for python in the middle
+        # of this test, so do this explictly
+        stream.offset_flush()
+        # connect again with a new client, but the same id
+        fluvio = Fluvio.connect()
+
+        # configure consumer
+        # even though offset_start says to start at the beginning
+        # it is expected start from the last offset consumed
+        # because the offset_consumer is set
+        builder = ConsumerConfigExtBuilder(self.topic)
+        builder.disable_continuous()
+        builder.offset_consumer(self.consumer_id)
+        builder.offset_start(Offset.beginning())
+        builder.offset_strategy(OffsetManagementStrategy.AUTO)
+
+        stream = fluvio.consumer_with_config(config)
+
+        num_items = 1
+        records = []
+        for _ in range(1):
+            records.append(bytearray(next(stream).value()).decode())
+        self.assertEqual(len(records), 1)
+        self.assertEqual(records[0], "record-y")
+
+        consumers = fluvio.consumer_offsets()
+        self.assertEqual(len(consumers), 1)
+        self.assertEqual(consumers[0].consumer_id, self.consumer_id)
+        self.assertEqual(consumers[0].topic, self.topic)
+        self.assertEqual(consumers[0].partition, 0)
+        self.assertEqual(consumers[0].offset, 1)
+
     def test_consume_at_offset_from_begin(self):
         fluvio = Fluvio.connect()
 
@@ -146,7 +281,7 @@ def test_consume_deprecated(self):
 
         config = ConsumerConfig()
         stream = consumer.stream_with_config(Offset.beginning(), config)
-        for i in range(2):
+        for _ in range(2):
             records.append(bytearray(next(stream).value()).decode())
 
         self.assertEqual(len(records), 2)