allow filtering on consume

Author: rerpha

README.md

@@ -20,9 +20,7 @@ See `saluki --help` for all options.
 
 Use the `-g` flag to go the other way, ie. in the above example to consume the 9 messages _after_ offset 123456
 
-### Consume X of a certain schema(s)
-
-TODO
+You can also filter out messages to specific schema(s) with the `-f` flag, like the example above for `listen`.
 
 ## `sniff` - List all topics and their high, low watermarks and number of messages
 `saluki sniff mybroker:9092`

src/saluki/consume.py

@@ -14,6 +14,7 @@ def consume(
     num_messages: int = 1,
     offset: int | None = None,
     go_forwards: bool = False,
+    schemas_to_filter_to: list[str] | None = None,
 ) -> None:
     """
     consume from a topic and deserialise each message
@@ -24,6 +25,7 @@ def consume(
     :param num_messages: number of messages to 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
     :return: None
     """
     c = Consumer(
@@ -57,7 +59,7 @@ def consume(
     try:
         logger.info(f"Consuming {num_messages} messages")
         msgs = c.consume(num_messages)
-        deserialise_and_print_messages(msgs, partition)
+        deserialise_and_print_messages(msgs, partition, schemas_to_filter_to)
     except Exception:
         logger.exception("Got exception while consuming:")
     finally:

src/saluki/listen.py

@@ -11,13 +11,14 @@ def listen(
     broker: str,
     topic: str,
     partition: int | None = None,
-    schemas_to_filter_out: list[str] | None = None,
+    schemas_to_filter_to: list[str] | None = None,
 ) -> None:
     """
     Listen to a topic and deserialise each message
     :param broker: the broker address, including the port
     :param topic: the topic to use
     :param partition: the partition to listen to (default is all partitions in a given topic)
+    :param schemas_to_filter_to: schemas to filter when listening to messages
     :return: None
     """
     c = Consumer(
@@ -36,7 +37,7 @@ def listen(
         while True:
             msg = c.poll(1.0)
             deserialise_and_print_messages(
-                [msg], partition, schemas_to_filter_out=schemas_to_filter_out
+                [msg], partition, schemas_to_filter_to=schemas_to_filter_to
             )
     except KeyboardInterrupt:
         logger.debug("finished listening")

src/saluki/main.py

@@ -72,13 +72,6 @@ def main() -> None:
         parents=[topic_parser, consumer_parser],
     )
 
-    #### NEW FEATURES HERE PLZ
-    # saluki consume x messages of y schema
-    # saluki consume -f pl72 mybroker:9092/XXX_runInfo -m 10 # get the last pl72 run starts
-
-    # 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#
-
     play_parser = sub_parsers.add_parser(
         _PLAY,
         help="replay mode - replay data into another topic",
@@ -111,12 +104,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,
+            broker, topic, args.partition, args.messages, args.offset, args.go_forwards, args.filter
         )
     elif args.command == _PLAY:
         src_broker, src_topic = parse_kafka_uri(args.topics[0])

src/saluki/utils.py

@@ -32,7 +32,7 @@ def fallback_deserialiser(payload: bytes) -> str:
 
 
 def deserialise_and_print_messages(
-    msgs: List[Message], partition: int | None, schemas_to_filter_out: list[str] | None
+    msgs: List[Message], partition: int | None, schemas_to_filter_to: list[str] | None = None
 ) -> None:
     for msg in msgs:
         try:
@@ -44,7 +44,7 @@ def deserialise_and_print_messages(
             if partition is not None and msg.partition() != partition:
                 continue
             schema, deserialised = _try_to_deserialise_message(msg.value())
-            if schemas_to_filter_out is not None and schema in schemas_to_filter_out:
+            if schemas_to_filter_to is not None and schema in schemas_to_filter_to:
                 break
             time = _parse_timestamp(msg)
             logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")