make a start on better error handlingand tests

Author: rerpha

src/saluki/__init__.py

@@ -1,60 +1 @@
-import logging
-from typing import Tuple, List
-import datetime
-from confluent_kafka import Message
-from streaming_data_types import DESERIALISERS
-from streaming_data_types.exceptions import StreamingDataTypesException
-from streaming_data_types.utils import get_schema
 
-logger = logging.getLogger("saluki")
-
-
-def _fallback_deserialiser(payload: bytes) -> str:
-    return payload.decode()
-
-
-def try_to_deserialise_message(payload: bytes) -> Tuple[str, str]:
-    logger.debug(f"got some data: {payload}")
-    schema = get_schema(payload)
-    deserialiser = (
-        _fallback_deserialiser  # Fall back to this if we need to so data isn't lost
-    )
-    try:
-        deserialiser = DESERIALISERS[schema]
-    except StreamingDataTypesException:
-        pass  # TODO
-    except KeyError:
-        pass
-    return schema, deserialiser(payload)
-
-
-def _deserialise_and_print_messages(msgs: List[Message], partition: int | None) -> None:
-    for msg in msgs:
-        if msg is None:
-            continue
-        if msg.error():
-            logger.error("Consumer error: {}".format(msg.error()))
-            continue
-        if partition is not None and msg.partition() != partition:
-            continue
-        schema, deserialised = try_to_deserialise_message(msg.value())
-        time = _parse_timestamp(msg)
-        logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")
-
-
-def _parse_timestamp(msg: Message) -> str:
-    """
-    Parse a message timestamp.
-
-    See https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.Message.timestamp
-    :param msg: the message to parse.
-    :return: either the string-formatted timestamp or "Unknown" if not able to parse.
-    """
-    timestamp_type, timestamp_ms_from_epoch = msg.timestamp()
-    if timestamp_type == 1:  # TIMESTAMP_CREATE_TIME
-        return datetime.datetime.fromtimestamp(timestamp_ms_from_epoch / 1000).strftime(
-            "%Y-%m-%d %H:%M:%S.%f"
-        )
-    else:
-        # TIMESTAMP_NOT_AVAILABLE or TIMESTAMP_LOG_APPEND_TIME
-        return "Unknown"

src/saluki/consume.py

@@ -1,7 +1,7 @@
 import logging
 
 from confluent_kafka import Consumer, TopicPartition
-from saluki import _deserialise_and_print_messages
+from saluki.utils import _deserialise_and_print_messages
 
 logger = logging.getLogger("saluki")
 
@@ -53,15 +53,15 @@ def consume(
                 - num_messages
             )
 
-    logger.info(f"starting at offset {start}")
+    logger.info(f"Starting at offset {start}")
     c.assign([TopicPartition(topic, partition, start)])
 
     try:
-        logger.info(f"consuming {num_messages} messages")
+        logger.info(f"Consuming {num_messages} messages")
         msgs = c.consume(num_messages)
         _deserialise_and_print_messages(msgs, partition)
-    except Exception as e:
-        logger.error(e)
+    except Exception:
+        logger.exception("Got exception while consuming:")
     finally:
-        logger.debug(f"closing consumer {c}")
+        logger.debug(f"Closing consumer {c}")
         c.close()

src/saluki/listen.py

@@ -1,7 +1,7 @@
 import logging
 
 from confluent_kafka import Consumer, TopicPartition
-from saluki import _deserialise_and_print_messages
+from saluki.utils import _deserialise_and_print_messages
 
 logger = logging.getLogger("saluki")
 
@@ -32,6 +32,8 @@ def listen(broker: str, topic: str, partition: int | None = None) -> None:
             _deserialise_and_print_messages([msg], partition)
     except KeyboardInterrupt:
         logger.debug("finished listening")
+    except Exception:
+        logger.exception("Got exception while listening:")
     finally:
         logger.debug(f"closing consumer {c}")
         c.close()

src/saluki/main.py

@@ -1,10 +1,10 @@
 import argparse
 import logging
 import sys
-from typing import Tuple
 
 from saluki.consume import consume
 from saluki.listen import listen
+from saluki.utils import parse_kafka_uri
 
 logger = logging.getLogger("saluki")
 logging.basicConfig(level=logging.INFO)
@@ -14,38 +14,16 @@
 _CONSUME = "consume"
 
 
-def parse_kafka_uri(uri: str) -> Tuple[str, str]:
-    """Parse Kafka connection URI.
-
-    A broker hostname/ip must be present.
-    If username is provided, a SASL mechanism must also be provided.
-    Any other validation must be performed in the calling code.
-    """
-    security_protocol, tail = uri.split("+") if "+" in uri else ("", uri)
-    sasl_mechanism, tail = tail.split("\\") if "\\" in tail else ("", tail)
-    username, tail = tail.split("@") if "@" in tail else ("", tail)
-    broker, topic = tail.split("/") if "/" in tail else (tail, "")
-    if not broker:
-        raise RuntimeError(
-            f"Unable to parse URI {uri}, broker not defined. URI should be of form"
-            f" [PROTOCOL+SASL_MECHANISM\\username@]broker:9092"
-        )
-    if username and not (security_protocol and sasl_mechanism):
-        raise RuntimeError(
-            f"Unable to parse URI {uri}, PROTOCOL or SASL_MECHANISM not defined."
-            f" URI should be of form [PROTOCOL+SASL_MECHANISM\\username@]broker:9092"
-        )
-    return broker, topic
-
-
 def main() -> None:
     parser = argparse.ArgumentParser(
         prog="saluki",
         description="serialise/de-serialise flatbuffers and consume/produce from/to kafka",
     )
 
     parent_parser = argparse.ArgumentParser(add_help=False)
-    parent_parser.add_argument("topic", type=str, help="Kafka topic. format is broker<:port>/topic")
+    parent_parser.add_argument(
+        "topic", type=str, help="Kafka topic. format is broker<:port>/topic"
+    )
 
     parent_parser.add_argument(
         "-X",
@@ -63,7 +41,9 @@ def main() -> None:
         type=argparse.FileType("a"),
     )
 
-    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"
+    )
 
     consumer_parser = argparse.ArgumentParser(add_help=False)
     consumer_parser.add_argument(
@@ -88,16 +68,24 @@ def main() -> None:
     consumer_mode_parser.add_argument(
         "-o", "--offset", help="offset to consume from", type=int, required=False
     )
-    consumer_mode_parser.add_argument("-s", "--schema", required=False, default="auto", type=str)
-    consumer_mode_parser.add_argument("-g", "--go-forwards", required=False, action="store_true")
-    consumer_mode_parser.add_argument("-p", "--partition", required=False, type=int, default=0)
+    consumer_mode_parser.add_argument(
+        "-s", "--schema", required=False, default="auto", type=str
+    )
+    consumer_mode_parser.add_argument(
+        "-g", "--go-forwards", required=False, action="store_true"
+    )
+    consumer_mode_parser.add_argument(
+        "-p", "--partition", required=False, type=int, default=0
+    )
 
     listen_parser = sub_parsers.add_parser(
         _LISTEN,
         help="listen mode - listen until KeyboardInterrupt",
         parents=[parent_parser, consumer_parser],
     )
-    listen_parser.add_argument("-p", "--partition", required=False, type=int, default=None)
+    listen_parser.add_argument(
+        "-p", "--partition", required=False, type=int, default=None
+    )
 
     if len(sys.argv) == 1:
         parser.print_help()

src/saluki/utils.py

@@ -0,0 +1,97 @@
+import datetime
+import logging
+from typing import Tuple, List
+
+from confluent_kafka import Message
+
+from streaming_data_types import DESERIALISERS
+from streaming_data_types.exceptions import ShortBufferException
+from streaming_data_types.utils import get_schema
+
+
+logger = logging.getLogger("saluki")
+
+
+def __try_to_deserialise_message(payload: bytes) -> Tuple[str | None, str | None]:
+    logger.debug(f"got some data: {payload}")
+    try:
+        schema = get_schema(payload)
+    except ShortBufferException:
+        schema = None
+
+    logger.debug(f"schema: {schema}")
+
+    try:
+        deserialiser = DESERIALISERS[schema]
+    except KeyError:
+        logger.exception(f"Invalid schema: {schema}, falling back to raw bytes decode")
+
+        def fallback_deserialiser(payload: bytes) -> str:
+            return payload.decode()
+
+        deserialiser = (
+            fallback_deserialiser  # Fall back to this if we need to so data isn't lost
+        )
+
+    logger.debug(f"Deserialiser: {deserialiser}")
+
+    try:
+        ret = deserialiser(payload)
+    except Exception as e:
+        raise e
+
+    return schema, ret
+
+
+def _deserialise_and_print_messages(msgs: List[Message], partition: int | None) -> None:
+    for msg in msgs:
+        try:
+            if msg is None:
+                continue
+            if msg.error():
+                logger.error("Consumer error: {}".format(msg.error()))
+                continue
+            if partition is not None and msg.partition() != partition:
+                continue
+            schema, deserialised = __try_to_deserialise_message(msg.value())
+            time = _parse_timestamp(msg)
+            logger.info(f"{msg.offset()} ({time}):({schema}) {deserialised}")
+        except Exception as e:
+            logger.exception(f"Got error while deserialising: {e}")
+
+
+def _parse_timestamp(msg: Message) -> str:
+    """
+    Parse a message timestamp.
+
+    See https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.Message.timestamp
+    :param msg: the message to parse.
+    :return: either the string-formatted timestamp or "Unknown" if not able to parse.
+    """
+    timestamp_type, timestamp_ms_from_epoch = msg.timestamp()
+    if timestamp_type == 1:  # TIMESTAMP_CREATE_TIME
+        return datetime.datetime.fromtimestamp(timestamp_ms_from_epoch / 1000).strftime(
+            "%Y-%m-%d %H:%M:%S.%f"
+        )
+    else:
+        # TIMESTAMP_NOT_AVAILABLE or TIMESTAMP_LOG_APPEND_TIME
+        return "Unknown"
+
+
+def parse_kafka_uri(uri: str) -> Tuple[str, str]:
+    """Parse Kafka connection URI.
+
+    A broker hostname/ip must be present.
+    If username is provided, a SASL mechanism must also be provided.
+    Any other validation must be performed in the calling code.
+    """
+    broker, topic = uri.split("/") if "/" in uri else (uri, "")
+    if not topic:
+        raise RuntimeError(
+            f"Unable to parse URI {uri}, topic not defined. URI should be of form"
+            f" broker[:port]/topic"
+        )
+    return (
+        broker,
+        topic,
+    )

tests/test_consume.py

@@ -0,0 +1,7 @@
+# test that tries go_forwards with no offset
+
+# test that tries going forwards that consumes from offset
+
+# test that checks start offset
+
+# test that catches exception

tests/test_deserialiser.py

@@ -0,0 +1,37 @@
+import pytest
+
+from saluki.utils import parse_kafka_uri
+
+
+# test with normal payload
+
+# test with fallback serialiser when schema not found (empty payload)
+
+# test with schema that looks ok, but not in list of deserialisers
+
+# test exception while deserialising
+
+# test _parse_timestamp
+
+def test_uri_with_broker_name_and_topic_successfully_split():
+    test_broker = "localhost"
+    test_topic = "some_topic"
+    test_uri = f"{test_broker}/{test_topic}"
+    broker, topic = parse_kafka_uri(test_uri)
+    assert broker == test_broker
+    assert topic == test_topic
+
+
+def test_uri_with_port_after_broker_is_included_in_broker_output():
+    test_broker = "localhost:9092"
+    test_topic = "some_topic"
+    test_uri = f"{test_broker}/{test_topic}"
+    broker, topic = parse_kafka_uri(test_uri)
+    assert broker == test_broker
+    assert topic == test_topic
+
+
+def test_uri_with_no_topic():
+    test_broker = "some_broker"
+    with pytest.raises(RuntimeError):
+        parse_kafka_uri(test_broker)