[DOP-25464] Make consumer message parsing more robust

Author: dolfinus

.env.local

@@ -8,10 +8,7 @@ export DATA_RENTGEN__KAFKA__SECURITY__USER=data_rentgen
 export DATA_RENTGEN__KAFKA__SECURITY__PASSWORD=changeme
 export DATA_RENTGEN__KAFKA__COMPRESSION=zstd
 #export DATA_RENTGEN__CONSUMER__MAX_RECORDS=100
-# Handling events with a lot of column lineage takes so much time
-# that Kafka coodrinator consider worker as dead. Limit by total message size.
-# This value depends on the number of OL version, complexity of Spark jobs, number of Kafka partitions and number of workers.
-export DATA_RENTGEN__CONSUMER__MAX_PARTITION_FETCH_BYTES=200Kb
+#export DATA_RENTGEN__CONSUMER__MAX_PARTITION_FETCH_BYTES=5MB
 
 export DATA_RENTGEN__SERVER__DEBUG=True
 

data_rentgen/__init__.py

@@ -3,7 +3,7 @@
 
 # _raw_version could contain pre-release version, like 0.0.1dev123
 # value is updated automatically by `poetry version ...` and poetry-bumpversion plugin
-_raw_version = "0.2.2"
+_raw_version = "0.3.0"
 
 # version always contain only release number like 0.0.1
 __version__ = ".".join(_raw_version.split(".")[:3])

data_rentgen/consumer/__init__.py

@@ -11,6 +11,7 @@
 from faststream._compat import ExceptionGroup
 from faststream.asgi import AsgiFastStream, AsgiResponse, get
 from faststream.kafka import KafkaBroker
+from faststream.kafka.publisher.asyncapi import AsyncAPIDefaultPublisher
 from sqlalchemy.ext.asyncio import AsyncSession
 
 import data_rentgen
@@ -38,15 +39,21 @@ def broker_factory(settings: ConsumerApplicationSettings) -> KafkaBroker:
     )
 
     # register subscribers using settings
-    consumer_settings = settings.consumer.model_dump(exclude={"topics_list", "topics_pattern"})
-    broker.subscriber(
+    consumer_settings = settings.consumer.model_dump(exclude={"topics_list", "topics_pattern", "malformed_topic"})
+
+    subscriber = broker.subscriber(
         *settings.consumer.topics_list,
         pattern=settings.consumer.topics_pattern,
         **consumer_settings,
         batch=True,
-    )(runs_events_subscriber)
+    )
+    publisher = broker.publisher(settings.producer.malformed_topic)
+
+    # perform registration
+    subscriber(runs_events_subscriber)
 
     dependency_provider.override(AsyncSession, create_session_factory(settings.database))
+    dependency_provider.override(AsyncAPIDefaultPublisher, lambda: publisher)
     return broker
 
 

data_rentgen/consumer/extractors/__init__.py

@@ -1,7 +1,8 @@
 # SPDX-FileCopyrightText: 2024-2025 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
 
-from data_rentgen.consumer.extractors.batch import BatchExtractionResult, extract_batch
+from data_rentgen.consumer.extractors.batch_extraction_result import BatchExtractionResult
+from data_rentgen.consumer.extractors.batch_extractor import BatchExtractor
 from data_rentgen.consumer.extractors.column_lineage import extract_column_lineage
 from data_rentgen.consumer.extractors.dataset import (
     connect_dataset_with_symlinks,
@@ -17,8 +18,8 @@
 
 __all__ = [
     "BatchExtractionResult",
+    "BatchExtractor",
     "connect_dataset_with_symlinks",
-    "extract_batch",
     "extract_column_lineage",
     "extract_dataset",
     "extract_dataset_and_symlinks",

data_rentgen/consumer/extractors/batch.py …er/extractors/batch_extraction_result.pydata_rentgen/consumer/extractors/batch.py renamed to data_rentgen/consumer/extractors/batch_extraction_result.py data_rentgen/consumer/extractors/batch.py renamed to data_rentgen/consumer/extractors/batch_extraction_result.py

@@ -4,13 +4,6 @@
 
 from typing import TypeVar
 
-from data_rentgen.consumer.extractors.column_lineage import extract_column_lineage
-from data_rentgen.consumer.extractors.input import extract_input
-from data_rentgen.consumer.extractors.operation import extract_operation
-from data_rentgen.consumer.extractors.output import extract_output
-from data_rentgen.consumer.extractors.run import extract_run
-from data_rentgen.consumer.openlineage.job_facets.job_type import OpenLineageJobType
-from data_rentgen.consumer.openlineage.run_event import OpenLineageRunEvent
 from data_rentgen.dto import (
     ColumnLineageDTO,
     DatasetDTO,
@@ -252,44 +245,3 @@ def schemas(self) -> list[SchemaDTO]:
 
     def users(self) -> list[UserDTO]:
         return list(map(self.get_user, self._users))
-
-
-def extract_batch(events: list[OpenLineageRunEvent]) -> BatchExtractionResult:
-    result = BatchExtractionResult()
-    dataset_cache: dict[tuple[str, str], DatasetDTO] = {}
-
-    for event in events:
-        if event.job.facets.jobType and event.job.facets.jobType.jobType == OpenLineageJobType.JOB:
-            operation = extract_operation(event)
-            result.add_operation(operation)
-
-            for input_dataset in event.inputs:
-                input_dto, symlink_dtos = extract_input(operation, input_dataset)
-
-                result.add_input(input_dto)
-                dataset_dto_cache_key = (input_dataset.namespace, input_dataset.name)
-                dataset_cache[dataset_dto_cache_key] = result.get_dataset(input_dto.dataset.unique_key)
-
-                for symlink_dto in symlink_dtos:
-                    result.add_dataset_symlink(symlink_dto)
-
-            for output_dataset in event.outputs:
-                output_dto, symlink_dtos = extract_output(operation, output_dataset)
-
-                result.add_output(output_dto)
-                dataset_dto_cache_key = (output_dataset.namespace, output_dataset.name)
-                dataset_cache[dataset_dto_cache_key] = result.get_dataset(output_dto.dataset.unique_key)
-
-                for symlink_dto in symlink_dtos:
-                    result.add_dataset_symlink(symlink_dto)
-
-            for dataset in event.inputs + event.outputs:
-                column_lineage = extract_column_lineage(operation, dataset, dataset_cache)
-                for item in column_lineage:
-                    result.add_column_lineage(item)
-
-        else:
-            run = extract_run(event)
-            result.add_run(run)
-
-    return result

data_rentgen/consumer/extractors/batch_extractor.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from data_rentgen.consumer.extractors.batch_extraction_result import BatchExtractionResult
+from data_rentgen.consumer.extractors.column_lineage import extract_column_lineage
+from data_rentgen.consumer.extractors.input import extract_input
+from data_rentgen.consumer.extractors.operation import extract_operation
+from data_rentgen.consumer.extractors.output import extract_output
+from data_rentgen.consumer.extractors.run import extract_run
+from data_rentgen.consumer.openlineage.job_facets.job_type import OpenLineageJobType
+from data_rentgen.consumer.openlineage.run_event import OpenLineageRunEvent
+from data_rentgen.dto import (
+    DatasetDTO,
+)
+
+
+class BatchExtractor:
+    def __init__(self) -> None:
+        self.dataset_cache: dict[tuple[str, str], DatasetDTO] = {}
+        self.result = BatchExtractionResult()
+
+    def add_events(self, events: list[OpenLineageRunEvent]) -> BatchExtractionResult:
+        for event in events:
+            if event.job.facets.jobType and event.job.facets.jobType.jobType == OpenLineageJobType.JOB:
+                self.extract_operation(event)
+            else:
+                self.extract_run(event)
+
+        return self.result
+
+    def extract_run(self, event: OpenLineageRunEvent) -> None:
+        run = extract_run(event)
+        self.result.add_run(run)
+
+    def extract_operation(self, event: OpenLineageRunEvent) -> None:
+        operation = extract_operation(event)
+        self.result.add_operation(operation)
+
+        for input_dataset in event.inputs:
+            input_dto, symlink_dtos = extract_input(operation, input_dataset)
+
+            self.result.add_input(input_dto)
+            dataset_dto_cache_key = (input_dataset.namespace, input_dataset.name)
+            self.dataset_cache[dataset_dto_cache_key] = self.result.get_dataset(input_dto.dataset.unique_key)
+
+            for symlink_dto in symlink_dtos:
+                self.result.add_dataset_symlink(symlink_dto)
+
+        for output_dataset in event.outputs:
+            output_dto, symlink_dtos = extract_output(operation, output_dataset)
+
+            self.result.add_output(output_dto)
+            dataset_dto_cache_key = (output_dataset.namespace, output_dataset.name)
+            self.dataset_cache[dataset_dto_cache_key] = self.result.get_dataset(output_dto.dataset.unique_key)
+
+            for symlink_dto in symlink_dtos:
+                self.result.add_dataset_symlink(symlink_dto)
+
+        for dataset in event.inputs + event.outputs:
+            column_lineage = extract_column_lineage(operation, dataset, self.dataset_cache)
+            for item in column_lineage:
+                self.result.add_column_lineage(item)

data_rentgen/consumer/settings/__init__.py

@@ -6,6 +6,7 @@
 
 from data_rentgen.consumer.settings.consumer import ConsumerSettings
 from data_rentgen.consumer.settings.kafka import KafkaSettings
+from data_rentgen.consumer.settings.producer import ProducerSettings
 from data_rentgen.db.settings import DatabaseSettings
 from data_rentgen.logging.settings import LoggingSettings
 
@@ -50,5 +51,9 @@ class ConsumerApplicationSettings(BaseSettings):
         default_factory=ConsumerSettings,
         description=":ref:`Consumer settings <configuration-consumer-specific>`",
     )
+    producer: ProducerSettings = Field(
+        default_factory=ProducerSettings,
+        description=":ref:`Producer settings <configuration-producer-specific>`",
+    )
 
     model_config = SettingsConfigDict(env_prefix="DATA_RENTGEN__", env_nested_delimiter="__", extra="forbid")

data_rentgen/consumer/settings/consumer.py

@@ -19,6 +19,7 @@ class ConsumerSettings(BaseModel):
     .. code-block:: bash
 
         DATA_RENTGEN__CONSUMER__TOPICS_LIST=["input.runs"]
+        DATA_RENTGEN__CONSUMER__MALFOMED_TOPIC="input.runs:malformed"
         DATA_RENTGEN__CONSUMER__GROUP_ID=data-rentgen
         DATA_RENTGEN__CONSUMER__FETCH_MAX_WAIT_MS=5000
         DATA_RENTGEN__CONSUMER__MAX_PARTITION_FETCH_BYTES=5MiB

data_rentgen/consumer/settings/producer.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import BaseModel, Field
+
+
+class ProducerSettings(BaseModel):
+    """Data.Rentgen producer-specific settings.
+
+    These options are passed directly to
+    `AIOKafkaProducer <https://aiokafka.readthedocs.io/en/stable/api.html#aiokafka.AIOKafkaProducer>`_.
+
+    Examples
+    --------
+
+    .. code-block:: bash
+
+        DATA_RENTGEN__PRODUCER__MALFOMED_TOPIC="input.runs__malformed"
+    """
+
+    malformed_topic: str = Field(
+        default="input.runs__malformed",
+        description="Topic to publish malformed messages to.",
+    )

data_rentgen/consumer/subscribers.py

@@ -3,10 +3,17 @@
 
 from __future__ import annotations
 
-from faststream import Depends, Logger
+import asyncio
+from typing import cast
+
+from aiokafka import ConsumerRecord
+from faststream import Depends, Logger, NoCast
+from faststream.kafka import KafkaMessage
+from faststream.kafka.publisher.asyncapi import AsyncAPIDefaultPublisher
+from pydantic import TypeAdapter
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from data_rentgen.consumer.extractors import BatchExtractionResult, extract_batch
+from data_rentgen.consumer.extractors import BatchExtractionResult, BatchExtractor
 from data_rentgen.consumer.openlineage.run_event import OpenLineageRunEvent
 from data_rentgen.dependencies import Stub
 from data_rentgen.services.uow import UnitOfWork
@@ -15,33 +22,75 @@
     "runs_events_subscriber",
 ]
 
-
-def get_unit_of_work(session: AsyncSession = Depends(Stub(AsyncSession))) -> UnitOfWork:
-    return UnitOfWork(session)
+OpenLineageRunEventAdapter = TypeAdapter(OpenLineageRunEvent)
 
 
 async def runs_events_subscriber(
-    events: list[OpenLineageRunEvent],
+    _events: NoCast[list[OpenLineageRunEvent]],
+    batch: KafkaMessage,
     logger: Logger,
-    unit_of_work: UnitOfWork = Depends(get_unit_of_work),
+    publisher: AsyncAPIDefaultPublisher = Depends(Stub(AsyncAPIDefaultPublisher)),
+    session: AsyncSession = Depends(Stub(AsyncSession)),
 ):
-    logger.info("Got %d events", len(events))
-    extracted = extract_batch(events)
-    logger.info("Extracted: %r", extracted)
+    logger.info("Extracting events")
+    parsed, malformed = await extract_events(batch, logger)
 
     logger.info("Saving to database")
-    await save_to_db(extracted, unit_of_work, logger)
+    await save_to_db(parsed, session, logger)
     logger.info("Saved successfully")
 
+    if malformed:
+        logger.warning("Malformed messages: %d", len(malformed))
+        await report_malformed(batch, malformed, publisher)
+
+
+async def extract_events(
+    raw_data: KafkaMessage,
+    logger: Logger,
+    await_every: int = 50,
+) -> tuple[BatchExtractionResult, list[ConsumerRecord]]:
+    messages = cast(tuple[ConsumerRecord], raw_data.raw_message)  # https://github.com/airtai/faststream/issues/2102
+    total_bytes = sum(len(message.value or "") for message in messages)
+    logger.info("Got %d messages (%dKiB)", len(messages), total_bytes / 1024)
+
+    extractor = BatchExtractor()
+    malformed: list[ConsumerRecord] = []
+
+    for i, message in enumerate(messages):
+        try:
+            if message.value is None:
+                msg = "Message value cannot be empty"
+                raise ValueError(msg)  # noqa: TRY301
+
+            event = OpenLineageRunEventAdapter.validate_json(message.value)
+            extractor.add_events([event])
+        except (ValueError, TypeError):
+            logger.error(  # noqa: TRY400
+                "Failed to parse message: ConsumerRecord(topic=%r, partition=%d, offset=%d)",
+                message.topic,
+                message.partition,
+                message.offset,
+            )
+            malformed.append(message)
+
+        if await_every and i >= await_every and i % await_every == 0:
+            # OpenLineage models are heavy, parsing is CPU bound task which may take some time.
+            # Blocking event loop is not a good idea, so we need to await sometimes,
+            await asyncio.sleep(0)
+
+    return extractor.result, malformed
+
 
 async def save_to_db(
     data: BatchExtractionResult,
-    unit_of_work: UnitOfWork,
+    session: AsyncSession,
     logger: Logger,
 ) -> None:
     # To avoid deadlocks when parallel consumer instances insert/update the same row,
     # commit changes for each row instead of committing the whole batch. Yes, this cloud be slow.
 
+    unit_of_work = UnitOfWork(session)
+
     logger.debug("Creating locations")
     for location_dto in data.locations():
         async with unit_of_work:
@@ -108,3 +157,25 @@ async def save_to_db(
 
         logger.debug("Creating column lineage")
         await unit_of_work.column_lineage.create_bulk(column_lineage)
+
+
+async def report_malformed(
+    batch: KafkaMessage,
+    messages: list[ConsumerRecord],
+    publisher: AsyncAPIDefaultPublisher,
+):
+    # Return malformed messages back to the broker
+    for message in messages:
+        headers: dict[str, str] = {}
+        if message.headers:
+            headers = {key: value.decode("utf-8") for key, value in message.headers}
+
+        await publisher.publish(
+            message.value,
+            key=message.key,
+            partition=message.partition,
+            timestamp_ms=message.timestamp,
+            headers=headers or None,
+            reply_to=batch.message_id,
+            correlation_id=batch.correlation_id,
+        )