From e17d631a05c1f87ebad7b93ec16072cad41752e4 Mon Sep 17 00:00:00 2001 From: Charles Paul Date: Fri, 14 Nov 2025 10:25:18 -0800 Subject: [PATCH] fix(typing): add types to tagstore I'm working in this neck of the woods and keep getting confused about what things are returning. Adding some (but not all) types to help myself get a stronger hold of things. Not fully typing due to prioritization + some not-expecting-None problems. TEST PLAN: `mypy` passes --- src/sentry/models/release.py | 2 +- src/sentry/models/releases/util.py | 2 +- src/sentry/tagstore/base.py | 4 +- src/sentry/tagstore/snuba/backend.py | 156 ++++++++++++++++----------- src/sentry/tagstore/types.py | 36 ++++++- src/sentry/utils/snuba.py | 8 +- 6 files changed, 131 insertions(+), 77 deletions(-) diff --git a/src/sentry/models/release.py b/src/sentry/models/release.py index e5ecee580d7c8c..23db2dc713af3f 100644 --- a/src/sentry/models/release.py +++ b/src/sentry/models/release.py @@ -134,7 +134,7 @@ def filter_by_stage( operator: str, value, project_ids: Sequence[int] | None = None, - environments: list[str] | None = None, + environments: Sequence[str | int] | None = None, ) -> models.QuerySet: return self.get_queryset().filter_by_stage( organization_id, operator, value, project_ids, environments diff --git a/src/sentry/models/releases/util.py b/src/sentry/models/releases/util.py index 66fa5b2829f5fe..946e928cc79b56 100644 --- a/src/sentry/models/releases/util.py +++ b/src/sentry/models/releases/util.py @@ -145,7 +145,7 @@ def filter_by_stage( operator: str, value, project_ids: Sequence[int] | None = None, - environments: list[str] | None = None, + environments: Sequence[str | int] | None = None, ) -> Self: from sentry.models.releaseprojectenvironment import ReleaseProjectEnvironment, ReleaseStages from sentry.search.events.filter import to_list diff --git a/src/sentry/tagstore/base.py b/src/sentry/tagstore/base.py index 6285167479c5d8..9928fc49f90b90 100644 --- a/src/sentry/tagstore/base.py +++ b/src/sentry/tagstore/base.py @@ -220,13 +220,13 @@ def get_tag_value_paginator_for_projects( def get_group_tag_value_iter( self, group: Group, - environment_ids: list[int | None], + environment_ids: Sequence[int | None], key: str, orderby: str = "-first_seen", limit: int = 1000, offset: int = 0, tenant_ids: dict[str, int | str] | None = None, - ) -> list[GroupTagValue]: + ) -> Sequence[GroupTagValue]: """ >>> get_group_tag_value_iter(group, 2, 3, 'environment') """ diff --git a/src/sentry/tagstore/snuba/backend.py b/src/sentry/tagstore/snuba/backend.py index e3f6c737af8a92..7a6b1c437f2223 100644 --- a/src/sentry/tagstore/snuba/backend.py +++ b/src/sentry/tagstore/snuba/backend.py @@ -4,7 +4,7 @@ import os import re from collections import defaultdict -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Never, Protocol, TypedDict @@ -101,7 +101,9 @@ def get_project_list(project_id: int | list[int]) -> list[int]: return project_id if isinstance(project_id, Iterable) else [project_id] -def _translate_filter_keys(project_ids, group_ids, environment_ids) -> dict[str, Any]: +def _translate_filter_keys( + project_ids, group_ids, environment_ids: Sequence[int] | None +) -> dict[str, Any]: filter_keys = {"project_id": project_ids} if environment_ids: @@ -164,16 +166,16 @@ class SnubaTagStorage(TagStorage): def __get_tag_key_and_top_values( self, project_id, - group, + group: Group | None, environment_id, key, - limit=3, + limit: int | None = 3, raise_on_empty=True, tenant_ids=None, **kwargs, - ): + ) -> GroupTagKey | TagKey: tag = self.format_string.format(key) - filters = {"project_id": get_project_list(project_id)} + filters: dict[str, Sequence[Any]] = {"project_id": get_project_list(project_id)} if environment_id: filters["environment"] = [environment_id] conditions = kwargs.get("conditions", []) @@ -231,17 +233,17 @@ def __get_tag_key_and_top_values( def __get_tag_keys( self, - project_id, - group, - environment_ids, - limit=1000, + project_id: int, + group: Group | None, + environment_ids: list[int] | None, + limit: int | None = 1000, keys: list[str] | None = None, - include_values_seen=True, + include_values_seen: bool = True, dataset: Dataset = Dataset.Events, - denylist=None, + denylist: frozenset[int | str] | None = None, tenant_ids=None, - **kwargs, - ): + **kwargs: Any, + ) -> set[TagKey | GroupTagKey]: optimize_kwargs: _OptimizeKwargs = {} if turbo := kwargs.get("turbo"): if isinstance(turbo, bool): @@ -269,17 +271,17 @@ def __get_tag_keys_for_projects( self, projects, group, - environments, - start, - end, - limit=1000, + environments: Sequence[int] | None, + start: datetime | None, + end: datetime | None, + limit: int | None = 1000, keys: list[str] | None = None, - include_values_seen=True, + include_values_seen: bool = True, use_cache=False, - denylist=None, + denylist: frozenset[int | str] | None = None, dataset: Dataset = Dataset.Discover, **kwargs, - ): + ) -> set[TagKey | GroupTagKey]: """Query snuba for tag keys based on projects When use_cache is passed, we'll attempt to use the cache. There's an exception if group_id was passed @@ -302,7 +304,7 @@ def __get_tag_keys_for_projects( aggregations = [["count()", "", "count"]] - filters = {"project_id": sorted(projects)} + filters: dict[str, Sequence[Any]] = {"project_id": sorted(projects)} if environments: filters["environment"] = sorted(environments) if group is not None: @@ -378,7 +380,7 @@ def __get_tag_keys_for_projects( else: ctor = functools.partial(GroupTagKey, group_id=group.id) - results = set() + results: set[TagKey | GroupTagKey] = set() for key, data in result.items(): # Ignore key (skip interaction) if it's in denylist @@ -405,7 +407,7 @@ def get_tag_key( status=TagKeyStatus.ACTIVE, tenant_ids=None, **kwargs, - ): + ) -> GroupTagKey | TagKey: assert status is TagKeyStatus.ACTIVE return self.__get_tag_key_and_top_values( project_id, None, environment_id, key, tenant_ids=tenant_ids, **kwargs @@ -413,14 +415,14 @@ def get_tag_key( def get_tag_keys( self, - project_id, - environment_id, - status=TagKeyStatus.ACTIVE, - include_values_seen=False, - denylist=None, + project_id: int, + environment_id: int | None, + status: int = TagKeyStatus.ACTIVE, + include_values_seen: bool = False, + denylist: frozenset[int | str] | None = None, tenant_ids=None, **kwargs, - ): + ) -> set[TagKey | GroupTagKey]: assert status is TagKeyStatus.ACTIVE optimize_kwargs: _OptimizeKwargs = {} @@ -437,7 +439,7 @@ def get_tag_keys( return self.__get_tag_keys( project_id, None, - environment_id and [environment_id], + environment_id if environment_id is None else [environment_id], denylist=denylist, tenant_ids=tenant_ids, include_values_seen=include_values_seen, @@ -447,15 +449,15 @@ def get_tag_keys( def get_tag_keys_for_projects( self, - projects, - environments, - start, - end, + projects: Sequence[int], + environments: Sequence[int] | None, + start: datetime, + end: datetime, dataset: Dataset = Dataset.Events, - status=TagKeyStatus.ACTIVE, + status: int = TagKeyStatus.ACTIVE, use_cache: bool = False, tenant_ids=None, - ): + ) -> set[TagKey | GroupTagKey]: max_unsampled_projects = _max_unsampled_projects # We want to disable FINAL in the snuba query to reduce load. optimize_kwargs: _OptimizeKwargs = {"turbo": True} @@ -525,12 +527,12 @@ def get_group_tag_key( def get_group_tag_keys( self, group, - environment_ids, - limit=None, + environment_ids: list[int], + limit: int | None = None, keys: list[str] | None = None, tenant_ids=None, **kwargs, - ): + ) -> set[TagKey | GroupTagKey]: """Get tag keys for a specific group""" return self.__get_tag_keys( group.project_id, @@ -548,7 +550,7 @@ def __get_group_list_tag_value( self, project_ids, group_id_list, - environment_ids, + environment_ids: list[int], key: str, value, dataset, @@ -587,7 +589,13 @@ def __get_group_list_tag_value( } def get_group_list_tag_value( - self, project_ids, group_id_list, environment_ids, key: str, value, tenant_ids=None + self, + project_ids, + group_id_list, + environment_ids: list[int], + key: str, + value, + tenant_ids=None, ): return self.__get_group_list_tag_value( project_ids, @@ -603,7 +611,13 @@ def get_group_list_tag_value( ) def get_generic_group_list_tag_value( - self, project_ids, group_id_list, environment_ids, key: str, value, tenant_ids=None + self, + project_ids, + group_id_list, + environment_ids: list[int], + key: str, + value, + tenant_ids=None, ): translated_params = _translate_filter_keys(project_ids, group_id_list, environment_ids) organization_id = get_organization_id_from_project_ids(project_ids) @@ -655,7 +669,7 @@ def get_generic_group_list_tag_value( for group_id, data in nested_groups.items() } - def apply_group_filters(self, group: Group | None, filters): + def apply_group_filters(self, group: Group | None, filters: MutableMapping[str, Sequence[Any]]): dataset = Dataset.Events if group: filters["group_id"] = [group.id] @@ -670,7 +684,7 @@ def get_group_tag_value_count( key: str, tenant_ids=None, ): - filters = {"project_id": get_project_list(group.project_id)} + filters: dict[str, Sequence[Any]] = {"project_id": get_project_list(group.project_id)} if environment_id: filters["environment"] = [environment_id] aggregations = [["count()", "", "count"]] @@ -689,7 +703,7 @@ def get_top_group_tag_values( group, environment_id, key: str, - limit=TOP_VALUES_DEFAULT_LIMIT, + limit: int = TOP_VALUES_DEFAULT_LIMIT, tenant_ids=None, ): tag = self.__get_tag_key_and_top_values( @@ -722,7 +736,9 @@ def get_group_tag_keys_and_top_values( ) # Then get the top values with first_seen/last_seen/count for each - filters: dict[str, list[Any]] = {"project_id": get_project_list(group.project_id)} + filters: MutableMapping[str, Sequence[Any]] = { + "project_id": get_project_list(group.project_id) + } conditions = kwargs.get("conditions", []) if environment_ids: @@ -800,12 +816,12 @@ def get_group_tag_keys_and_top_values( def __get_empty_value_stats_map( self, dataset: Dataset, - filters: dict[str, list[Any]], + filters: MutableMapping[str, Sequence[Any]], conditions: list, keys_to_check: list[str], tenant_ids: dict[str, int | str] | None, - start, - end, + start: datetime | None, + end: datetime | None, ) -> dict[str, dict[str, Any]]: stats_map: dict[str, dict[str, Any]] = {} if not keys_to_check: @@ -882,7 +898,9 @@ def get_release_tags(self, organization_id, project_ids, environment_id, version return set(values) - def get_min_start_date(self, organization_id, project_ids, environment_id, versions): + def get_min_start_date( + self, organization_id, project_ids, environment_id, versions + ) -> datetime | None: rpe_query = ReleaseProjectEnvironment.objects.filter( project_id__in=project_ids, release__version__in=versions, @@ -902,9 +920,9 @@ def __get_groups_user_counts( self, project_ids, group_ids, - environment_ids, - start=None, - end=None, + environment_ids: Sequence[int] | None, + start: datetime | None = None, + end: datetime | None = None, dataset=Dataset.Events, extra_aggregations=None, referrer=Referrer.TAGSTORE_GET_GROUPS_USER_COUNTS.value, @@ -1045,8 +1063,8 @@ def _get_semver_versions_for_package(self, projects, organization_id, package): def _get_tag_values_for_semver( self, - projects: Sequence[int], - environments: Sequence[str] | None, + projects: list[int], + environments: Sequence[int] | None, query: str | None, ): from sentry.api.paginator import SequencePaginator @@ -1113,7 +1131,9 @@ def _get_tag_values_for_semver( ] ) - def _get_tag_values_for_semver_package(self, projects, environments, package): + def _get_tag_values_for_semver_package( + self, projects, environments: Sequence[int] | None, package: str | None + ): from sentry.api.paginator import SequencePaginator package = package if package else "" @@ -1136,7 +1156,9 @@ def _get_tag_values_for_semver_package(self, projects, environments, package): ] ) - def _get_tag_values_for_release_stages(self, projects, environments, query): + def _get_tag_values_for_release_stages( + self, projects, environments: Sequence[int] | None, query: str | None + ): from sentry.api.paginator import SequencePaginator organization_id = Project.objects.filter(id=projects[0]).values_list( @@ -1164,7 +1186,9 @@ def _get_tag_values_for_release_stages(self, projects, environments, query): ] ) - def _get_tag_values_for_semver_build(self, projects, environments, build): + def _get_tag_values_for_semver_build( + self, projects, environments: Sequence[int] | None, build: str | None + ): from sentry.api.paginator import SequencePaginator build = build if build else "" @@ -1190,7 +1214,9 @@ def _get_tag_values_for_semver_build(self, projects, environments, build): [(i, TagValue(SEMVER_BUILD_ALIAS, v, None, None, None)) for i, v in enumerate(packages)] ) - def _get_tag_values_for_releases_across_all_datasets(self, projects, environments, query): + def _get_tag_values_for_releases_across_all_datasets( + self, projects, environments: Sequence[int] | None, query: str | None + ): from sentry.api.paginator import SequencePaginator organization_id = Project.objects.filter(id=projects[0]).values_list( @@ -1226,7 +1252,7 @@ def get_snuba_column_name(self, key: str, dataset: Dataset): def get_tag_value_paginator_for_projects( self, projects, - environments, + environments: Sequence[int] | None, key: str, start=None, end=None, @@ -1466,14 +1492,14 @@ def score_field_to_int(tv: TagValue) -> int: def get_group_tag_value_iter( self, group: Group, - environment_ids: list[int | None], + environment_ids: Sequence[int | None], key: str, orderby: str = "-first_seen", limit: int = 1000, offset: int = 0, tenant_ids: dict[str, int | str] | None = None, - ) -> list[GroupTagValue]: - filters: dict[str, list[Any]] = { + ) -> Sequence[GroupTagValue]: + filters: MutableMapping[str, Sequence[Any]] = { "project_id": get_project_list(group.project_id), } dataset, filters = self.apply_group_filters(group, filters) @@ -1508,7 +1534,7 @@ def get_group_tag_value_iter( def get_group_tag_value_paginator( self, group, - environment_ids, + environment_ids: list[int], key: str, order_by="-id", tenant_ids=None, diff --git a/src/sentry/tagstore/types.py b/src/sentry/tagstore/types.py index 1b14cbe1f8d166..a668dbae50bc7c 100644 --- a/src/sentry/tagstore/types.py +++ b/src/sentry/tagstore/types.py @@ -1,6 +1,7 @@ from __future__ import annotations import functools +from datetime import datetime from typing import Any, ClassVar, TypedDict from sentry.api.serializers import Serializer, register, serialize @@ -43,7 +44,12 @@ class TagKey(TagType): _sort_key = "values_seen" def __init__( - self, key, values_seen=None, status=TagKeyStatus.ACTIVE, count=None, top_values=None + self, + key: str, + values_seen: int | None = None, + status: int = TagKeyStatus.ACTIVE, + count: int | None = None, + top_values=None, ): self.key = key self.values_seen = values_seen @@ -59,7 +65,14 @@ class TagValue(TagType): __slots__ = ("key", "value", "times_seen", "first_seen", "last_seen") _sort_key = "value" - def __init__(self, key, value, times_seen, first_seen, last_seen): + def __init__( + self, + key: str, + value, + times_seen: int | None, + first_seen: datetime | None, + last_seen: datetime | None, + ): self.key = key self.value = value self.times_seen = times_seen @@ -71,7 +84,14 @@ class GroupTagKey(TagType): __slots__ = ("group_id", "key", "values_seen", "count", "top_values") _sort_key = "values_seen" - def __init__(self, group_id, key, values_seen=None, count=None, top_values=None): + def __init__( + self, + group_id: int, + key: str, + values_seen: int | None = None, + count: int | None = None, + top_values=None, + ): self.group_id = group_id self.key = key self.values_seen = values_seen @@ -83,7 +103,15 @@ class GroupTagValue(TagType): __slots__ = ("group_id", "key", "value", "times_seen", "first_seen", "last_seen") _sort_key = "value" - def __init__(self, group_id, key, value, times_seen, first_seen, last_seen): + def __init__( + self, + group_id: int, + key: str, + value, + times_seen: int, + first_seen, + last_seen, + ): self.group_id = group_id self.key = key self.value = value diff --git a/src/sentry/utils/snuba.py b/src/sentry/utils/snuba.py index 48785b7a7b2373..d5f63b1c7f6397 100644 --- a/src/sentry/utils/snuba.py +++ b/src/sentry/utils/snuba.py @@ -893,8 +893,8 @@ class SnubaQueryParams: def __init__( self, dataset=None, - start=None, - end=None, + start: datetime | None = None, + end: datetime | None = None, groupby=None, conditions=None, filter_keys=None, @@ -1501,8 +1501,8 @@ def _raw_snql_query(request: Request, headers: Mapping[str, str]) -> urllib3.res def query( dataset=None, - start=None, - end=None, + start: datetime | None = None, + end: datetime | None = None, groupby=None, conditions=None, filter_keys=None,