docs, formatting, fix listen -f

Author: rerpha

README.md

@@ -8,14 +8,14 @@ Also allows replaying data in a topic.
 # Usage
 See `saluki --help` for all options. 
 
-## Listen to a topic for updates
+## `listen` - Listen to a topic for updates
 `saluki listen mybroker:9092/mytopic` - This will listen for updates for `mytopic` on `mybroker`. 
 
 ### Filter to specific schemas
 
-TODO
+`saluki listen mybroker:9092/mytopic -f f144 -f f142` - This will listen for updates but ignore messages with schema IDs of `f142` or `f144`
 
-## Consume from a topic
+## `consume`- Consume from a topic
 `saluki consume mybroker:9092/mytopic -p 1 -o 123456 -m 10` - This will print 9 messages before (and inclusively the offset specified) offset `123456` of `mytopic` on `mybroker`, in partition 1.
 
 Use the `-g` flag to go the other way, ie. in the above example to consume the 9 messages _after_ offset 123456
@@ -24,10 +24,10 @@ Use the `-g` flag to go the other way, ie. in the above example to consume the 9
 
 TODO
 
-## List all topics and their high, low watermarks and number of messages
-TODO
+## `sniff` - List all topics and their high, low watermarks and number of messages
+`saluki sniff mybroker:9092`
 
-## Replay data from one topic to another
+## `play` - Replay data from one topic to another
 
 ### Between offsets
 

src/saluki/consume.py

