Allow skipping validation

sentry_streams/sentry_streams/metrics/metrics.py

@@ -84,6 +84,9 @@ class Metric(Enum):
     # This counts how many errors were encountered in the step in the pipeline.
     # Tags: step, pipeline, error_type
     ERRORS = "errors"
+    # This counts how many schema validation failures occurred while parsing messages.
+    # Tags: step, pipeline
+    PARSER_VALIDATION = "parser.validation"
 
 
 @runtime_checkable

sentry_streams/sentry_streams/metrics/stats.py

@@ -14,6 +14,7 @@ def __init__(self, metrics: Metrics) -> None:
 
         self._exec_buffer: dict[str, int] = defaultdict(int)
         self._error_buffer: dict[str, int] = defaultdict(int)
+        self._validation_buffer: dict[str, int] = defaultdict(int)
         self._timing_buffer: dict[str, float] = defaultdict(float)
 
         self.__last_flush_time = 0.0
@@ -26,6 +27,10 @@ def step_error(self, step: str) -> None:
         self._error_buffer[step] += 1
         self._maybe_flush()
 
+    def parser_validation_failure(self, step: str) -> None:
+        self._validation_buffer[step] += 1
+        self._maybe_flush()
+
     def step_timing(self, step: str, value: float) -> None:
         if self._timing_buffer[step] < value:
             # TODO: turn this into a moving average.
@@ -41,12 +46,16 @@ def _maybe_flush(self) -> None:
             for step, value in self._error_buffer.items():
                 tags = {"step": step}
                 self._metrics.increment(Metric.ERRORS, value, tags)
+            for step, value in self._validation_buffer.items():
+                tags = {"step": step}
+                self._metrics.increment(Metric.PARSER_VALIDATION, value, tags)
             for step, fvalue in self._timing_buffer.items():
                 tags = {"step": step}
                 self._metrics.timing(Metric.DURATION, fvalue, tags)
 
             self._exec_buffer = defaultdict(int)
             self._error_buffer = defaultdict(int)
+            self._validation_buffer = defaultdict(int)
             self._timing_buffer = defaultdict(float)
 
 

sentry_streams/sentry_streams/pipeline/msg_codecs.py

@@ -15,8 +15,9 @@
 import polars as pl
 from polars import Schema as PolarsSchema
 from sentry_kafka_schemas import get_codec
-from sentry_kafka_schemas.codecs import Codec
+from sentry_kafka_schemas.codecs import Codec, ValidationError
 
+from sentry_streams.metrics.stats import get_stats
 from sentry_streams.pipeline.datatypes import (
     DataType,
 )
@@ -42,18 +43,35 @@ def _get_codec_from_msg(msg: Message[Any]) -> Codec[Any]:
     return codec
 
 
-def msg_parser(msg: Message[bytes]) -> Any:
+def msg_parser(msg: Message[bytes], skip_validation: bool = False, step_name: str = "") -> Any:
     codec = _get_codec_from_msg(msg)
-    payload = msg.payload
-    decoded = codec.decode(payload, True)
-
+    decoded = codec.decode(msg.payload, validate=False)
+    try:
+        codec.validate(decoded)
+    except ValidationError:
+        get_stats().parser_validation_failure(step_name)
+        if not skip_validation:
+            raise
     return decoded
 
 
-def batch_msg_parser(msg: Message[Sequence[bytes]]) -> Sequence[Any]:
-    payloads = msg.payload
+def batch_msg_parser(
+    msg: Message[Sequence[bytes]],
+    skip_validation: bool = False,
+    step_name: str = "",
+) -> Sequence[Any]:
     codec = _get_codec_from_msg(msg)
-    return [codec.decode(payload, True) for payload in payloads]
+    decoded: list[Any] = []
+    for payload in msg.payload:
+        decoded_payload = codec.decode(payload, validate=False)
+        try:
+            codec.validate(decoded_payload)
+        except ValidationError:
+            get_stats().parser_validation_failure(step_name)
+            if not skip_validation:
+                raise
+        decoded.append(decoded_payload)
+    return decoded
 
 
 def msg_serializer(msg: Message[Any], dt_format: Optional[str] = None) -> bytes:

