backing up again

Author: rerpha

src/saluki/main.py

@@ -2,12 +2,11 @@
 import logging
 import sys
 
-from dateutil.parser import parse, ParserError
 from saluki.consume import consume
 from saluki.listen import listen
 from saluki.play import play
 from saluki.sniff import sniff
-from saluki.utils import parse_kafka_uri
+from saluki.utils import parse_kafka_uri, dateutil_parsable_or_unix_timestamp
 
 logger = logging.getLogger("saluki")
 logging.basicConfig(level=logging.INFO)
@@ -18,22 +17,14 @@
 _SNIFF = "sniff"
 
 
-def _dateutil_parsable_or_unix_timestamp(inp: str) -> float:
-    try:
-        try:
-            return parse(inp).timestamp()
-        except ParserError:
-            return float(inp)
-    except ValueError:
-        raise argparse.ArgumentTypeError(f"timestamp {inp} is not parsable by dateutil.parse() and is not a unix timestamp")
-
-
 def main() -> None:
     parser = argparse.ArgumentParser(
         prog="saluki",
         description="serialise/de-serialise flatbuffers and consume/produce from/to kafka",
     )
-    parser.add_argument(
+    common_options = argparse.ArgumentParser(add_help=False)
+    common_options.add_argument("-v", "--verbose", help="show DEBUG logs", action='store_true')
+    common_options.add_argument(
         "-l",
         "--log-file",
         help="filename to output all data to",
@@ -57,8 +48,8 @@ def main() -> None:
 
     sub_parsers = parser.add_subparsers(help="sub-command help", required=True, dest="command")
 
-    sniff_parser = sub_parsers.add_parser(_SNIFF, help="sniff - broker metadata")
-    sniff_parser.add_argument("broker", type=str, help="broker, optionally suffixed with a topic name")
+    sniff_parser = sub_parsers.add_parser(_SNIFF, help="sniff - broker metadata", parents=[common_options])
+    sniff_parser.add_argument("broker", type=str, help="broker, optionally suffixed with a topic name to filter to")
 
     consumer_parser = argparse.ArgumentParser(add_help=False)
     consumer_parser.add_argument(
@@ -70,7 +61,7 @@ def main() -> None:
     )
 
     consumer_mode_parser = sub_parsers.add_parser(
-        _CONSUME, help="consumer mode", parents=[topic_parser, consumer_parser]
+        _CONSUME, help="consumer mode", parents=[topic_parser, consumer_parser, common_options]
     )
     consumer_mode_parser.add_argument(
         "-m",
@@ -88,13 +79,13 @@ def main() -> None:
     listen_parser = sub_parsers.add_parser(  # noqa: F841
         _LISTEN,
         help="listen mode - listen until KeyboardInterrupt",
-        parents=[topic_parser, consumer_parser],
+        parents=[topic_parser, consumer_parser, common_options],
     )
 
     play_parser = sub_parsers.add_parser(
         _PLAY,
         help="replay mode - replay data into another topic",
-        parents=[],
+        parents=[common_options],
     )
     play_parser.add_argument("topics", type=str, nargs=2, help="SRC topic DEST topic")
     g = play_parser.add_mutually_exclusive_group(required=True)
@@ -105,13 +96,16 @@ def main() -> None:
         type=int,
         nargs=2,
     )
-    g.add_argument("-t", "--timestamps", help='timestamps to replay between in ISO8601 or RFC3339 format ie. "2025-11-17 07:00:00"  ', type=_dateutil_parsable_or_unix_timestamp, nargs=2)
+    g.add_argument("-t", "--timestamps", help='timestamps to replay between in ISO8601 or RFC3339 format ie. "2025-11-17 07:00:00 or as a unix timestamp"  ', type=dateutil_parsable_or_unix_timestamp, nargs=2)
 
     if len(sys.argv) == 1:
         parser.print_help()
         sys.exit(1)
     args = parser.parse_args()
 
+    if args.verbose:
+        logger.setLevel(logging.DEBUG)
+
     if args.log_file:
         logger.addHandler(logging.FileHandler(args.log_file.name))
 
@@ -141,8 +135,10 @@ def main() -> None:
     elif args.command == _SNIFF:
         try:
             broker, topic = parse_kafka_uri(args.broker)
+            logger.debug(f"Sniffing single topic {topic} on broker {broker}")
             sniff(broker, topic)
         except RuntimeError:
+            logger.debug(f"Sniffing whole broker {args.broker}")
             sniff(args.broker)
 
 

src/saluki/play.py

@@ -1,6 +1,6 @@
 import logging
 import uuid
-
+from time import sleep
 from confluent_kafka import Consumer, Producer, TopicPartition
 
 logger = logging.getLogger("saluki")
@@ -40,27 +40,53 @@ def play(
     src_partition = 0
 
     if timestamps is not None:
-        start_offset, stop_offset = consumer.offsets_for_times(
+        logger.debug(f"getting offsets for times: {timestamps[0]} and {timestamps[1]}")
+        start_offset = consumer.offsets_for_times(
             [
                 TopicPartition(src_topic, src_partition, timestamps[0]),
-                TopicPartition(src_topic, src_partition, timestamps[1]),
+
             ]
-        )
+        )[0]
+        # See https://github.com/confluentinc/confluent-kafka-python/issues/1178 as to why offsets_for_times is called twice.
+        stop_offset = consumer.offsets_for_times([TopicPartition(src_topic, src_partition, timestamps[1])])[0]
     elif offsets is not None:
         start_offset = TopicPartition(src_topic, src_partition, offsets[0])
         stop_offset = TopicPartition(src_topic, src_partition, offsets[1])
     else:
         raise ValueError("offsets and timestamps cannot both be None")
 
+    logger.debug(f"start_offset: {start_offset.offset}, stop_offset: {stop_offset.offset}")
+
+    logger.debug(f"assigning to offset {start_offset.offset}")
     consumer.assign([start_offset])
 
     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.produce(dest_topic, message.value(), message.key())
-        producer.flush()
+            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()
+        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)}")
+
     except Exception:
         logger.exception("Got exception while replaying:")
     finally:

src/saluki/utils.py

@@ -1,9 +1,11 @@
 import datetime
 import logging
+from argparse import ArgumentTypeError
 from typing import List, Tuple
 from zoneinfo import ZoneInfo
 
 from confluent_kafka import Message
+from dateutil.parser import parse, ParserError
 from streaming_data_types import DESERIALISERS
 from streaming_data_types.exceptions import ShortBufferException
 from streaming_data_types.utils import get_schema
@@ -47,7 +49,7 @@ def deserialise_and_print_messages(
             if schemas_to_filter_to is not None and schema not in schemas_to_filter_to:
                 continue
             time = _parse_timestamp(msg)
-            logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")
+            logger.info(f"(o:{msg.offset()},t:{time},s:{schema}) {deserialised}")
         except Exception as e:
             logger.exception(f"Got error while deserialising: {e}")
 
@@ -89,3 +91,18 @@ def parse_kafka_uri(uri: str) -> Tuple[str, str]:
         broker,
         topic,
     )
+
+
+def dateutil_parsable_or_unix_timestamp(inp: str) -> int:
+    """
+    Parse a dateutil string, if this fails then try to parse a unix timestamp.
+    This returns a unix timestamp as an int
+    """
+    try:
+        try:
+            return int(round(parse(inp).timestamp() * 1000))
+        except ParserError:
+            logger.debug(f"Failed to parse {inp} as a dateutil parsable. Falling back to unix timestamp")
+            return int(inp)
+    except ValueError:
+        raise ArgumentTypeError(f"timestamp {inp} is not parsable by dateutil.parse() and is not a unix timestamp")