def process_query(self, query: Query, request_settings: RequestSettings) -> None: # If the settings don't already have a project rate limit, add one existing = request_settings.get_rate_limit_params() for ex in existing: if ex.rate_limit_name == PROJECT_RATE_LIMIT_NAME: return project_ids = get_object_ids_in_query_ast(query, self.project_column) if not project_ids: return # TODO: Use all the projects, not just one project_id = project_ids.pop() prl, pcl = get_configs([("project_per_second_limit", 1000), ("project_concurrent_limit", 1000)]) # Specific projects can have their rate limits overridden (per_second, concurr) = get_configs([ ("project_per_second_limit_{}".format(project_id), prl), ("project_concurrent_limit_{}".format(project_id), pcl), ]) rate_limit = RateLimitParameters( rate_limit_name=PROJECT_RATE_LIMIT_NAME, bucket=str(project_id), per_second_limit=per_second, concurrent_limit=concurr, ) request_settings.add_rate_limit(rate_limit)
def get_object_id( self, query: Query, query_settings: QuerySettings ) -> Optional[str]: obj_ids = get_object_ids_in_query_ast(query, self.object_column) if not obj_ids: return None # TODO: Add logic for multiple IDs obj_id = str(obj_ids.pop()) return str(obj_id)
def test_find_projects( query_body: MutableMapping[str, Any], expected_projects: Optional[Set[int]] ) -> None: events = get_dataset("events") if expected_projects is None: with pytest.raises(ParsingException): snql_query = json_to_snql(query_body, "events") query, _ = parse_snql_query(str(snql_query), events) identity_translate(query) else: snql_query = json_to_snql(query_body, "events") query, _ = parse_snql_query(str(snql_query), events) query = identity_translate(query) project_ids_ast = get_object_ids_in_query_ast(query, "project_id") assert project_ids_ast == expected_projects
def get_object_id( self, query: Query, query_settings: QuerySettings ) -> Optional[str]: obj_ids = get_object_ids_in_query_ast(query, self.object_column) if not obj_ids: return None # TODO: Add logic for multiple IDs obj_id = str(obj_ids.pop()) if self.query_settings_field is not None: query_settings_field_val = getattr( query_settings, self.query_settings_field, None ) if query_settings_field_val is not None: obj_id = f"{obj_id}_{query_settings_field_val}" return str(obj_id)
def is_in_experiment(query: LogicalQuery, referrer: str) -> bool: if referrer != "tagstore.__get_tag_keys": return False project_ids = get_object_ids_in_query_ast(query, "project_id") if not project_ids: return False test_projects_raw = state.get_config("snuplicator-sampling-projects", "") test_projects = set() if (isinstance(test_projects_raw, str) and test_projects_raw != ""): # should be in the form [1,2,3] test_projects_raw = test_projects_raw[1:-1] test_projects = set(int(p) for p in test_projects_raw.split(",") if p) elif isinstance(test_projects_raw, (int, float)): test_projects = {int(test_projects_raw)} return project_ids.issubset(test_projects)
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 _groups_to_exclude( self, query: Query, group_ids_to_exclude: Set[int] ) -> Set[int]: """ Given a Query and the group ids to exclude for any project this query touches, returns the intersection of the group ids from the replacements and the group ids this Query explicitly queries for, if any. Eg. - The query specifically looks for group ids: {1, 2} - The replacements on the projects require exclusion of groups: {1, 3, 4} - The query only needs to exclude group id 1 """ groups_in_query = get_object_ids_in_query_ast(query, self.__groups_column) if groups_in_query: group_ids_to_exclude = group_ids_to_exclude.intersection(groups_in_query) return group_ids_to_exclude
def process_query(self, query: Query, request_settings: RequestSettings) -> None: if request_settings.get_turbo(): return project_ids = get_object_ids_in_query_ast(query, self.__project_column) set_final = False if project_ids: final, exclude_group_ids = get_projects_query_flags( list(project_ids), self.__replacer_state_name, ) if final: metrics.increment("final", tags={"cause": "final_flag"}) 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) assert isinstance(max_group_ids_exclude, int) if len(exclude_group_ids) > max_group_ids_exclude: metrics.increment("final", tags={"cause": "max_groups"}) set_final = True else: query.add_condition_to_ast( not_in_condition( FunctionCall(None, "assumeNotNull", (Column(None, None, "group_id"), )), [Literal(None, p) for p in exclude_group_ids], )) else: set_final = final query.set_from_clause(replace(query.get_from_clause(), final=set_final))
def process_query(self, query: Query, query_settings: QuerySettings) -> None: if query_settings.get_turbo(): return project_ids = get_object_ids_in_query_ast(query, self.__project_column) if project_ids is None: self._set_query_final(query, False) return flags: ProjectsQueryFlags = ProjectsQueryFlags.load_from_redis( list(project_ids), self.__replacer_state_name ) query_overlaps_replacement = self._query_overlaps_replacements( query, flags.latest_replacement_time ) if not query_overlaps_replacement: self._set_query_final(query, False) return tags = self._initialize_tags(query_settings, flags) set_final = False if flags.needs_final: tags["cause"] = "final_flag" metrics.increment( name=FINAL_METRIC, tags=tags, ) set_final = True elif flags.group_ids_to_exclude: # 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, ) assert isinstance(max_group_ids_exclude, int) groups_to_exclude = self._groups_to_exclude( query, flags.group_ids_to_exclude ) if len(groups_to_exclude) > max_group_ids_exclude: tags["cause"] = "max_groups" metrics.increment( name=FINAL_METRIC, tags=tags, ) set_final = True elif groups_to_exclude: query.add_condition_to_ast( not_in_condition( FunctionCall( None, "assumeNotNull", (Column(None, None, self.__groups_column),), ), [Literal(None, p) for p in groups_to_exclude], ) ) self._set_query_final(query, set_final)
def build_request( body: MutableMapping[str, Any], parser: Parser, settings_class: Union[Type[HTTPQuerySettings], Type[SubscriptionQuerySettings]], schema: RequestSchema, dataset: Dataset, timer: Timer, referrer: str, custom_processing: Optional[CustomProcessors] = None, ) -> Request: with sentry_sdk.start_span(description="build_request", op="validate") as span: try: request_parts = schema.validate(body) if settings_class == HTTPQuerySettings: query_settings: MutableMapping[str, bool | str] = { **request_parts.query_settings, "consistent": _consistent_override( request_parts.query_settings.get("consistent", False), referrer ), } query_settings["referrer"] = referrer # TODO: referrer probably doesn't need to be passed in, it should be from the body settings_obj: Union[ HTTPQuerySettings, SubscriptionQuerySettings ] = settings_class( **query_settings, ) elif settings_class == SubscriptionQuerySettings: settings_obj = settings_class( consistent=_consistent_override(True, referrer), ) query, snql_anonymized = parser( request_parts, settings_obj, dataset, custom_processing ) project_ids = get_object_ids_in_query_ast(query, "project_id") if project_ids is not None and len(project_ids) == 1: sentry_sdk.set_tag("snuba_project_id", project_ids.pop()) org_ids = get_object_ids_in_query_ast(query, "org_id") if org_ids is not None and len(org_ids) == 1: sentry_sdk.set_tag("snuba_org_id", org_ids.pop()) attribution_info = dict(request_parts.attribution_info) # TODO: clean this up attribution_info["app_id"] = get_app_id( request_parts.attribution_info["app_id"] ) attribution_info["referrer"] = referrer request_id = uuid.uuid4().hex request = Request( id=request_id, # TODO: Replace this with the actual query raw body. # this can have an impact on subscriptions so we need # to be careful with the change. original_body=body, query=query, attribution_info=AttributionInfo(**attribution_info), query_settings=settings_obj, snql_anonymized=snql_anonymized, ) except (InvalidJsonRequestException, InvalidQueryException) as exception: record_invalid_request(timer, referrer) raise exception except Exception as exception: record_error_building_request(timer, referrer) raise exception span.set_data( "snuba_query_parsed", repr(query).split("\n"), ) span.set_data( "snuba_query_raw", textwrap.wrap(repr(request.original_body), 100, break_long_words=False), ) sentry_sdk.add_breadcrumb( category="query_info", level="info", message="snuba_query_raw", data={ "query": textwrap.wrap( repr(request.original_body), 100, break_long_words=False ) }, ) timer.mark("validate_schema") return request
def build_request( body: MutableMapping[str, Any], parser: Parser, settings_class: Union[Type[HTTPRequestSettings], Type[SubscriptionRequestSettings]], schema: RequestSchema, dataset: Dataset, timer: Timer, referrer: str, ) -> Request: with sentry_sdk.start_span(description="build_request", op="validate") as span: try: request_parts = schema.validate(body) if settings_class == HTTPRequestSettings: settings = { **request_parts.settings, "consistent": _consistent_override( request_parts.settings.get("consistent", False), referrer), } settings_obj: Union[ HTTPRequestSettings, SubscriptionRequestSettings] = settings_class(**settings) elif settings_class == SubscriptionRequestSettings: settings_obj = settings_class( consistent=_consistent_override(True, referrer)) query = parser(request_parts, settings_obj, dataset) project_ids = get_object_ids_in_query_ast(query, "project_id") if project_ids is not None and len(project_ids) == 1: sentry_sdk.set_tag("snuba_project_id", project_ids.pop()) org_ids = get_object_ids_in_query_ast(query, "org_id") if org_ids is not None and len(org_ids) == 1: sentry_sdk.set_tag("snuba_org_id", org_ids.pop()) request_id = uuid.uuid4().hex request = Request( request_id, # TODO: Replace this with the actual query raw body. # this can have an impact on subscriptions so we need # to be careful with the change. ChainMap(request_parts.query, *request_parts.extensions.values()), query, settings_obj, referrer, ) except (InvalidJsonRequestException, InvalidQueryException) as exception: record_invalid_request(timer, referrer) raise exception except Exception as exception: record_error_building_request(timer, referrer) raise exception span.set_data("snuba_query", request.body) timer.mark("validate_schema") return request
def _visit_simple_query(self, data_source: ProcessableQuery[Entity]) -> Set[int]: return get_object_ids_in_query_ast(data_source, "project_id") or set()
def test_find_projects(query_body: MutableMapping[str, Any], expected_projects: Set[int]) -> None: events = get_dataset("events") query = identity_translate(parse_query(query_body, events)) project_ids_ast = get_object_ids_in_query_ast(query, "project_id") assert project_ids_ast == expected_projects