get timestamps working for replay and consume

Author: rerpha

pyproject.toml

@@ -7,7 +7,7 @@ name = "saluki"
 dynamic = ["version"]
 dependencies = [
     "ess-streaming-data-types",
-    "confluent-kafka",
+    "confluent-kafka>=2.12.1", # for produce_batch in play()
     "python-dateutil",
     "tzdata"
 ]

src/saluki/consume.py

@@ -16,6 +16,7 @@ def consume(
     offset: int | None = None,
     go_forwards: bool = False,
     schemas_to_filter_to: list[str] | None = None,
+    timestamp: int|None = None,
 ) -> None:
     """
     consume from a topic and deserialise each message
@@ -27,6 +28,7 @@ def consume(
     :param offset: offset to consume from/to
     :param go_forwards: whether to consume forwards or backwards
     :param schemas_to_filter_to: schemas in messages to filter to
+    :param timestamp: optionally a timestamp as a starting point
     :return: None
     """
     c = Consumer(
@@ -41,6 +43,12 @@ def consume(
         }
     )
 
+    if timestamp is not None:
+        offset = c.offsets_for_times(
+            [TopicPartition(topic, partition, timestamp)]
+        )[0].offset
+        logger.debug(f"offset for timestamp {timestamp} is {offset}")
+
     if go_forwards:
         if offset is None:
             raise ValueError("Can't go forwards without an offset")

src/saluki/main.py

@@ -71,10 +71,13 @@ def main() -> None:
         required=False,
         default=1,
     )
-    consumer_mode_parser.add_argument(
-        "-o", "--offset", help="offset to consume from", type=int, required=False
-    )
+
     consumer_mode_parser.add_argument("-g", "--go-forwards", required=False, action="store_true")
+    cg = consumer_mode_parser.add_mutually_exclusive_group(required=False)
+    cg.add_argument(
+        "-o", "--offset", help="offset to consume from", type=int,
+    )
+    cg.add_argument("-t", "--timestamp", help="timestamp to consume from", type=dateutil_parsable_or_unix_timestamp)
 
     listen_parser = sub_parsers.add_parser(  # noqa: F841
         _LISTEN,
@@ -118,7 +121,7 @@ def main() -> None:
     elif args.command == _CONSUME:
         broker, topic = parse_kafka_uri(args.topic)
         consume(
-            broker, topic, args.partition, args.messages, args.offset, args.go_forwards, args.filter
+            broker, topic, args.partition, args.messages, args.offset, args.go_forwards, args.filter, args.timestamp
         )
     elif args.command == _PLAY:
         src_broker, src_topic = parse_kafka_uri(args.topics[0])

src/saluki/play.py

@@ -1,6 +1,5 @@
 import logging
 import uuid
-from time import sleep
 from confluent_kafka import Consumer, Producer, TopicPartition
 
 logger = logging.getLogger("saluki")
@@ -17,6 +16,7 @@ def play(
     """
     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 uses partition 0.
+    It also does not copy message timestamps.
 
     :param src_broker: The source broker, including port.
     :param src_topic: The topic to replay data from.
@@ -62,27 +62,12 @@ def play(
 
     num_messages = stop_offset.offset - start_offset.offset + 1
 
-    def delivery_report(err, msg):
-        """ Called once for each message produced to indicate delivery result.
-            Triggered by poll() or flush()."""
-        if err is not None:
-            logger.error('Message delivery failed: {}'.format(err))
-        else:
-            logger.debug('Message delivered to {} [{}]'.format(msg.topic(), msg.partition()))
-
     try:
         msgs = consumer.consume(num_messages)
         logger.debug(f"finished consuming {num_messages} messages")
         consumer.close()
-        # logger.debug(f"{msgs}")
-        for message in msgs:
-            producer.poll(0)
-            producer.produce(dest_topic, message.value(), message.key(), callback=delivery_report)
-        # producer.produce_batch(dest_topic, [{'key': message.key(), 'value': message.value()} for message in msgs])
-        # producer.poll()
+        producer.produce_batch(dest_topic, [{'key': message.key(), 'value': message.value()} for message in msgs])
         logger.debug(f"flushing producer. len(p): {len(producer)}")
-        # while len(producer): producer.flush()
-
         producer.flush(timeout=10)
 
         logger.debug(f"length after flushing: {len(producer)}")

src/saluki/utils.py

@@ -101,7 +101,7 @@ def dateutil_parsable_or_unix_timestamp(inp: str) -> int:
     try:
         try:
             return int(round(parse(inp).timestamp() * 1000))
-        except ParserError:
+        except (ParserError, OverflowError):
             logger.debug(f"Failed to parse {inp} as a dateutil parsable. Falling back to unix timestamp")
             return int(inp)
     except ValueError: