Kafka: Expand default_msg_processor into a miniature decoding unit

amotl · amotl · commit 4a77c6e030af · 2025-07-13T18:32:52.000+02:00
- Accept a bunch of decoding options per `KafkaDecodingOptions`
- Provide a bunch of output formatting options per `KafkaEvent`
- Tie both elements together using `KafkaEventProcessor`

The machinery is effectively the same like before, but provides a few
more options to allow type decoding for Kafka event's key/value slots,
a selection mechanism to limit the output to specific fields only, and
a small projection mechanism to optionally drill down into a specific
field.

In combination, those decoding options allow users to relay
JSON-encoded Kafka event values directly into a destination table,
without any metadata wrappings.

The output formatter provides three different variants out of the box.
More variants can be added in the future, as other users or use cases
may have different requirements in the same area.

Most importantly, the decoding unit is very compact, so relevant tasks
don't need a corresponding transformation unit down the pipeline, to
keep the whole ensemble lean, in the very spirit of ingestr.
diff --git a/ingestr/src/kafka/__init__.py b/ingestr/src/kafka/__init__.py
@@ -16,8 +16,8 @@
 
 from .helpers import (
     KafkaCredentials,
+    KafkaEventProcessor,
     OffsetTracker,
-    default_msg_processor,
 )
 
 
@@ -29,9 +29,7 @@
 def kafka_consumer(
     topics: Union[str, List[str]],
     credentials: Union[KafkaCredentials, Consumer] = dlt.secrets.value,
-    msg_processor: Optional[
-        Callable[[Message], Dict[str, Any]]
-    ] = default_msg_processor,
+    msg_processor: Optional[Callable[[Message], Dict[str, Any]]] = None,
     batch_size: Optional[int] = 3000,
     batch_timeout: Optional[int] = 3,
     start_from: Optional[TAnyDateTime] = None,
@@ -60,6 +58,8 @@ def kafka_consumer(
     Yields:
         Iterable[TDataItem]: Kafka messages.
     """
+    msg_processor = msg_processor or KafkaEventProcessor().process
+
     if not isinstance(topics, list):
         topics = [topics]
 
diff --git a/ingestr/src/kafka/helpers.py b/ingestr/src/kafka/helpers.py
@@ -1,53 +1,86 @@
 from typing import Any, Dict, List, Optional
 
+import orjson
 from confluent_kafka import Consumer, Message, TopicPartition  # type: ignore
 from confluent_kafka.admin import TopicMetadata  # type: ignore
 from dlt import config
 from dlt.common import pendulum
 from dlt.common.configuration import configspec
 from dlt.common.configuration.specs import CredentialsConfiguration
-from dlt.common.time import ensure_pendulum_datetime
 from dlt.common.typing import DictStrAny, TSecretValue
-from dlt.common.utils import digest128
 
+from ingestr.src.kafka.model import KafkaDecodingOptions, KafkaEvent
 
-def default_msg_processor(msg: Message) -> Dict[str, Any]:
-    """Basic Kafka message processor.
 
-    Returns the message value and metadata. Timestamp consists of two values:
-    (type of the timestamp, timestamp). Type represents one of the Python
-    Kafka constants:
-        TIMESTAMP_NOT_AVAILABLE - Timestamps not supported by broker.
-        TIMESTAMP_CREATE_TIME - Message creation time (or source / producer time).
-        TIMESTAMP_LOG_APPEND_TIME - Broker receive time.
-
-    Args:
-        msg (confluent_kafka.Message): A single Kafka message.
+class KafkaEventProcessor:
+    """
+    A processor for Kafka events with processing stages and configuration capabilities.
 
-    Returns:
-        dict: Processed Kafka message.
+    It cycles through "decode", "deserialize" and "format".
     """
-    ts = msg.timestamp()
-    topic = msg.topic()
-    partition = msg.partition()
-    key = msg.key()
-    if key is not None:
-        key = key.decode("utf-8")
-
-    return {
-        "_kafka": {
-            "partition": partition,
-            "topic": topic,
-            "key": key,
-            "offset": msg.offset(),
-            "ts": {
-                "type": ts[0],
-                "value": ensure_pendulum_datetime(ts[1] / 1e3),
-            },
-            "data": msg.value().decode("utf-8"),
-        },
-        "_kafka_msg_id": digest128(topic + str(partition) + str(key)),
-    }
+
+    def __init__(self, options: Optional[KafkaDecodingOptions] = None):
+        self.options = options or KafkaDecodingOptions()
+
+    def process(self, msg: Message) -> Dict[str, Any]:
+        """
+        Progress Kafka event.
+
+        Returns the message value and metadata. Timestamp consists of two values:
+        (type of the timestamp, timestamp). Type represents one of the Python
+        Kafka constants:
+            TIMESTAMP_NOT_AVAILABLE - Timestamps not supported by broker.
+            TIMESTAMP_CREATE_TIME - Message creation time (or source / producer time).
+            TIMESTAMP_LOG_APPEND_TIME - Broker receive time.
+
+        Args:
+            msg (confluent_kafka.Message): A single Kafka message.
+
+        Returns:
+            dict: Processed Kafka message.
+        """
+
+        # Decode.
+        event = self.decode(msg)
+
+        # Deserialize.
+        self.deserialize(event)
+
+        # Format egress message based on input options.
+        return event.to_dict(self.options)
+
+    def decode(self, msg: Message) -> KafkaEvent:
+        """
+        Translate from Confluent library's `Message` instance to `Event` instance.
+        """
+        return KafkaEvent(
+            ts=msg.timestamp(),
+            topic=msg.topic(),
+            partition=msg.partition(),
+            offset=msg.offset(),
+            key=msg.key(),
+            value=msg.value(),
+        )
+
+    def deserialize(self, event: KafkaEvent) -> KafkaEvent:
+        """
+        Deserialize event key and value according to decoding options.
+        """
+        if self.options.key_type is not None:
+            if self.options.key_type == "json":
+                event.key = orjson.loads(event.key)
+            else:
+                raise NotImplementedError(
+                    f"Unknown key type: {self.options.value_type}"
+                )
+        if self.options.value_type is not None:
+            if self.options.value_type == "json":
+                event.value = orjson.loads(event.value)
+            else:
+                raise NotImplementedError(
+                    f"Unknown value type: {self.options.value_type}"
+                )
+        return event
 
 
 class OffsetTracker(dict):  # type: ignore
diff --git a/ingestr/src/kafka/model.py b/ingestr/src/kafka/model.py
@@ -0,0 +1,138 @@
+from typing import List, Optional, Tuple, Union
+
+import toolz  # type: ignore[import-untyped]
+from attrs import define
+from dlt.common.time import ensure_pendulum_datetime
+from dlt.common.utils import digest128
+
+
+@define
+class KafkaDecodingOptions:
+    """
+    Options that control decoding of the Kafka event.
+    """
+
+    key_type: Optional[str] = None
+    value_type: Optional[str] = None
+    format: Optional[str] = None
+    include: Optional[str] = None
+    select: Optional[str] = None
+
+    def __attrs_post_init__(self):
+        if self.format is None:
+            self.format = "standard_v1"
+
+    @classmethod
+    def from_params(
+        cls,
+        key_type: List[str],
+        value_type: List[str],
+        format: List[str],
+        include: List[str],
+        select: List[str],
+    ):
+        output_format = None
+        include_fields = None
+        select_field = None
+        if format:
+            output_format = format[0]
+        else:
+            output_format = "standard_v1"
+        if include:
+            output_format = "flexible"
+            include_fields = toolz.apply(str.strip, include[0].split(","))
+        if select:
+            output_format = "flexible"
+            select_field = select[0]
+        return cls(
+            key_type=key_type and key_type[0] or None,
+            value_type=value_type and value_type[0] or None,
+            format=output_format,
+            include=include_fields,
+            select=select_field,
+        )
+
+
+@define
+class KafkaEvent:
+    """
+    Manage details of a typical Kafka event/message/record.
+
+    https://kafka.apache.org/intro#intro_concepts_and_terms
+    """
+
+    ts: Tuple[int, int]
+    topic: str
+    partition: int
+    offset: int
+    key: Union[bytes, str]
+    value: Union[bytes, str]
+
+    def decode_text(self):
+        if self.key is not None and isinstance(self.key, bytes):
+            self.key = self.key.decode("utf-8")
+        if self.value is not None and isinstance(self.value, bytes):
+            self.value = self.value.decode("utf-8")
+
+    def to_dict(self, options: KafkaDecodingOptions):
+        # TODO: Make decoding from text optional.
+        self.decode_text()
+
+        message_id = digest128(self.topic + str(self.partition) + str(self.key))
+
+        # The standard message layout as defined per dlt and ingestr.
+        standard_payload = {
+            "partition": self.partition,
+            "topic": self.topic,
+            "key": self.key,
+            "offset": self.offset,
+            "ts": {
+                "type": self.ts[0],
+                "value": ensure_pendulum_datetime(self.ts[1] / 1e3),
+            },
+            "data": self.value,
+        }
+
+        # Basic Kafka message processors providing two formats.
+        # Returns the message value and metadata.
+        # The legacy format `standard_v1` uses the field `_kafka_msg_id`,
+        # while the future `standard_v2` format uses `_kafka__msg_id`,
+        # better aligned with all the other fields.
+        #
+        # Currently, as of July 2025, `standard_v1` is used as the
+        # default to not cause any breaking changes.
+
+        if options.format == "standard_v1":
+            UserWarning(
+                "Future versions of ingestr will use the `standard_v2` output format. "
+                "To retain compatibility, make sure to start using `format=standard_v1` early."
+            )
+            return {
+                "_kafka": standard_payload,
+                "_kafka_msg_id": message_id,
+            }
+
+        if options.format == "standard_v2":
+            standard_payload["msg_id"] = message_id
+            return {
+                "_kafka": standard_payload,
+            }
+
+        # Slightly advanced Kafka message processor providing basic means of projections.
+        # include: A list of column names to include.
+        # select: A single column name to select and drill down into.
+        if options.format == "flexible":
+            # TODO: Refactor by caching preparation steps.
+            if options.include:
+                include_keys = [
+                    key == "value" and "data" or key for key in options.include
+                ]
+                return toolz.keyfilter(lambda k: k in include_keys, standard_payload)
+            if options.select:
+                # TODO: Instead of a simple dictionary getter, `jsonpointer` or `jqlang`
+                #       can provide easy access to deeper levels of nested data structures.
+                key = options.select.replace("value", "data")
+                return standard_payload.get(key)
+            return standard_payload
+
+        raise NotImplementedError(f"Unknown message processor format: {options.format}")
diff --git a/ingestr/src/sources.py b/ingestr/src/sources.py
@@ -1180,6 +1180,13 @@ def dlt_source(self, uri: str, table: str, **kwargs):
         batch_size = source_params.get("batch_size", [3000])
         batch_timeout = source_params.get("batch_timeout", [3])
 
+        # Decoding options.
+        key_type = source_params.get("key_type", [])
+        value_type = source_params.get("value_type", [])
+        format = source_params.get("format", [])
+        include = source_params.get("include", [])
+        select = source_params.get("select", [])
+
         if not bootstrap_servers:
             raise ValueError(
                 "bootstrap_servers in the URI is required to connect to kafka"
@@ -1189,8 +1196,21 @@ def dlt_source(self, uri: str, table: str, **kwargs):
             raise ValueError("group_id in the URI is required to connect to kafka")
 
         start_date = kwargs.get("interval_start")
+
         from ingestr.src.kafka import kafka_consumer
-        from ingestr.src.kafka.helpers import KafkaCredentials
+        from ingestr.src.kafka.helpers import (
+            KafkaCredentials,
+            KafkaEventProcessor,
+        )
+        from ingestr.src.kafka.model import KafkaDecodingOptions
+
+        options = KafkaDecodingOptions.from_params(
+            key_type=key_type,
+            value_type=value_type,
+            format=format,
+            include=include,
+            select=select,
+        )
 
         return kafka_consumer(
             topics=[table],
@@ -1206,6 +1226,7 @@ def dlt_source(self, uri: str, table: str, **kwargs):
                 sasl_username=sasl_username[0] if len(sasl_username) > 0 else None,  # type: ignore
                 sasl_password=sasl_password[0] if len(sasl_password) > 0 else None,  # type: ignore
             ),
+            msg_processor=KafkaEventProcessor(options=options).process,
             start_from=start_date,
             batch_size=int(batch_size[0]),
             batch_timeout=int(batch_timeout[0]),
diff --git a/requirements.in b/requirements.in
@@ -56,3 +56,4 @@ google-cloud-spanner==3.54.0
 smartsheet-python-sdk==3.0.5
 paramiko==3.5.1
 python-quickbooks==0.9.2
+toolz==1.0.0
diff --git a/requirements.txt b/requirements.txt
@@ -579,6 +579,8 @@ tomlkit==0.13.2
     # via
     #   dlt
     #   snowflake-connector-python
+toolz==1.0.0
+    # via -r requirements.in
 tqdm==4.67.1
     # via -r requirements.in
 typer==0.13.1
diff --git a/requirements_arm64.txt b/requirements_arm64.txt
@@ -574,6 +574,8 @@ tomlkit==0.13.2
     # via
     #   dlt
     #   snowflake-connector-python
+toolz==1.0.0
+    # via -r requirements.in
 tqdm==4.67.1
     # via -r requirements.in
 typer==0.13.1
