play working for offsets

Author: rerpha

src/saluki/listen.py

@@ -23,7 +23,7 @@ def listen(
     c = Consumer(
         {
             "bootstrap.servers": broker,
-            "group.id": "saluki",
+            "group.id": "saluki-listen",
             "auto.offset.reset": "latest",
             "enable.auto.commit": False,
         }

src/saluki/main.py

@@ -93,13 +93,6 @@ def main() -> None:
     # saluki consume x messages of y or z schema
     # saluki consume -f pl72,6s4t mybroker:9092/XXX_runInfo -m 10 # get the last pl72 run starts or 6s4t run stops#
 
-    # saluki bury - dump data on topic to file
-    # saluki bury mybroker:9092/topicname -p 0 -o startoffset finishoffset outputfile
-    # saluki bury mybroker:9092/topicname -p 0 -t starttimestamp finishtimestamp outputfile
-
-    # saluki dig - push data from dump generated by saluki bury to topic
-    # saluki dig mybroker:9092/topicname -p 0 outputfile
-
     play_parser = sub_parsers.add_parser(
         _PLAY,
         help="replay mode - replay data into another topic",
@@ -117,14 +110,6 @@ def main() -> None:
     g.add_argument(
         "-t", "--timestamps", help="timestamps to replay between", type=str, nargs=2
     )
-    g.add_argument(
-        "-c",
-        "--chunk",
-        help="forward in chunks. ie to avoid storing a huge list in memory",
-        default=0,
-        type=int,
-        required=False,
-    )
 
     if len(sys.argv) == 1:
         parser.print_help()
@@ -171,7 +156,6 @@ def main() -> None:
                 dest_topic,
                 args.offsets,
                 args.timestamps,
-                args.chunks,
             )
             print("replayed")
     elif args.command == _SNIFF:

src/saluki/play.py

@@ -1,4 +1,7 @@
-from confluent_kafka import Consumer, Producer
+import logging
+from confluent_kafka import Consumer, Producer, TopicPartition
+
+logger = logging.getLogger("saluki")
 
 
 def play(
@@ -8,8 +11,21 @@ def play(
     dest_topic: str,
     offsets: list[int] | None,
     timestamps: list[int] | None,
-    chunks: int,
 ) -> None:
+    """
+    Replay data from src_topic to dest_topic between the offsets OR timestamps specified.
+    This currently assumes contiguous data in a topic (ie. no log compaction) and only uses partition 0.
+
+    :param src_broker: The source broker, including port.
+    :param src_topic: The topic to replay data from.
+    :param dest_broker: The destination broker, including port.
+    :param dest_topic: The topic to replay data to.
+    :param offsets: The start and finish offsets to replay data from.
+    :param timestamps: The start and finish timestamps to replay data from.
+    """
+
+    print(f"ARGS: {src_broker}, {src_topic}, {dest_broker}, {dest_topic}, {offsets}, {timestamps}")
+
     consumer = Consumer(
         {
             "bootstrap.servers": src_broker,
@@ -21,4 +37,27 @@ def play(
             "bootstrap.servers": dest_broker,
         }
     )
-    pass
+    src_partition = 0
+
+    if timestamps is not None:
+        start_offset, stop_offset = consumer.offsets_for_times([
+            TopicPartition(src_topic, src_partition, timestamps[0]),
+            TopicPartition(src_topic, src_partition, timestamps[1]),
+        ])
+    else:
+        start_offset = TopicPartition(src_topic, src_partition, offsets[0])
+        stop_offset = TopicPartition(src_topic, src_partition, offsets[1])
+    consumer.assign([start_offset])
+
+    num_messages = stop_offset.offset - start_offset.offset + 1
+
+    try:
+        msgs = consumer.consume(num_messages)
+        [producer.produce(dest_topic, message.value(), message.key()) for message in msgs]
+        producer.flush()
+    except Exception:
+        logger.exception("Got exception while consuming:")
+    finally:
+        logger.debug(f"Closing consumer {consumer}")
+        consumer.close()
+