def __build_tick_consumer(self) -> CommitLogTickConsumer:
consumer_configuration = build_kafka_consumer_configuration(
self.__commit_log_topic_spec.topic,
self.__consumer_group,
auto_offset_reset=self.__auto_offset_reset,
strict_offset_reset=self.__strict_offset_reset,
)
# Collect metrics from librdkafka if we have stats_collection_freq_ms set
# for the consumer group, or use the default.
stats_collection_frequency_ms = get_config(
f"stats_collection_freq_ms_{self.__consumer_group}",
get_config("stats_collection_freq_ms", 0),
)
if stats_collection_frequency_ms and stats_collection_frequency_ms > 0:
def stats_callback(stats_json: str) -> None:
stats = rapidjson.loads(stats_json)
self.__metrics.gauge("librdkafka.total_queue_size",
stats.get("replyq", 0))
consumer_configuration.update({
"statistics.interval.ms": stats_collection_frequency_ms,
"stats_cb": stats_callback,
})
return CommitLogTickConsumer(
KafkaConsumer(consumer_configuration),
followed_consumer_group=self.__followed_consumer_group,
time_shift=(timedelta(
seconds=self.__delay_seconds *
-1) if self.__delay_seconds is not None else None),
)
def selector_func(_query: Query,
referrer: str) -> Tuple[str, List[str]]:
# In case something goes wrong, set this to 1 to revert to the events storage.
kill_rollout = state.get_config("errors_rollout_killswitch", 0)
assert isinstance(kill_rollout, (int, str))
if int(kill_rollout):
return "events", []
if referrer in settings.ERRORS_ROLLOUT_BY_REFERRER:
return "discover", []
if settings.ERRORS_ROLLOUT_ALL:
return "discover", []
default_threshold = state.get_config("discover_query_percentage",
0)
assert isinstance(default_threshold, (float, int, str))
threshold = settings.ERRORS_QUERY_PERCENTAGE_BY_REFERRER.get(
referrer, default_threshold)
if random.random() < float(threshold):
return "events", ["discover"]
return "events", []
def set_configs() -> Generator[None, None, None]:
old_max = state.get_config("max_days")
old_align = state.get_config("date_align_seconds")
state.set_config("max_days", 5)
state.set_config("date_align_seconds", 3600)
yield
state.set_config("max_days", old_max)
state.set_config("date_align_seconds", old_align)
def build_request(
self,
dataset: Dataset,
timestamp: datetime,
offset: Optional[int],
timer: Timer,
metrics: Optional[MetricsBackend] = None,
) -> Request:
try:
if metrics is not None:
metrics.increment("snql.subscription.delegate.incoming")
snql_rollout_pct = state.get_config(
"snql_subscription_rollout_pct", 0.0)
assert isinstance(snql_rollout_pct, float)
snql_rollout_projects_raw = state.get_config(
"snql_subscription_rollout_projects", "")
snql_rollout_projects: Set[int]
if isinstance(snql_rollout_projects_raw, int):
snql_rollout_projects = {snql_rollout_projects_raw}
elif isinstance(snql_rollout_projects_raw, str):
snql_rollout_projects = (set([
int(s.strip())
for s in snql_rollout_projects_raw.split(",")
]) if snql_rollout_projects_raw else set())
else:
raise ValueError(
f"invalid project setting: '{snql_rollout_projects_raw}'")
use_snql = self.project_id in snql_rollout_projects or (
snql_rollout_pct > 0.0 and random.random() <= snql_rollout_pct)
if use_snql:
if metrics is not None:
metrics.increment("snql.subscription.delegate.use_snql")
return self.to_snql().build_request(dataset, timestamp, offset,
timer)
except Exception as e:
if metrics is not None:
metrics.increment("snql.subscription.delegate.error")
logger.warning(
f"failed snql subscription: {e}",
exc_info=e,
extra={
"error": str(e),
"project": self.project_id,
"query": self.query,
},
)
if metrics is not None:
metrics.increment("snql.subscription.delegate.use_legacy")
return self.to_legacy().build_request(dataset, timestamp, offset,
timer)
def __is_query_rolled_out(
self,
referrer: str,
config_referrer_prefix: str,
general_rollout_config: str,
) -> bool:
rollout_percentage = get_config(
f"rollout_upgraded_{self.__config_prefix}_{config_referrer_prefix}_{referrer}",
None,
)
if rollout_percentage is None:
rollout_percentage = get_config(general_rollout_config, 0.0)
return random() <= cast(float, rollout_percentage)
def process_query(self, query: Query,
query_settings: QuerySettings) -> None:
having_clause = query.get_having()
if not having_clause:
return None
selected_columns = query.get_selected_columns()
uniq_matcher = Param("function", FunctionCallMatch(String("uniq")))
found_functions = []
for exp in having_clause:
match = uniq_matcher.match(exp)
if match is not None:
found_functions.append(match.expression("function"))
if found_functions is not None:
matcher = _ExpressionOrAliasMatcher(found_functions)
for col in selected_columns:
col.expression.accept(matcher)
if not all(matcher.found_expressions):
should_throw = get_config("throw_on_uniq_select_and_having",
False)
error = MismatchedAggregationException(
"Aggregation is in HAVING clause but not SELECT",
query=str(query))
if should_throw:
raise error
else:
logging.warning(
"Aggregation is in HAVING clause but not SELECT",
exc_info=True,
extra=cast(Dict[str, Any], error.to_dict()),
)
def skip_kafka_message(message: Message[KafkaPayload]) -> bool:
# expected format is "[topic:partition_index:offset,...]" eg [snuba-metrics:0:1,snuba-metrics:0:3]
messages_to_skip = (get_config("kafka_messages_to_skip")
or "[]")[1:-1].split(",")
return (
f"{message.partition.topic.name}:{message.partition.index}:{message.offset}"
in messages_to_skip)
def _get_cache_partition(reader: Reader) -> Cache[Result]:
enable_cache_partitioning = state.get_config("enable_cache_partitioning", 1)
if not enable_cache_partitioning:
return cache_partitions[DEFAULT_CACHE_PARTITION_ID]
partition_id = reader.cache_partition_id
if partition_id is not None and partition_id not in cache_partitions:
with cache_partitions_lock:
# This condition was checked before as this lock should be acquired only
# during the first query. So, for the vast majority of queries, the overhead
# of acquiring the lock is not needed.
if partition_id not in cache_partitions:
exception = (
TigerExecutionTimeoutError
if "tiger" in partition_id
else ExecutionTimeoutError
)
cache_partitions[partition_id] = RedisCache(
redis_client,
f"snuba-query-cache:{partition_id}:",
ResultCacheCodec(),
ThreadPoolExecutor(),
exception,
)
return cache_partitions[
partition_id if partition_id is not None else DEFAULT_CACHE_PARTITION_ID
]
def test_config(self):
state.set_config("foo", 1)
state.set_configs({"bar": 2, "baz": 3})
assert state.get_config("foo") == 1
assert state.get_config("bar") == 2
assert state.get_config("noexist", 4) == 4
all_configs = state.get_all_configs()
assert all(all_configs[k] == v
for k, v in [("foo", 1), ("bar", 2), ("baz", 3)])
assert state.get_configs([("foo", 100), ("bar", 200), ("noexist", 300),
("noexist-2", None)]) == [1, 2, 300, None]
state.set_configs({"bar": "quux"})
all_configs = state.get_all_configs()
assert all(all_configs[k] == v
for k, v in [("foo", 1), ("bar", "quux"), ("baz", 3)])
def do_post_processing(
self,
project_ids: Sequence[int],
query: Query,
request_settings: RequestSettings,
) -> None:
if not request_settings.get_turbo():
final, exclude_group_ids = get_projects_query_flags(
project_ids, self.__replacer_state_name)
if not final and exclude_group_ids:
# If the number of groups to exclude exceeds our limit, the query
# should just use final instead of the exclusion set.
max_group_ids_exclude = get_config(
"max_group_ids_exclude",
settings.REPLACER_MAX_GROUP_IDS_TO_EXCLUDE)
if len(exclude_group_ids) > max_group_ids_exclude:
query.set_final(True)
else:
query.add_conditions([(["assumeNotNull", ["group_id"]],
"NOT IN", exclude_group_ids)])
query.add_condition_to_ast(
not_in_condition(
None,
FunctionCall(None, "assumeNotNull",
(Column(None, "group_id", None), )),
[Literal(None, p) for p in exclude_group_ids],
))
else:
query.set_final(final)
def test_config(self):
state.set_config('foo', 1)
state.set_configs({'bar': 2, 'baz': 3})
assert state.get_config('foo') == 1
assert state.get_config('bar') == 2
assert state.get_config('noexist', 4) == 4
all_configs = state.get_all_configs()
assert all(all_configs[k] == v
for k, v in [('foo', 1), ('bar', 2), ('baz', 3)])
assert state.get_configs([('foo', 100), ('bar', 200), ('noexist', 300),
('noexist-2', None)]) == [1, 2, 300, None]
state.set_configs({'bar': 'quux'})
all_configs = state.get_all_configs()
assert all(all_configs[k] == v
for k, v in [('foo', 1), ('bar', 'quux'), ('baz', 3)])
def process_query(self, query: Query,
request_settings: RequestSettings) -> None:
missing_checkers = {checker for checker in self.__condition_checkers}
def inspect_expression(condition: Expression) -> None:
top_level = get_first_level_and_conditions(condition)
for condition in top_level:
for checker in self.__condition_checkers:
if checker in missing_checkers:
if checker.check(condition):
missing_checkers.remove(checker)
condition = query.get_condition()
if condition is not None:
inspect_expression(condition)
prewhere = query.get_prewhere_ast()
if prewhere is not None:
inspect_expression(prewhere)
missing_ids = {checker.get_id() for checker in missing_checkers}
if get_config("mandatory_condition_enforce", 0):
assert (
not missing_checkers
), f"Missing mandatory columns in query. Missing {missing_ids}"
else:
if missing_checkers:
logger.error(
"Query is missing mandatory columns",
extra={"missing_checkers": missing_ids},
)
def process_query(self, query: Query, request_settings: RequestSettings) -> None:
if not get_config(self.__killswitch, 1):
return
cond_class = ConditionClass.IRRELEVANT
condition = query.get_condition()
if condition is not None:
cond_class = self.__classify_combined_conditions(condition)
if cond_class == ConditionClass.NOT_OPTIMIZABLE:
return
having_cond_class = ConditionClass.IRRELEVANT
having_cond = query.get_having()
if having_cond is not None:
having_cond_class = self.__classify_combined_conditions(having_cond)
if having_cond_class == ConditionClass.NOT_OPTIMIZABLE:
return
if not (
cond_class == ConditionClass.OPTIMIZABLE
or having_cond_class == ConditionClass.OPTIMIZABLE
):
return
metrics.increment("optimizable_query")
if condition is not None:
query.set_ast_condition(condition.transform(self.__replace_with_hash))
if having_cond is not None:
query.set_ast_having(having_cond.transform(self.__replace_with_hash))
def select_storage(
self, query: Query, request_settings: RequestSettings
) -> StorageAndMappers:
granularity = extract_granularity_from_query(query, "started") or 3600
use_materialized_storage = granularity >= 3600 and (granularity % 3600) == 0
metrics.increment(
"query.selector",
tags={
"selected_storage": "materialized"
if use_materialized_storage
else "raw",
},
)
allow_subhour_sessions = state.get_config("allow_subhour_sessions", 0)
if not allow_subhour_sessions:
use_materialized_storage = True
if use_materialized_storage:
return StorageAndMappers(
self.materialized_storage, sessions_hourly_translators
)
else:
return StorageAndMappers(self.raw_storage, sessions_raw_translators)
def process_query(self, query: Query, query_settings: QuerySettings) -> None:
if not get_config(self.__killswitch, 1):
return
condition, cond_class = self.__get_reduced_and_classified_query_clause(
query.get_condition(), query
)
query.set_ast_condition(condition)
if cond_class == ConditionClass.NOT_OPTIMIZABLE:
return
having_cond, having_cond_class = self.__get_reduced_and_classified_query_clause(
query.get_having(), query
)
query.set_ast_having(having_cond)
if having_cond_class == ConditionClass.NOT_OPTIMIZABLE:
return
if not (
cond_class == ConditionClass.OPTIMIZABLE
or having_cond_class == ConditionClass.OPTIMIZABLE
):
return
metrics.increment("optimizable_query")
query.add_experiment("tags_hashmap_applied", 1)
if condition is not None:
query.set_ast_condition(condition.transform(self.__replace_with_hash))
if having_cond is not None:
query.set_ast_having(having_cond.transform(self.__replace_with_hash))
def __enter__(self) -> Tuple[RateLimitResult, int]:
limit = (
state.get_config(f"{RATE_LIMIT_PER_SEC_KEY_PREFIX}{self.__bucket}", None)
if not self.__max_rate_per_sec
else self.__max_rate_per_sec
)
if not limit:
return (RateLimitResult.OFF, 0)
with self.__lock:
current_time = time.time()
current_epoch = int(current_time)
if (
self.__bucket_epoch is None
or self.__bucket_attempts is None
or current_epoch != self.__bucket_epoch
):
self.__bucket_epoch = current_epoch
self.__bucket_attempts = 1
ret_state = RateLimitResult.WITHIN_QUOTA
elif self.__bucket_attempts >= limit:
new_epoch = current_epoch + 1
time.sleep(new_epoch - current_time)
self.__bucket_epoch = new_epoch
self.__bucket_attempts = 1
ret_state = RateLimitResult.THROTTLED
else:
self.__bucket_epoch = current_epoch
self.__bucket_attempts += 1
ret_state = RateLimitResult.WITHIN_QUOTA
return (ret_state, self.__bucket_attempts)
def should_run(self) -> bool:
try:
unaliaser_config_percentage = float(
cast(float, get_config("tuple_unaliaser_rollout", 0)))
return random.random() < unaliaser_config_percentage
except ValueError:
return False
def execute(
self,
query: Query,
request_settings: RequestSettings,
runner: QueryRunner,
) -> QueryResult:
def process_and_run_query(
query: Query,
request_settings: RequestSettings) -> QueryResult:
for processor in self.__query_processors:
with sentry_sdk.start_span(
description=type(processor).__name__, op="processor"):
processor.process_query(query, request_settings)
return runner(query, request_settings, self.__cluster.get_reader())
use_split = state.get_config("use_split", 1)
if use_split:
for splitter in self.__splitters:
with sentry_sdk.start_span(description=type(splitter).__name__,
op="splitter"):
result = splitter.execute(query, request_settings,
process_and_run_query)
if result is not None:
return result
return process_and_run_query(query, request_settings)
def process_message(
self, message: Message[KafkaPayload]) -> Optional[Replacement]:
metadata = ReplacementMessageMetadata(
partition_index=message.partition.index,
offset=message.offset,
consumer_group=self.__consumer_group,
)
if self._message_already_processed(metadata):
logger.warning(
f"Replacer ignored a message, consumer group: {self.__consumer_group}",
extra={
"partition": metadata.partition_index,
"offset": metadata.offset,
},
)
if get_config("skip_seen_offsets", False):
return None
seq_message = json.loads(message.payload.value)
[version, action_type, data] = seq_message
if version == 2:
return self.__replacer_processor.process_message(
ReplacementMessage(
action_type=action_type,
data=data,
metadata=metadata,
))
else:
raise InvalidMessageVersion("Unknown message format: " +
str(seq_message))
def __get_filter_tags(self, query: Query) -> List[str]:
"""
Identifies the tag names we can apply the arrayFilter optimization on.
Which means: if the tags_key column is in the select clause and there are
one or more top level conditions on the tags_key column.
We can only apply the arrayFilter optimization to tag keys conditions
that are not in OR with other columns. To simplify the problem, we only
consider those conditions that are included in the first level of the query:
[['tagskey' '=' 'a'],['col' '=' 'b'],['col2' '=' 'c']] works
[[['tagskey' '=' 'a'], ['col2' '=' 'b']], ['tagskey' '=' 'c']] does not
"""
if not state.get_config("ast_tag_processor_enabled", 1):
return []
tags_key_found = any(
"tags_key" in columns_in_expr(expression)
for expression in query.get_selected_columns() or []
)
if not tags_key_found:
return []
def extract_tags_from_condition(
cond: Sequence[Condition],
) -> Optional[List[str]]:
if not cond:
return []
ret = []
for c in cond:
if not is_condition(c):
# This is an OR
return None
if c[1] == "=" and c[0] == "tags_key" and isinstance(c[2], str):
ret.append(str(c[2]))
elif (
c[1] == "IN"
and c[0] == "tags_key"
and isinstance(c[2], (list, tuple))
):
ret.extend([str(tag) for tag in c[2]])
return ret
cond_tags_key = extract_tags_from_condition(query.get_conditions() or [])
if cond_tags_key is None:
# This means we found an OR. Cowardly we give up even though there could
# be cases where this condition is still optimizable.
return []
having_tags_key = extract_tags_from_condition(query.get_having() or [])
if having_tags_key is None:
# Same as above
return []
return cond_tags_key + having_tags_key
def select_storage(self, query: Query,
request_settings: RequestSettings) -> StorageAndMappers:
use_readonly_storage = (state.get_config(
"enable_events_readonly_table", False)
and not request_settings.get_consistent())
storage = (self.__events_ro_table
if use_readonly_storage else self.__events_table)
return StorageAndMappers(storage, event_translator)
def should_write_every_node(self) -> bool:
project_rollout_setting = get_config("write_node_replacements_projects", "")
if project_rollout_setting:
# The expected for mat is [project,project,...]
project_rollout_setting = project_rollout_setting[1:-1]
if project_rollout_setting:
rolled_out_projects = [
int(p.strip()) for p in project_rollout_setting.split(",")
]
if self.get_project_id() in rolled_out_projects:
return True
global_rollout_setting = get_config("write_node_replacements_global", 0.0)
assert isinstance(global_rollout_setting, float)
if random.random() < global_rollout_setting:
return True
return False
def select_storage(self, query: Query,
query_settings: QuerySettings) -> StorageAndMappers:
readonly_referrer = (
query_settings.referrer
in settings.TRANSACTIONS_DIRECT_TO_READONLY_REFERRERS)
use_readonly_storage = readonly_referrer or state.get_config(
"enable_transactions_readonly_table", False)
storage = (self.__transactions_ro_table
if use_readonly_storage else self.__transactions_table)
return StorageAndMappers(storage, self.__mappers)
def validate(self, func_name: str, parameters: Sequence[Expression],
data_source: DataSource) -> None:
if is_valid_global_function(func_name):
return
if state.get_config("function-validator.enabled", False):
raise InvalidFunctionCall(f"Invalid function name: {func_name}")
else:
metrics.increment("invalid_funcs", tags={"func_name": func_name})
def build_planner(
self,
query: LogicalQuery,
settings: RequestSettings,
) -> EntityQueryPlanner:
new_query = deepcopy(query)
sampling_rate = state.get_config("snuplicator-sampling-rate", 1.0)
assert isinstance(sampling_rate, float)
new_query.set_sample(sampling_rate)
return super().build_planner(new_query, settings)
def sampling_selector_func(query: LogicalQuery,
referrer: str) -> Tuple[str, List[str]]:
if is_in_experiment(query, referrer):
sample_query_rate = state.get_config(
"snuplicator-sampling-experiment-rate", 0.0)
assert isinstance(sample_query_rate, float)
if random.random() < sample_query_rate:
return "primary", ["sampler"]
return "primary", []
def __build_consumer(
self, strategy_factory: ProcessingStrategyFactory[KafkaPayload]
) -> StreamProcessor[KafkaPayload]:
configuration = build_kafka_consumer_configuration(
self.storage.get_table_writer().get_stream_loader().
get_default_topic_spec().topic,
bootstrap_servers=self.bootstrap_servers,
group_id=self.group_id,
auto_offset_reset=self.auto_offset_reset,
strict_offset_reset=self.strict_offset_reset,
queued_max_messages_kbytes=self.queued_max_messages_kbytes,
queued_min_messages=self.queued_min_messages,
)
if self.__cooperative_rebalancing is True:
configuration[
"partition.assignment.strategy"] = "cooperative-sticky"
stats_collection_frequency_ms = get_config(
f"stats_collection_freq_ms_{self.group_id}",
get_config("stats_collection_freq_ms", 0),
)
if stats_collection_frequency_ms and stats_collection_frequency_ms > 0:
configuration.update({
"statistics.interval.ms": stats_collection_frequency_ms,
"stats_cb": self.stats_callback,
})
if self.commit_log_topic is None:
consumer = KafkaConsumer(
configuration,
commit_retry_policy=self.__commit_retry_policy,
)
else:
consumer = KafkaConsumerWithCommitLog(
configuration,
producer=self.producer,
commit_log_topic=self.commit_log_topic,
commit_retry_policy=self.__commit_retry_policy,
)
return StreamProcessor(consumer, self.raw_topic, strategy_factory)
def submit(self, message: Message[KafkaPayload]) -> None:
assert not self.__closed
# If there are max_concurrent_queries + 10 pending futures in the queue,
# we will start raising MessageRejected to slow down the consumer as
# it means our executor cannot keep up
queue_size_factor = state.get_config("executor_queue_size_factor", 10)
assert (queue_size_factor
is not None), "Invalid executor_queue_size_factor config"
max_queue_size = self.__max_concurrent_queries * queue_size_factor
# Tell the consumer to pause until we have removed some futures from
# the queue
if len(self.__queue) >= max_queue_size:
raise MessageRejected
task = self.__encoder.decode(message.payload)
tick_upper_offset = task.task.tick_upper_offset
entity = task.task.entity
entity_name = entity.value
should_execute = entity_name in self.__entity_names
# Don't execute stale subscriptions
if (self.__stale_threshold_seconds is not None
and time.time() - datetime.timestamp(task.timestamp) >=
self.__stale_threshold_seconds):
should_execute = False
if should_execute:
self.__queue.append((
message,
SubscriptionTaskResultFuture(
task,
self.__executor.submit(self.__execute_query, task,
tick_upper_offset),
),
))
else:
self.__metrics.increment("skipped_execution",
tags={"entity": entity_name})
# Periodically commit offsets if we haven't started rollout yet
self.__commit_data[message.partition] = Position(
message.next_offset, message.timestamp)
now = time.time()
if (self.__last_committed is None
or now - self.__last_committed >= COMMIT_FREQUENCY_SEC):
self.__commit(self.__commit_data)
self.__last_committed = now
self.__commit_data = {}
def process_query(self, query: Query, query_settings: QuerySettings) -> None:
enabled = get_config(ENABLED_CONFIG, 1)
if not enabled:
return
project_ids = get_object_ids_in_query_ast(query, self.__project_field)
if not project_ids:
return
# TODO: Like for the rate limiter Add logic for multiple IDs
project_id = str(project_ids.pop())
thread_quota = get_config(
f"{REFERRER_PROJECT_CONFIG}_{query_settings.referrer}_{project_id}"
)
if not thread_quota:
return
assert isinstance(thread_quota, int)
query_settings.set_resource_quota(ResourceQuota(max_threads=thread_quota))
def _consistent_override(original_setting: bool, referrer: str) -> bool:
consistent_config = state.get_config("consistent_override", None)
if isinstance(consistent_config, str):
referrers_override = consistent_config.split(";")
for config in referrers_override:
referrer_config, percentage = config.split("=")
if referrer_config == referrer:
if random.random() > float(percentage):
return False
return original_setting