@@ -47,7 +47,9 @@ def consume(
             start = offset - num_messages + 1
         else:
             start = (
-                c.get_watermark_offsets(TopicPartition(topic, partition), cached=False)[1]
+                c.get_watermark_offsets(TopicPartition(topic, partition), cached=False)[
+                    1
+                ]
                 - num_messages
             )
 

src/saluki/listen.py

@@ -7,7 +7,12 @@
 logger = logging.getLogger("saluki")
 
 
-def listen(broker: str, topic: str, partition: int | None = None, filter: list[str] | None = None) -> None:
+def listen(
+    broker: str,
+    topic: str,
+    partition: int | None = None,
+    schemas_to_filter_out: list[str] | None = None,
+) -> None:
     """
     Listen to a topic and deserialise each message
     :param broker: the broker address, including the port
@@ -30,7 +35,9 @@ def listen(broker: str, topic: str, partition: int | None = None, filter: list[s
         logger.info(f"listening to {broker}/{topic}")
         while True:
             msg = c.poll(1.0)
-            deserialise_and_print_messages([msg], partition)
+            deserialise_and_print_messages(
+                [msg], partition, schemas_to_filter_out=schemas_to_filter_out
+            )
     except KeyboardInterrupt:
         logger.debug("finished listening")
     finally:

src/saluki/main.py

@@ -18,16 +18,17 @@
 _BURY = "bury"
 _DIG = "dig"
 
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         prog="saluki",
         description="serialise/de-serialise flatbuffers and consume/produce from/to kafka",
     )
 
     topic_parser = argparse.ArgumentParser(add_help=False)
-    #TODO this needs restructuring. consider having a topic_parser with partition options etc. 
-    # because -m and -g does not make sense for saluki listen. 
-    topic_parser.add_argument("topic", type=str, help="Kafka topic. format is broker<:port>/topic")
+    topic_parser.add_argument(
+        "topic", type=str, help="Kafka topic. format is broker<:port>/topic"
+    )
 
     topic_parser.add_argument(
         "-X",
@@ -37,8 +38,11 @@ def main() -> None:
         default=None,
     )
     topic_parser.add_argument("-p", "--partition", required=False, type=int, default=0)
+    topic_parser.add_argument("-f", "--filter", required=False, action="append")
 
-    sub_parsers = parser.add_subparsers(help="sub-command help", required=True, dest="command")
+    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)
@@ -66,10 +70,11 @@ def main() -> None:
     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")
-    consumer_mode_parser.add_argument("-f", "--filter", required=False, action="append")
+    consumer_mode_parser.add_argument(
+        "-g", "--go-forwards", required=False, action="store_true"
+    )
 
-    listen_parser = sub_parsers.add_parser( # noqa: F841
+    listen_parser = sub_parsers.add_parser(  # noqa: F841
         _LISTEN,
         help="listen mode - listen until KeyboardInterrupt",
         parents=[topic_parser, consumer_parser],
@@ -92,7 +97,7 @@ def main() -> None:
     # 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 - push data from dump generated by saluki bury to topic
     # saluki dig mybroker:9092/topicname -p 0 outputfile
 
     play_parser = sub_parsers.add_parser(
@@ -102,16 +107,31 @@ def main() -> None:
     )
     play_parser.add_argument("topics", type=str, nargs=2, help="SRC topic DEST topic")
     g = play_parser.add_mutually_exclusive_group(required=True)
-    g.add_argument("-o", "--offsets", help="offsets to replay between (inclusive)", type=int, nargs=2)
-    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)
+    g.add_argument(
+        "-o",
+        "--offsets",
+        help="offsets to replay between (inclusive)",
+        type=int,
+        nargs=2,
+    )
+    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()
         sys.exit(1)
     args = parser.parse_args()
 
-    if 'kafka_config' in args and args.kafka_config is not None:
+    if "kafka_config" in args and args.kafka_config is not None:
         raise NotImplementedError("-X is not implemented yet.")
 
     if args.command == _LISTEN:
@@ -131,14 +151,28 @@ def main() -> None:
         src_broker, src_topic = parse_kafka_uri(args.topics[0])
         dest_broker, dest_topic = parse_kafka_uri(args.topics[1])
 
-        print(f"SOURCE BROKER: {src_broker}, SOURCE TOPIC: {src_topic}, DEST BROKER: {dest_broker}, DEST TOPIC: {dest_topic} ")
+        print(
+            f"SOURCE BROKER: {src_broker}, SOURCE TOPIC: {src_topic}, DEST BROKER: {dest_broker}, DEST TOPIC: {dest_topic} "
+        )
         if args.offsets is not None:
-            print(f"Replaying {src_broker}/{src_topic} between offsets {args.offsets[0]} and {args.offsets[1]} to {dest_broker}/{dest_topic}")
+            print(
+                f"Replaying {src_broker}/{src_topic} between offsets {args.offsets[0]} and {args.offsets[1]} to {dest_broker}/{dest_topic}"
+            )
         elif args.timestamps is not None:
-            print(f"Replaying {src_broker}/{src_topic} between timestamps {args.timestamps[0]} and {args.timestamps[1]} to {dest_broker}/{dest_topic}")
-
-        if input("OK? (y/n)").lower() == 'y':
-            play(src_broker, src_topic, dest_broker, dest_topic, args.offsets, args.timestamps, args.chunks)
+            print(
+                f"Replaying {src_broker}/{src_topic} between timestamps {args.timestamps[0]} and {args.timestamps[1]} to {dest_broker}/{dest_topic}"
+            )
+
+        if input("OK? (y/n)").lower() == "y":
+            play(
+                src_broker,
+                src_topic,
+                dest_broker,
+                dest_topic,
+                args.offsets,
+                args.timestamps,
+                args.chunks,
+            )
             print("replayed")
     elif args.command == _SNIFF:
         sniff(args.broker)
@@ -147,5 +181,6 @@ def main() -> None:
     elif args.command == _DIG:
         pass
 
+
 if __name__ == "__main__":
     main()

src/saluki/play.py

@@ -1,12 +1,24 @@
 from confluent_kafka import Consumer, Producer
 
 
-def play(src_broker: str, src_topic: str, dest_broker: str, dest_topic: str, offsets:list[int]|None, timestamps:list[int]|None, chunks:int) -> None:
-    consumer = Consumer(        {
+def play(
+    src_broker: str,
+    src_topic: str,
+    dest_broker: str,
+    dest_topic: str,
+    offsets: list[int] | None,
+    timestamps: list[int] | None,
+    chunks: int,
+) -> None:
+    consumer = Consumer(
+        {
             "bootstrap.servers": src_broker,
             "group.id": "saluki-play",
-        })
-    producer = Producer({
-        "bootstrap.servers": dest_broker,
-    })
+        }
+    )
+    producer = Producer(
+        {
+            "bootstrap.servers": dest_broker,
+        }
+    )
     pass

src/saluki/sniff.py

@@ -1,20 +1,21 @@
 from confluent_kafka import Consumer, TopicPartition
 from confluent_kafka.admin import AdminClient
 
+
 def sniff(broker: str):
-    a = AdminClient({'bootstrap.servers': broker})
-    c = Consumer({'bootstrap.servers': broker, 'group.id': 'saluki-sniff'})
+    a = AdminClient({"bootstrap.servers": broker})
+    c = Consumer({"bootstrap.servers": broker, "group.id": "saluki-sniff"})
     t = a.list_topics(timeout=5)
     print(f"Cluster ID: {t.cluster_id}")
     print(f"Brokers:")
     [print(f"\t{value}") for value in t.brokers.values()]
 
     print(f"Topics:")
 
-    for k,v in t.topics.items():
+    for k, v in t.topics.items():
         partitions = v.partitions.keys()
         print(f"\t{k}:")
         for p in partitions:
             tp = TopicPartition(k, p)
             low, high = c.get_watermark_offsets(tp)
-            print(f"\t\tlow:{low}, high:{high}, num_messages:{high-low}")
+            print(f"\t\tlow:{low}, high:{high}, num_messages:{high - low}")

src/saluki/utils.py

@@ -23,15 +23,19 @@ def _try_to_deserialise_message(payload: bytes) -> Tuple[str | None, str | None]
     def fallback_deserialiser(payload: bytes) -> str:
         return payload.decode()
 
-    deserialiser = DESERIALISERS.get(schema if schema is not None else "", fallback_deserialiser)
+    deserialiser = DESERIALISERS.get(
+        schema if schema is not None else "", fallback_deserialiser
+    )
     logger.debug(f"Deserialiser: {deserialiser}")
 
     ret = deserialiser(payload)
 
     return schema, ret
 
 
-def deserialise_and_print_messages(msgs: List[Message], partition: int | None, filter: list[str] | None) -> None:
+def deserialise_and_print_messages(
+    msgs: List[Message], partition: int | None, schemas_to_filter_out: list[str] | None
+) -> None:
     for msg in msgs:
         try:
             if msg is None:
@@ -42,8 +46,7 @@ def deserialise_and_print_messages(msgs: List[Message], partition: int | None, f
             if partition is not None and msg.partition() != partition:
                 continue
             schema, deserialised = _try_to_deserialise_message(msg.value())
-            if filter is not None and schema not in filter:
-                # ignore
+            if schemas_to_filter_out is not None and schema in schemas_to_filter_out:
                 break
             time = _parse_timestamp(msg)
             logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")

tests/test_listen.py

@@ -16,7 +16,9 @@ def test_listen_with_partition_assigns_to_partition():
         mock.patch("saluki.listen.Consumer") as c,
     ):
         listen("somebroker", "sometopic", partition=expected_partition)
-        c.return_value.assign.assert_called_with([TopicPartition(topic, expected_partition)])
+        c.return_value.assign.assert_called_with(
+            [TopicPartition(topic, expected_partition)]
+        )
 
 
 def test_keyboard_interrupt_causes_consumer_to_close():