sentry_streams/sentry_streams/pipeline/pipeline.py

@@ -724,13 +724,20 @@ class Parser(ComplexStep[bytes, TransformFuncReturnType], Generic[TransformFuncR
     is supported by sentry-kafka-schemas. See examples/ for usage, this step can be plugged in
     flexibly into a pipeline. Keep in mind, data up until this step will simply be bytes.
 
-    Supports both JSON and protobuf.
+    Supports both JSON and protobuf. Set ``skip_validation=True`` to decode without
+    schema validation.
     """
 
+    skip_validation: bool = False
+
     def convert(self) -> Transform[bytes, TransformFuncReturnType]:
         return Map[bytes, TransformFuncReturnType](
             name=self.name,
-            function=msg_parser,
+            function=partial(
+                msg_parser,
+                skip_validation=self.skip_validation,
+                step_name=self.name,
+            ),
         )
 
 
@@ -739,11 +746,21 @@ class BatchParser(
     ComplexStep[Sequence[bytes], Sequence[TransformFuncReturnType]],
     Generic[TransformFuncReturnType],
 ):
+    """
+    Like :class:`Parser`, but for batches of byte payloads. Set ``skip_validation=True``
+    to decode without schema validation.
+    """
+
+    skip_validation: bool = False
 
     def convert(self) -> Transform[Sequence[bytes], Sequence[TransformFuncReturnType]]:
         return Map[Sequence[bytes], Sequence[TransformFuncReturnType]](
             name=self.name,
-            function=batch_msg_parser,
+            function=partial(
+                batch_msg_parser,
+                skip_validation=self.skip_validation,
+                step_name=self.name,
+            ),
         )
 
 

sentry_streams/tests/metrics/test_stats.py

@@ -24,6 +24,7 @@ def test_correct_values_are_flushed(
     stats.step_exec("in_step")
     stats.step_exec("in_step")
     stats.step_error("err_step")
+    stats.parser_validation_failure("parser_step")
     stats.step_timing("timer_step", 0.1)
     _mock_time.return_value = 120.0
     stats.step_timing("timer_step", 0.05)  # max is 0.1
@@ -32,6 +33,7 @@ def test_correct_values_are_flushed(
         [
             call(Metric.INPUT_MESSAGES.value, 2, tags={"step": "in_step"}),
             call(Metric.ERRORS.value, 1, tags={"step": "err_step"}),
+            call(Metric.PARSER_VALIDATION.value, 1, tags={"step": "parser_step"}),
         ],
         any_order=True,
     )

sentry_streams/tests/pipeline/test_metrics.py

@@ -42,6 +42,7 @@ def test_metric_enum_values() -> None:
     assert Metric.OUTPUT_BYTES.value == "output.bytes"
     assert Metric.DURATION.value == "duration"
     assert Metric.ERRORS.value == "errors"
+    assert Metric.PARSER_VALIDATION.value == "parser.validation"
 
 
 def test_dummy_metrics_backend_increment() -> None:

sentry_streams/tests/pipeline/test_msg_codecs.py

@@ -1,15 +1,16 @@
 import json
 from collections.abc import Iterable
 from datetime import datetime
-from importlib import resources
 from io import BytesIO
 from typing import Any, Mapping, Sequence, Union
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import polars as pl
 import pytest
 from polars import Schema as PolarsSchema
 from polars.testing import assert_frame_equal
+from sentry_kafka_schemas.codecs import ValidationError
+from sentry_kafka_schemas.schema_types.ingest_metrics_v1 import IngestMetric
 
 from sentry_streams.pipeline.datatypes import (
     Field,
@@ -26,6 +27,22 @@
     resolve_polars_schema,
     serialize_to_parquet,
 )
+from sentry_streams.pipeline.pipeline import BatchParser, Map, Parser
+
+INGEST_METRIC: IngestMetric = {
+    "org_id": 420,
+    "project_id": 420,
+    "name": "s:sessions/user@none",
+    "tags": {
+        "sdk": "raven-node/2.6.3",
+        "environment": "production",
+        "release": "sentry-test@1.0.0",
+    },
+    "timestamp": 1846062325,
+    "type": "s",
+    "retention_days": 90,
+    "value": [1617781333],
+}
 
 
 def test_msg_serializer_default_isoformat() -> None:
@@ -50,15 +67,7 @@ def test_msg_serializer_custom_dt_format() -> None:
 
 
 def test_batch_msg_parser_nominal_case() -> None:
-    with (
-        resources.files("sentry_kafka_schemas.examples.ingest-metrics.1")
-        .joinpath("base64-set.json")
-        .open("r") as f
-    ):
-        data = json.load(f)
-    expected = [data]
-
-    payload: Sequence[bytes] = [json.dumps(data).encode("utf-8")]
+    payload: Sequence[bytes] = [json.dumps(INGEST_METRIC).encode("utf-8")]
 
     msg = PyMessage(
         payload=payload,
@@ -68,7 +77,56 @@ def test_batch_msg_parser_nominal_case() -> None:
     )
 
     result = batch_msg_parser(msg)
-    assert result == expected
+    assert result == [INGEST_METRIC]
+
+
+def test_parser_validates_by_default() -> None:
+    step = Parser[IngestMetric]("parser").convert()
+    assert isinstance(step, Map)
+    msg = PyMessage(payload=b"{}", schema="ingest-metrics", headers=[], timestamp=0.0)
+    mock_stats = MagicMock()
+    with (
+        patch("sentry_streams.pipeline.msg_codecs.get_stats", return_value=mock_stats),
+        pytest.raises(ValidationError),
+    ):
+        step.resolved_function(msg)
+    mock_stats.parser_validation_failure.assert_called_once_with("parser")
+
+
+def test_parser_skip_validation() -> None:
+    step = Parser[IngestMetric]("parser", skip_validation=True).convert()
+    assert isinstance(step, Map)
+    msg = PyMessage(payload=b"{}", schema="ingest-metrics", headers=[], timestamp=0.0)
+    mock_stats = MagicMock()
+    with patch("sentry_streams.pipeline.msg_codecs.get_stats", return_value=mock_stats):
+        assert step.resolved_function(msg) == {}
+    mock_stats.parser_validation_failure.assert_called_once_with("parser")
+
+
+def test_batch_parser_validates_by_default() -> None:
+    step = BatchParser[IngestMetric]("batch-parser").convert()
+    assert isinstance(step, Map)
+    payloads: Sequence[bytes] = [b"{}"]
+    msg = PyMessage(payload=payloads, schema="ingest-metrics", headers=[], timestamp=0.0)
+    mock_stats = MagicMock()
+    with (
+        patch("sentry_streams.pipeline.msg_codecs.get_stats", return_value=mock_stats),
+        pytest.raises(ValidationError),
+    ):
+        step.resolved_function(msg)
+    mock_stats.parser_validation_failure.assert_called_once_with("batch-parser")
+
+
+def test_batch_parser_skip_validation() -> None:
+    step = BatchParser[IngestMetric]("batch-parser", skip_validation=True).convert()
+    assert isinstance(step, Map)
+    payloads: Sequence[bytes] = [b"{}", b'{"extra": true}']
+    msg = PyMessage(payload=payloads, schema="ingest-metrics", headers=[], timestamp=0.0)
+    mock_stats = MagicMock()
+    with patch("sentry_streams.pipeline.msg_codecs.get_stats", return_value=mock_stats):
+        assert step.resolved_function(msg) == [{}, {"extra": True}]
+    assert mock_stats.parser_validation_failure.call_count == 2
+    mock_stats.parser_validation_failure.assert_called_with("batch-parser")
 
 
 def test_msg_no_schema() -> None: