diff --git a/rust_snuba/src/config.rs b/rust_snuba/src/config.rs index c40f64e9ef..ff05473a65 100644 --- a/rust_snuba/src/config.rs +++ b/rust_snuba/src/config.rs @@ -105,6 +105,7 @@ pub struct EnvConfig { pub default_retention_days: u16, pub lower_retention_days: u16, pub valid_retention_days: HashSet, + pub record_cogs: bool, } impl Default for EnvConfig { @@ -116,6 +117,7 @@ impl Default for EnvConfig { default_retention_days: 90, lower_retention_days: 30, valid_retention_days: [30, 90].iter().cloned().collect(), + record_cogs: false, } } } diff --git a/rust_snuba/src/factory.rs b/rust_snuba/src/factory.rs index 07f5994613..ebe432af20 100644 --- a/rust_snuba/src/factory.rs +++ b/rust_snuba/src/factory.rs @@ -131,10 +131,10 @@ impl ProcessingStrategyFactory for ConsumerStrategyFactory { let cogs_label = get_cogs_label(&self.storage_config.message_processor.python_class_name); - // Produce cogs if generic metrics and we are not skipping writes + // Produce cogs if generic metrics AND we are not skipping writes AND record_cogs is true let next_step: Box> = - match (self.skip_write, cogs_label) { - (false, Some(resource_id)) => Box::new(RecordCogs::new( + match (self.skip_write, self.env_config.record_cogs, cogs_label) { + (false, true, Some(resource_id)) => Box::new(RecordCogs::new( next_step, resource_id, self.accountant_topic_config.broker_config.clone(), diff --git a/snuba/consumers/consumer_config.py b/snuba/consumers/consumer_config.py index 406edb28d5..3cc083ee5a 100644 --- a/snuba/consumers/consumer_config.py +++ b/snuba/consumers/consumer_config.py @@ -50,6 +50,7 @@ class EnvConfig: default_retention_days: int lower_retention_days: int valid_retention_days: list[int] + record_cogs: bool @dataclass(frozen=True) @@ -124,6 +125,7 @@ def _resolve_env_config() -> EnvConfig: default_retention_days = settings.DEFAULT_RETENTION_DAYS lower_retention_days = settings.LOWER_RETENTION_DAYS valid_retention_days = list(settings.VALID_RETENTION_DAYS) + record_cogs = settings.RECORD_COGS return EnvConfig( sentry_dsn=sentry_dsn, dogstatsd_host=dogstatsd_host, @@ -131,6 +133,7 @@ def _resolve_env_config() -> EnvConfig: default_retention_days=default_retention_days, lower_retention_days=lower_retention_days, valid_retention_days=valid_retention_days, + record_cogs=record_cogs, ) diff --git a/snuba/datasets/processors/generic_metrics_processor.py b/snuba/datasets/processors/generic_metrics_processor.py index 4aa2efb739..eb09fdbb92 100644 --- a/snuba/datasets/processors/generic_metrics_processor.py +++ b/snuba/datasets/processors/generic_metrics_processor.py @@ -3,7 +3,6 @@ import zlib from abc import ABC, abstractmethod, abstractproperty from datetime import datetime, timezone -from random import random from typing import ( Any, Iterable, @@ -18,6 +17,7 @@ from sentry_kafka_schemas.schema_types.snuba_generic_metrics_v1 import GenericMetric from usageaccountant import UsageUnit +from snuba import settings from snuba.cogs.accountant import record_cogs from snuba.consumers.types import KafkaMessageMetadata from snuba.datasets.events_format import EventTooOld, enforce_retention @@ -37,7 +37,6 @@ ) from snuba.datasets.processors import DatasetMessageProcessor from snuba.processor import InsertBatch, ProcessedMessage, _ensure_valid_date -from snuba.state import get_config logger = logging.getLogger(__name__) @@ -167,7 +166,7 @@ def process_message( ) def __record_cogs(self, message: GenericMetric) -> None: - if random() < (get_config("gen_metrics_processor_cogs_probability") or 0): + if settings.RECORD_COGS: record_cogs( resource_id=self._resource_id, app_feature=f"genericmetrics_{message['use_case_id']}", diff --git a/snuba/settings/__init__.py b/snuba/settings/__init__.py index 498a54cf6d..1006523e6e 100644 --- a/snuba/settings/__init__.py +++ b/snuba/settings/__init__.py @@ -188,6 +188,9 @@ class RedisClusters(TypedDict): # Query Recording Options RECORD_QUERIES = False +# Record COGS +RECORD_COGS = False + # Runtime Config Options CONFIG_MEMOIZE_TIMEOUT = 10 CONFIG_STATE: Mapping[str, Optional[Any]] = {} diff --git a/tests/datasets/test_generic_metrics_processor.py b/tests/datasets/test_generic_metrics_processor.py index f98b6c8060..5e76d7a8bc 100644 --- a/tests/datasets/test_generic_metrics_processor.py +++ b/tests/datasets/test_generic_metrics_processor.py @@ -6,7 +6,7 @@ import pytest from usageaccountant import UsageUnit -from snuba import state +from snuba import settings from snuba.datasets.metrics_messages import InputType from snuba.datasets.processors.generic_metrics_processor import ( GenericDistributionsMetricsProcessor, @@ -154,7 +154,8 @@ def test_record_cogs(dis_processor: GenericDistributionsMetricsProcessor) -> Non "retention_days": 22, "mapping_meta": MAPPING_META_COMMON, } - state.set_config("gen_metrics_processor_cogs_probability", 1.0) + settings.RECORD_COGS = True + with mock.patch( "snuba.datasets.processors.generic_metrics_processor.record_cogs" ) as record_cogs: