feat: A small step towards schema enforcement (#4405)

Devserver now enforces the schema for the events topic

snuba/cli/consumer.py

@@ -116,6 +116,13 @@
  type=int,
 )
 @click.option("--join-timeout", type=int, help="Join timeout in seconds.", default=5)
+@click.option(
+ "--enforce-schema",
+ type=bool,
+ is_flag=True,
+ default=False,
+ help="Enforce schema on the raw events topic.",
+)
 @click.option(
  "--profile-path", type=click.Path(dir_okay=True, file_okay=False, exists=True)
 )
@@ -144,6 +151,7 @@ def consumer(
  input_block_size: Optional[int],
  output_block_size: Optional[int],
  join_timeout: int = 5,
+ enforce_schema: bool = False,
  log_level: Optional[str] = None,
  profile_path: Optional[str] = None,
  max_poll_interval_ms: Optional[int] = None,
@@ -201,6 +209,7 @@ def consumer(
  slice_id=slice_id,
  join_timeout=join_timeout,
  max_poll_interval_ms=max_poll_interval_ms,
+ enforce_schema=enforce_schema,
  )
 
  consumer = consumer_builder.build_base_consumer()

snuba/cli/devserver.py

@@ -81,6 +81,7 @@ def devserver(*, bootstrap: bool, workers: bool) -> None:
  "--no-strict-offset-reset",
  "--log-level=debug",
  "--storage=errors",
+ "--enforce-schema",
  ],
  ),
  (

snuba/cli/dlq_consumer.py

@@ -169,6 +169,7 @@ def handler(signum: int, frame: Any) -> None:
  metrics=metrics,
  slice_id=instruction.slice_id,
  join_timeout=None,
+ enforce_schema=False,
  )
 
  consumer = consumer_builder.build_dlq_consumer(instruction)

snuba/consumers/consumer.py

@@ -523,6 +523,7 @@ def process_message(
  processor: MessageProcessor,
  consumer_group: str,
  snuba_logical_topic: SnubaTopic,
+ enforce_schema: bool,
  message: Message[KafkaPayload],
 ) -> Union[None, BytesInsertBatch, ReplacementBatch]:
  local_metrics = MetricsWrapper(
@@ -573,6 +574,8 @@ def process_message(
  _LAST_INVALID_MESSAGE[snuba_logical_topic.name] = start
  sentry_sdk.set_tag("invalid_message_schema", "true")
  logger.warning(err, exc_info=True)
+ if enforce_schema:
+ raise
 
  # TODO: this is not the most efficient place to emit a metric, but
  # as long as should_validate is behind a sample rate it should be
@@ -604,8 +607,6 @@ def process_message(
  value = message.value
  raise InvalidMessage(value.partition, value.offset) from err
 
- return None
-
  if isinstance(result, InsertBatch):
  return BytesInsertBatch(
  [json_row_encoder.encode(row) for row in result.rows],

snuba/consumers/consumer_builder.py

@@ -71,6 +71,7 @@ def __init__(
  metrics: MetricsBackend,
  slice_id: Optional[int],
  join_timeout: Optional[float],
+ enforce_schema: bool,
  profile_path: Optional[str] = None,
  max_poll_interval_ms: Optional[int] = None,
  ) -> None:
@@ -83,6 +84,7 @@ def __init__(
  self.__consumer_config = consumer_config
  self.__kafka_params = kafka_params
  self.consumer_group = kafka_params.group_id
+ self.__enforce_schema = enforce_schema
 
  broker_config = build_kafka_consumer_configuration(
  self.__consumer_config.raw_topic.broker_config,
@@ -213,6 +215,7 @@ def build_streaming_strategy_factory(
  processor,
  self.consumer_group,
  logical_topic,
+ self.__enforce_schema,
  ),
  collector=build_batch_writer(
  table_writer,

snuba/web/views.py

@@ -644,6 +644,7 @@ def commit(
  stream_loader.get_processor(),
  "consumer_grouup",
  stream_loader.get_default_topic_spec().topic,
+ False,
  ),
  build_batch_writer(table_writer, metrics=metrics),
  max_batch_size=1,

tests/consumers/test_consumer_builder.py

@@ -18,7 +18,7 @@
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.utils.metrics.backends.abstract import MetricsBackend
 from snuba.utils.metrics.wrapper import MetricsWrapper
-from tests.fixtures import get_raw_event
+from tests.fixtures import get_raw_error_message
 from tests.test_consumer import get_row_count
 
 test_storage_key = StorageKey("errors")
@@ -61,6 +61,7 @@
  ),
  slice_id=None,
  join_timeout=5,
+ enforce_schema=True,
 )
 
 optional_consumer_config = resolve_consumer_config(
@@ -104,6 +105,7 @@
  ),
  slice_id=None,
  join_timeout=5,
+ enforce_schema=True,
 )
 
 
@@ -160,8 +162,7 @@ def test_run_processing_strategy() -> None:
  strategy_factory = consumer_builder.build_streaming_strategy_factory()
  strategy = strategy_factory.create_with_partitions(commit, partitions)
 
- raw_message = get_raw_event()
- json_string = json.dumps([2, "insert", raw_message, []])
+ json_string = json.dumps(get_raw_error_message())
 
  message = Message(
  BrokerValue(

tests/test_consumer.py

@@ -31,13 +31,16 @@
 from snuba.utils.streams.topics import Topic as SnubaTopic
 from tests.assertions import assert_changes
 from tests.backends.metrics import TestingMetricsBackend, Timing
+from tests.fixtures import get_raw_error_message
 
 
 def test_streaming_consumer_strategy() -> None:
  messages = (
  Message(
  BrokerValue(
- KafkaPayload(None, b"{}", []),
+ KafkaPayload(
+ None, json.dumps(get_raw_error_message()).encode("utf-8"), []
+ ),
  Partition(Topic("events"), 0),
  i,
  datetime.now(),
@@ -72,7 +75,7 @@ def write_step() -> ProcessedMessageBatchWriter:
  factory = KafkaConsumerStrategyFactory(
  None,
  functools.partial(
- process_message, processor, "consumer_group", SnubaTopic.EVENTS
+ process_message, processor, "consumer_group", SnubaTopic.EVENTS, True
  ),
  write_step,
  max_batch_size=10,