ref(metric alerts): Use Kafka Schema for Message Validation (#46296)

Use the new kafka schema and python types for message validation,
falling back to the old way and logging if we do so for safety. Payload
v2 isn't being sent by anything anymore, so this also removes code that
supported it.

requirements-base.txt

@@ -56,7 +56,7 @@ rfc3339-validator>=0.1.2
 rfc3986-validator>=0.1.1
 # [end] jsonschema format validators
 sentry-arroyo[json]>=2.7.1
-sentry-kafka-schemas>=0.0.12
+sentry-kafka-schemas>=0.0.19
 sentry-relay>=0.8.19
 sentry-sdk>=1.17.0
 snuba-sdk>=1.0.5

requirements-dev-frozen.txt

@@ -164,7 +164,7 @@ s3transfer==0.5.2
 selenium==4.3.0
 sentry-arroyo==2.7.1
 sentry-cli==2.16.0
-sentry-kafka-schemas==0.0.12
+sentry-kafka-schemas==0.0.19
 sentry-relay==0.8.19
 sentry-sdk==1.17.0
 simplejson==3.17.6

requirements-frozen.txt

@@ -116,7 +116,7 @@ rsa==4.8
 s3transfer==0.5.2
 selenium==4.3.0
 sentry-arroyo==2.7.1
-sentry-kafka-schemas==0.0.12
+sentry-kafka-schemas==0.0.19
 sentry-relay==0.8.19
 sentry-sdk==1.17.0
 simplejson==3.17.6

src/sentry/incidents/tasks.py

@@ -116,7 +116,7 @@ def handle_subscription_metrics_logger(subscription_update, subscription):
     """
     Logs results from a `QuerySubscription`.
     :param subscription_update: dict formatted according to schemas in
-    sentry.snuba.json_schemas.SUBSCRIPTION_PAYLOAD_VERSIONS
+    sentry_kafka_schemas.schema_types.events_subscription_results_v1.SubscriptionResults
     :param subscription: The `QuerySubscription` that this update is for
     """
     from sentry.incidents.subscription_processor import SubscriptionProcessor
@@ -147,7 +147,7 @@ def handle_snuba_query_update(subscription_update, subscription):
     """
     Handles a subscription update for a `QuerySubscription`.
     :param subscription_update: dict formatted according to schemas in
-    sentry.snuba.json_schemas.SUBSCRIPTION_PAYLOAD_VERSIONS
+    sentry_kafka_schemas.schema_types.events_subscription_results_v1.SubscriptionResults
     :param subscription: The `QuerySubscription` that this update is for
     """
     from sentry.incidents.subscription_processor import SubscriptionProcessor

src/sentry/snuba/json_schemas.py

@@ -9,31 +9,6 @@ SUBSCRIPTION_WRAPPER_SCHEMA = {
 
 
 SUBSCRIPTION_PAYLOAD_VERSIONS: Mapping[int, Mapping[str, Any]] = {
-    2: {
-        "type": "object",
-        "properties": {
-            "subscription_id": {"type": "string", "minLength": 1},
-            "request": {"type": "object"},
-            "result": {
-                "type": "object",
-                "properties": {
-                    "data": {
-                        "type": "array",
-                        "minItems": 1,
-                        "items": {
-                            "type": "object",
-                            "minProperties": 1,
-                            "additionalProperties": {"type": ["number", "null"]},
-                        },
-                    }
-                },
-                "required": ["data"],
-            },
-            "timestamp": {"type": "string", "format": "date-time"},
-        },
-        "required": ["subscription_id", "request", "result", "timestamp"],
-        "additionalProperties": False,
-    },
     3: {
         "type": "object",
         "properties": {

src/sentry/snuba/query_subscription_consumer.py

@@ -1,7 +1,6 @@
 import logging
-import re
 from random import random
-from typing import Any, Callable, Dict, Mapping
+from typing import Callable, Dict, Mapping, cast
 
 import jsonschema
 import pytz
@@ -17,9 +16,16 @@ from arroyo.processing.strategies import (
     ProcessingStrategyFactory,
     RunTask,
 )
+from arroyo.processing.strategies.decoder.base import ValidationError
+from arroyo.processing.strategies.decoder.json import JsonCodec
 from arroyo.types import BrokerValue, Commit, Message, Partition
 from dateutil.parser import parse as parse_date
 from django.conf import settings
+from sentry_kafka_schemas import get_schema
+from sentry_kafka_schemas.schema_types.events_subscription_results_v1 import (
+    PayloadV3,
+    SubscriptionResult,
+)
 
 from sentry import options
 from sentry.snuba.dataset import Dataset, EntityKey
@@ -31,7 +37,7 @@ from sentry.utils.arroyo import MetricsWrapper
 
 logger = logging.getLogger(__name__)
 
-TQuerySubscriptionCallable = Callable[[Dict[str, Any], QuerySubscription], None]
+TQuerySubscriptionCallable = Callable[[PayloadV3, QuerySubscription], None]
 
 subscriber_registry: Dict[str, TQuerySubscriptionCallable] = {}
 
@@ -45,6 +51,14 @@ topic_to_dataset: Dict[str, Dataset] = {
     settings.KAFKA_METRICS_SUBSCRIPTIONS_RESULTS: Dataset.Metrics,
 }
 
+dataset_to_logical_topic = {
+    Dataset.Events: "events-subscription-results",
+    Dataset.Transactions: "transactions-subscription-results",
+    Dataset.PerformanceMetrics: "generic-metrics-subscription-results",
+    Dataset.Sessions: "sessions-subscription-results",
+    Dataset.Metrics: "metrics-subscription-results",
+}
+
 
 def register_subscriber(
     subscriber_key: str,
@@ -58,51 +72,60 @@ def register_subscriber(
     return inner
 
 
-def parse_message_value(value: str) -> Dict[str, Any]:
+def parse_message_value(value: bytes, jsoncodec: JsonCodec) -> PayloadV3:
     """
     Parses the value received via the Kafka consumer and verifies that it
     matches the expected schema.
     :param value: A json formatted string
     :return: A dict with the parsed message
     """
-    with metrics.timer("snuba_query_subscriber.parse_message_value.json_parse"):
-        wrapper: Dict[str, Any] = json.loads(value)
+    old_version = False
 
     with metrics.timer("snuba_query_subscriber.parse_message_value.json_validate_wrapper"):
         try:
-            jsonschema.validate(wrapper, SUBSCRIPTION_WRAPPER_SCHEMA)
-        except jsonschema.ValidationError:
-            metrics.incr("snuba_query_subscriber.message_wrapper_invalid")
-            raise InvalidSchemaError("Message wrapper does not match schema")
-
-    schema_version: int = wrapper["version"]
-    if schema_version not in SUBSCRIPTION_PAYLOAD_VERSIONS:
-        metrics.incr("snuba_query_subscriber.message_wrapper_invalid_version")
-        raise InvalidMessageError("Version specified in wrapper has no schema")
-
-    payload: Dict[str, Any] = wrapper["payload"]
-    with metrics.timer("snuba_query_subscriber.parse_message_value.json_validate_payload"):
-        try:
-            jsonschema.validate(payload, SUBSCRIPTION_PAYLOAD_VERSIONS[schema_version])
-        except jsonschema.ValidationError:
-            metrics.incr("snuba_query_subscriber.message_payload_invalid")
-            raise InvalidSchemaError("Message payload does not match schema")
+            wrapper = cast(SubscriptionResult, jsoncodec.decode(value, validate=True))
+        except ValidationError:
+            old_version = True
+            metrics.incr("snuba_query_subscriber.message_wrapper.old_validation")
+            try:
+                wrapper = json.loads(value)
+                jsonschema.validate(wrapper, SUBSCRIPTION_WRAPPER_SCHEMA)
+            except jsonschema.ValidationError:
+                metrics.incr("snuba_query_subscriber.message_wrapper_invalid")
+                raise InvalidSchemaError("Message wrapper does not match schema")
+
+    payload: PayloadV3 = wrapper["payload"]
+
+    if old_version:
+        schema_version: int = wrapper["version"]
+        if schema_version not in SUBSCRIPTION_PAYLOAD_VERSIONS:
+            metrics.incr("snuba_query_subscriber.message_wrapper_invalid_version")
+            raise InvalidMessageError("Version specified in wrapper has no schema")
+
+        with metrics.timer("snuba_query_subscriber.parse_message_value.json_validate_payload"):
+            try:
+                jsonschema.validate(payload, SUBSCRIPTION_PAYLOAD_VERSIONS[schema_version])
+            except jsonschema.ValidationError:
+                metrics.incr("snuba_query_subscriber.message_payload_invalid")
+                raise InvalidSchemaError("Message payload does not match schema")
+
     # XXX: Since we just return the raw dict here, when the payload changes it'll
     # break things. This should convert the payload into a class rather than passing
     # the dict around, but until we get time to refactor we can keep things working
     # here.
-    payload.setdefault("values", payload.get("result"))
+    payload.setdefault("values", payload.get("result"))  # type: ignore
 
-    payload["timestamp"] = parse_date(payload["timestamp"]).replace(tzinfo=pytz.utc)
+    payload["timestamp"] = parse_date(payload["timestamp"]).replace(tzinfo=pytz.utc)  # type: ignore
     return payload
 
 
 def handle_message(
-    message_value: str,
+    message_value: bytes,
     message_offset: int,
     message_partition: int,
     topic: str,
     dataset: str,
+    jsoncodec: JsonCodec,
 ) -> None:
     """
     Parses the value from Kafka, and if valid passes the payload to the callback defined by the
@@ -116,7 +139,7 @@ def handle_message(
             with metrics.timer(
                 "snuba_query_subscriber.parse_message_value", tags={"dataset": dataset}
             ):
-                contents = parse_message_value(message_value)
+                contents = parse_message_value(message_value, jsoncodec)
         except InvalidMessageError:
             # If the message is in an invalid format, just log the error
             # and continue
@@ -154,23 +177,11 @@ def handle_message(
                 },
             )
             try:
-                if "entity" in contents:
-                    entity_key = contents["entity"]
-                else:
-                    # XXX(ahmed): Remove this logic. This was kept here as backwards compat
-                    # for subscription updates with schema version `2`. However schema version 3
-                    # sends the "entity" in the payload
-                    metrics.incr("query_subscription_consumer.message_value.v2")
-                    entity_regex = r"^(MATCH|match)[ ]*\(([^)]+)\)"
-                    entity_match = re.match(entity_regex, contents["request"]["query"])
-                    if not entity_match:
-                        raise InvalidMessageError("Unable to fetch entity from query in message")
-                    entity_key = entity_match.group(2)
                 if topic in topic_to_dataset:
                     _delete_from_snuba(
                         topic_to_dataset[topic],
                         contents["subscription_id"],
-                        EntityKey(entity_key),
+                        EntityKey(contents["entity"]),
                     )
                 else:
                     logger.error(
@@ -232,6 +243,8 @@ class QuerySubscriptionStrategyFactory(ProcessingStrategyFactory[KafkaPayload]):
     def __init__(self, topic: str):
         self.topic = topic
         self.dataset = topic_to_dataset[self.topic]
+        self.logical_topic = dataset_to_logical_topic[self.dataset]
+        self.jsoncodec = JsonCodec(get_schema(self.logical_topic)["schema"])
 
     def create_with_partitions(
         self,
@@ -258,6 +271,7 @@ class QuerySubscriptionStrategyFactory(ProcessingStrategyFactory[KafkaPayload]):
                         partition,
                         self.topic,
                         self.dataset.value,
+                        self.jsoncodec,
                     )
                 except Exception:
                     # This is a failsafe to make sure that no individual message will block this

tests/sentry/snuba/test_query_subscription_consumer.py

@@ -7,9 +7,11 @@ from unittest import mock
 import pytest
 import pytz
 from arroyo.backends.kafka import KafkaPayload
+from arroyo.processing.strategies.decoder.json import JsonCodec
 from arroyo.types import BrokerValue, Message, Partition, Topic
 from dateutil.parser import parse as parse_date
 from django.conf import settings
+from sentry_kafka_schemas import get_schema
 
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.models import SnubaQuery
@@ -31,6 +33,10 @@ class BaseQuerySubscriptionTest:
     def topic(self):
         return settings.KAFKA_METRICS_SUBSCRIPTIONS_RESULTS
 
+    @cached_property
+    def jsoncodec(self):
+        return JsonCodec(get_schema(self.topic)["schema"])
+
     @cached_property
     def valid_wrapper(self):
         return {"version": 3, "payload": self.valid_payload}
@@ -39,7 +45,10 @@ class BaseQuerySubscriptionTest:
     def valid_payload(self):
         return {
             "subscription_id": "1234",
-            "result": {"data": [{"hello": 50}]},
+            "result": {
+                "data": [{"hello": 50}],
+                "meta": [{"name": "count", "type": "UInt64"}],
+            },
             "request": {
                 "some": "data",
                 "query": """MATCH (metrics_counters) SELECT sum(value) AS value BY
@@ -50,15 +59,6 @@ class BaseQuerySubscriptionTest:
             "timestamp": "2020-01-01T01:23:45.1234",
         }
 
-    @cached_property
-    def old_payload(self):
-        return {
-            "subscription_id": "1234",
-            "result": {"data": [{"hello": 50}]},
-            "request": {"some": "data"},
-            "timestamp": "2020-01-01T01:23:45.1234",
-        }
-
     def build_mock_message(self, data, topic=None):
         message = mock.Mock()
         message.value.return_value = json.dumps(data)
@@ -120,7 +120,7 @@ class HandleMessageTest(BaseQuerySubscriptionTest, TestCase):
 
 class ParseMessageValueTest(BaseQuerySubscriptionTest, unittest.TestCase):
     def run_test(self, message):
-        parse_message_value(json.dumps(message))
+        parse_message_value(json.dumps(message), self.jsoncodec)
 
     def run_invalid_schema_test(self, message):
         with pytest.raises(InvalidSchemaError):
@@ -148,7 +148,7 @@ class ParseMessageValueTest(BaseQuerySubscriptionTest, unittest.TestCase):
 
     def test_invalid_version(self):
         with pytest.raises(InvalidMessageError) as excinfo:
-            self.run_test({"version": 50, "payload": {}})
+            self.run_test({"version": 50, "payload": self.valid_payload})
         assert str(excinfo.value) == "Version specified in wrapper has no schema"
 
     def test_valid(self):
@@ -159,9 +159,6 @@ class ParseMessageValueTest(BaseQuerySubscriptionTest, unittest.TestCase):
         payload["result"]["data"][0]["hello"] = float("nan")
         self.run_test({"version": 3, "payload": payload})
 
-    def test_old_version(self):
-        self.run_test({"version": 2, "payload": self.old_payload})
-
     def test_invalid_wrapper(self):
         self.run_invalid_schema_test({})
         self.run_invalid_schema_test({"version": 1})