typing(issue_search): Add typing to issue search backends and executors (#29902)

Was in this area and thought this could do with type hints since it's complex

mypy.ini

@@ -51,6 +51,9 @@ files = src/sentry/api/bases/external_actor.py,
         src/sentry/notifications/**/*.py,
         src/sentry/processing/realtime_metrics/,
         src/sentry/release_health/**/*.py,
+        src/sentry/search/base.py,
+        src/sentry/search/events/constants.py,
+        src/sentry/search/snuba/*.py,
         src/sentry/sentry_metrics/**/*.py,
         src/sentry/shared_integrations/constants.py,
         src/sentry/snuba/outcomes.py,

src/sentry/search/base.py

@@ -1,26 +1,38 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, FrozenSet, Mapping, Optional, Sequence
+
 from sentry.utils.services import Service
 
 ANY = object()
 
+if TYPE_CHECKING:
+    from sentry.api.event_search import SearchFilter
+    from sentry.models import Environment, Project
+    from sentry.utils.cursors import Cursor, CursorResult
+
 
-class SearchBackend(Service):
+class SearchBackend(Service):  # type: ignore
     __read_methods__ = frozenset(["query"])
-    __write_methods__ = frozenset()
+    __write_methods__: FrozenSet[str] = frozenset()
     __all__ = __read_methods__ | __write_methods__
 
-    def __init__(self, **options):
+    def __init__(self, **options: Optional[Mapping[str, Any]]):
         pass
 
     def query(
         self,
-        projects,
-        tags=None,
-        environments=None,
-        sort_by="date",
-        limit=100,
-        cursor=None,
-        count_hits=False,
-        paginator_options=None,
-        **parameters,
-    ):
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]] = None,
+        sort_by: str = "date",
+        limit: int = 100,
+        cursor: Optional[Cursor] = None,
+        count_hits: bool = False,
+        paginator_options: Optional[Mapping[str, Any]] = None,
+        search_filters: Optional[Sequence[SearchFilter]] = None,
+        date_from: Optional[datetime] = None,
+        date_to: Optional[datetime] = None,
+        max_hits: Optional[int] = None,
+    ) -> CursorResult:
         raise NotImplementedError

src/sentry/search/snuba/backend.py

@@ -1,16 +1,20 @@
+from __future__ import annotations
+
 import functools
 from abc import ABCMeta, abstractmethod
 from collections import defaultdict
-from datetime import timedelta
-from typing import Sequence
+from datetime import datetime, timedelta
+from typing import Any, Callable, Dict, Mapping, Optional, Sequence
 
-from django.db.models import Q
+from django.db.models import Q, QuerySet
 from django.utils import timezone
 from django.utils.functional import SimpleLazyObject
 
 from sentry import quotas
+from sentry.api.event_search import SearchFilter
 from sentry.exceptions import InvalidSearchQuery
 from sentry.models import (
+    Environment,
     Group,
     GroupAssignee,
     GroupEnvironment,
@@ -30,10 +34,17 @@ from sentry.models import (
 )
 from sentry.search.base import SearchBackend
 from sentry.search.events.constants import EQUALITY_OPERATORS, OPERATOR_TO_DJANGO
-from sentry.search.snuba.executors import CdcPostgresSnubaQueryExecutor, PostgresSnubaQueryExecutor
+from sentry.search.snuba.executors import (
+    AbstractQueryExecutor,
+    CdcPostgresSnubaQueryExecutor,
+    PostgresSnubaQueryExecutor,
+)
+from sentry.utils.cursors import Cursor, CursorResult
 
 
-def assigned_to_filter(actors, projects, field_filter="id"):
+def assigned_to_filter(
+    actors: Sequence[User | Team | None], projects: Sequence[Project], field_filter: str = "id"
+) -> Q:
     from sentry.models import OrganizationMember, OrganizationMemberTeam, Team
 
     include_none = False
@@ -88,7 +99,7 @@ def assigned_to_filter(actors, projects, field_filter="id"):
     return query
 
 
-def unassigned_filter(unassigned, projects, field_filter="id"):
+def unassigned_filter(unassigned: bool, projects: Sequence[Project], field_filter: str = "id") -> Q:
     query = Q(
         **{
             f"{field_filter}__in": GroupAssignee.objects.filter(
@@ -101,7 +112,7 @@ def unassigned_filter(unassigned, projects, field_filter="id"):
     return query
 
 
-def linked_filter(linked, projects):
+def linked_filter(linked: bool, projects: Sequence[Project]) -> Q:
     """
     Builds a filter for whether or not a Group has an issue linked via either
     a PlatformExternalIssue or an ExternalIssue.
@@ -138,7 +149,9 @@ def linked_filter(linked, projects):
     return query
 
 
-def first_release_all_environments_filter(versions, projects):
+def first_release_all_environments_filter(
+    versions: Sequence[str], projects: Sequence[Project]
+) -> Q:
     releases = {
         id_: version
         for id_, version in Release.objects.filter(
@@ -164,7 +177,7 @@ def first_release_all_environments_filter(versions, projects):
     )
 
 
-def inbox_filter(inbox, projects):
+def inbox_filter(inbox: bool, projects: Sequence[Project]) -> Q:
     query = Q(groupinbox__id__isnull=False)
     if not inbox:
         query = ~query
@@ -173,7 +186,9 @@ def inbox_filter(inbox, projects):
     return query
 
 
-def assigned_or_suggested_filter(owners, projects, field_filter="id"):
+def assigned_or_suggested_filter(
+    owners: Sequence[User | Team | None], projects: Sequence[Project], field_filter: str = "id"
+) -> Q:
     organization_id = projects[0].organization_id
     project_ids = [p.id for p in projects]
 
@@ -269,15 +284,15 @@ class Condition:
     ``QuerySetBuilder``.
     """
 
-    def apply(self, queryset, name, parameters):
+    def apply(self, queryset: QuerySet, search_filter: SearchFilter) -> QuerySet:
         raise NotImplementedError
 
 
 class QCallbackCondition(Condition):
-    def __init__(self, callback):
+    def __init__(self, callback: Callable[[Any], QuerySet]):
         self.callback = callback
 
-    def apply(self, queryset, search_filter):
+    def apply(self, queryset: QuerySet, search_filter: SearchFilter) -> QuerySet:
         value = search_filter.value.raw_value
         q = self.callback(value)
         if search_filter.operator not in ("=", "!=", "IN", "NOT IN"):
@@ -297,17 +312,17 @@ class ScalarCondition(Condition):
     instances
     """
 
-    def __init__(self, field, extra=None):
+    def __init__(self, field: str, extra: Optional[dict[str, Sequence[int]]] = None):
         self.field = field
         self.extra = extra
 
-    def _get_operator(self, search_filter):
+    def _get_operator(self, search_filter: SearchFilter) -> str:
         django_operator = OPERATOR_TO_DJANGO.get(search_filter.operator, "")
         if django_operator:
             django_operator = f"__{django_operator}"
         return django_operator
 
-    def apply(self, queryset, search_filter):
+    def apply(self, queryset: QuerySet, search_filter: SearchFilter) -> QuerySet:
         django_operator = self._get_operator(search_filter)
         qs_method = queryset.exclude if search_filter.operator == "!=" else queryset.filter
 
@@ -319,10 +334,10 @@ class ScalarCondition(Condition):
 
 
 class QuerySetBuilder:
-    def __init__(self, conditions):
+    def __init__(self, conditions: Mapping[str, Condition]):
         self.conditions = conditions
 
-    def build(self, queryset, search_filters):
+    def build(self, queryset: QuerySet, search_filters: Sequence[SearchFilter]) -> QuerySet:
         for search_filter in search_filters:
             name = search_filter.key.name
             if name in self.conditions:
@@ -334,18 +349,18 @@ class QuerySetBuilder:
 class SnubaSearchBackendBase(SearchBackend, metaclass=ABCMeta):
     def query(
         self,
-        projects,
-        environments=None,
-        sort_by="date",
-        limit=100,
-        cursor=None,
-        count_hits=False,
-        paginator_options=None,
-        search_filters=None,
-        date_from=None,
-        date_to=None,
-        max_hits=None,
-    ):
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]] = None,
+        sort_by: str = "date",
+        limit: int = 100,
+        cursor: Optional[Cursor] = None,
+        count_hits: bool = False,
+        paginator_options: Optional[Mapping[str, Any]] = None,
+        search_filters: Optional[Sequence[SearchFilter]] = None,
+        date_from: Optional[datetime] = None,
+        date_to: Optional[datetime] = None,
+        max_hits: Optional[int] = None,
+    ) -> CursorResult:
         search_filters = search_filters if search_filters is not None else []
 
         # ensure projects are from same org
@@ -401,8 +416,14 @@ class SnubaSearchBackendBase(SearchBackend, metaclass=ABCMeta):
         )
 
     def _build_group_queryset(
-        self, projects, environments, search_filters, retention_window_start, *args, **kwargs
-    ):
+        self,
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]],
+        search_filters: Sequence[SearchFilter],
+        retention_window_start: Optional[datetime],
+        *args: Any,
+        **kwargs: Any,
+    ) -> QuerySet:
         """This method should return a QuerySet of the Group model.
         How you implement it is up to you, but we generally take in the various search parameters
         and filter Group's down using the field's we want to query on in Postgres."""
@@ -419,8 +440,12 @@ class SnubaSearchBackendBase(SearchBackend, metaclass=ABCMeta):
         return group_queryset
 
     def _initialize_group_queryset(
-        self, projects, environments, retention_window_start, search_filters
-    ):
+        self,
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]],
+        retention_window_start: Optional[datetime],
+        search_filters: Sequence[SearchFilter],
+    ) -> QuerySet:
         group_queryset = Group.objects.filter(project__in=projects).exclude(
             status__in=[
                 GroupStatus.PENDING_DELETION,
@@ -442,25 +467,41 @@ class SnubaSearchBackendBase(SearchBackend, metaclass=ABCMeta):
         return group_queryset
 
     @abstractmethod
-    def _get_queryset_conditions(self, projects, environments, search_filters):
+    def _get_queryset_conditions(
+        self,
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]],
+        search_filters: Sequence[SearchFilter],
+    ) -> Mapping[str, Condition]:
         """This method should return a dict of query set fields and a "Condition" to apply on that field."""
         raise NotImplementedError
 
     @abstractmethod
     def _get_query_executor(
-        self, group_queryset, projects, environments, search_filters, date_from, date_to
-    ):
+        self,
+        group_queryset: QuerySet,
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]],
+        search_filters: Sequence[SearchFilter],
+        date_from: Optional[datetime],
+        date_to: Optional[datetime],
+    ) -> AbstractQueryExecutor:
         """This method should return an implementation of the AbstractQueryExecutor
         We will end up calling .query() on the class returned by this method"""
         raise NotImplementedError
 
 
 class EventsDatasetSnubaSearchBackend(SnubaSearchBackendBase):
-    def _get_query_executor(self, *args, **kwargs):
+    def _get_query_executor(self, *args: Any, **kwargs: Any) -> AbstractQueryExecutor:
         return PostgresSnubaQueryExecutor()
 
-    def _get_queryset_conditions(self, projects, environments, search_filters):
-        queryset_conditions = {
+    def _get_queryset_conditions(
+        self,
+        projects: Sequence[Project],
+        environments: Optional[Sequence[Environment]],
+        search_filters: Sequence[SearchFilter],
+    ) -> Mapping[str, Condition]:
+        queryset_conditions: Dict[str, Condition] = {
             "status": QCallbackCondition(lambda statuses: Q(status__in=statuses)),
             "bookmarked_by": QCallbackCondition(
                 lambda users: Q(bookmark_set__project__in=projects, bookmark_set__user__in=users)
@@ -525,5 +566,5 @@ class EventsDatasetSnubaSearchBackend(SnubaSearchBackendBase):
 
 
 class CdcEventsDatasetSnubaSearchBackend(EventsDatasetSnubaSearchBackend):
-    def _get_query_executor(self, *args, **kwargs):
+    def _get_query_executor(self, *args: Any, **kwargs: Any) -> CdcPostgresSnubaQueryExecutor:
         return CdcPostgresSnubaQueryExecutor()

src/sentry/search/snuba/executors.py

@@ -1,15 +1,18 @@
+from __future__ import annotations
+
 import logging
 import time
-from abc import ABCMeta, abstractmethod, abstractproperty
+from abc import ABCMeta, abstractmethod
 from dataclasses import replace
 from datetime import datetime, timedelta
 from hashlib import md5
-from typing import Any, Mapping, Sequence
+from typing import Any, List, Mapping, Sequence, Set, Tuple
 
 import sentry_sdk
 from django.db.models import QuerySet
 from django.utils import timezone
 from snuba_sdk import Direction, Op
+from snuba_sdk.expressions import Expression
 from snuba_sdk.query import Column, Condition, Entity, Function, Join, Limit, OrderBy, Query
 from snuba_sdk.relationships import Relationship
 
@@ -25,7 +28,7 @@ from sentry.utils import json, metrics, snuba
 from sentry.utils.cursors import Cursor, CursorResult
 
 
-def get_search_filter(search_filters, name, operator):
+def get_search_filter(search_filters: Sequence[SearchFilter], name: str, operator: str) -> Any:
     """
     Finds the value of a search filter with the passed name and operator. If
     multiple values are found, returns the most restrictive value
@@ -57,65 +60,78 @@ class AbstractQueryExecutor(metaclass=ABCMeta):
 
     TABLE_ALIAS = ""
 
-    @abstractproperty
-    def aggregation_defs(self):
+    @property
+    @abstractmethod
+    def aggregation_defs(self) -> Sequence[str] | Expression:
         """This method should return a dict of key:value
         where key is a field name for your aggregation
         and value is the aggregation function"""
         raise NotImplementedError
 
-    @abstractproperty
-    def dependency_aggregations(self):
+    @property
+    @abstractmethod
+    def dependency_aggregations(self) -> Mapping[str, List[str]]:
         """This method should return a dict of key:value
         where key is an aggregation_def field name
         and value is a list of aggregation field names that the 'key' aggregation requires."""
         raise NotImplementedError
 
     @property
-    def empty_result(self):
+    def empty_result(self) -> CursorResult:
         return Paginator(Group.objects.none()).get_result()
 
     @property
     @abstractmethod
-    def dataset(self):
+    def dataset(self) -> snuba.Dataset:
         """ "This function should return an enum from snuba.Dataset (like snuba.Dataset.Events)"""
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def sort_strategies(self) -> Mapping[str, str]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def postgres_only_fields(self) -> Set[str]:
+        raise NotImplementedError
+
     @abstractmethod
     def query(
         self,
-        projects,
-        retention_window_start,
-        group_queryset,
-        environments,
-        sort_by,
-        limit,
-        cursor,
-        count_hits,
-        paginator_options,
-        search_filters,
-        date_from,
-        date_to,
-    ):
+        projects: Sequence[Project],
+        retention_window_start: Optional[datetime],
+        group_queryset: QuerySet,
+        environments: Optional[Sequence[Environment]],
+        sort_by: str,
+        limit: int,
+        cursor: Cursor,
+        count_hits: bool,
+        paginator_options: Optional[Mapping[str, Any]],
+        search_filters: Optional[Sequence[SearchFilter]],
+        date_from: Optional[datetime],
+        date_to: Optional[datetime],
+        max_hits: Optional[int] = None,
+    ) -> CursorResult:
         """This function runs your actual query and returns the results
         We usually return a paginator object, which contains the results and the number of hits"""
         raise NotImplementedError
 
     def snuba_search(
         self,
-        start,
-        end,
-        project_ids,
-        environment_ids,
-        sort_field,
-        organization_id,
-        cursor=None,
-        group_ids=None,
-        limit=None,
-        offset=0,
-        get_sample=False,
-        search_filters=None,
-    ):
+        start: datetime,
+        end: datetime,
+        project_ids: Sequence[int],
+        environment_ids: Sequence[int],
+        sort_field: str,
+        organization_id: int,
+        cursor: Optional[Cursor] = None,
+        group_ids: Optional[Sequence[int]] = None,
+        limit: Optional[int] = None,
+        offset: int = 0,
+        get_sample: bool = False,
+        search_filters: Optional[Sequence[SearchFilter]] = None,
+    ) -> Tuple[List[Tuple[int, Any]], int]:
         """
         Returns a tuple of:
         * a sorted list of (group_id, group_score) tuples sorted descending by score,
@@ -227,8 +243,12 @@ class AbstractQueryExecutor(metaclass=ABCMeta):
         return [(row["group_id"], row[sort_field]) for row in rows], total
 
     def _transform_converted_filter(
-        self, search_filter, converted_filter, project_ids, environment_ids=None
-    ):
+        self,
+        search_filter: Sequence[SearchFilter],
+        converted_filter: Optional[Sequence[any]],
+        project_ids: Sequence[int],
+        environment_ids: Optional[Sequence[int]] = None,
+    ) -> Optional[Sequence[any]]:
         """
         This method serves as a hook - after we convert the search_filter into a
         snuba compatible filter (which converts it in a general dataset
@@ -239,13 +259,13 @@ class AbstractQueryExecutor(metaclass=ABCMeta):
         """
         return converted_filter
 
-    def has_sort_strategy(self, sort_by):
+    def has_sort_strategy(self, sort_by: str) -> bool:
         return sort_by in self.sort_strategies.keys()
 
 
-def trend_aggregation(start, end):
-    middle = start + timedelta(seconds=(end - start).total_seconds() * 0.5)
-    middle = datetime.strftime(middle, DateArg.date_format)
+def trend_aggregation(start: datetime, end: datetime) -> Sequence[str]:
+    middle_date = start + timedelta(seconds=(end - start).total_seconds() * 0.5)
+    middle = datetime.strftime(middle_date, DateArg.date_format)
 
     agg_range_1 = f"countIf(greater(toDateTime('{middle}'), timestamp))"
     agg_range_2 = f"countIf(lessOrEquals(toDateTime('{middle}'), timestamp))"
@@ -298,25 +318,25 @@ class PostgresSnubaQueryExecutor(AbstractQueryExecutor):
     }
 
     @property
-    def dataset(self):
+    def dataset(self) -> snuba.Dataset:
         return snuba.Dataset.Events
 
     def query(
         self,
-        projects,
-        retention_window_start,
-        group_queryset,
-        environments,
-        sort_by,
-        limit,
-        cursor,
-        count_hits,
-        paginator_options,
-        search_filters,
-        date_from,
-        date_to,
-        max_hits=None,
-    ):
+        projects: Sequence[Project],
+        retention_window_start: Optional[datetime],
+        group_queryset: QuerySet,
+        environments: Optional[Sequence[Environment]],
+        sort_by: str,
+        limit: int,
+        cursor: Cursor,
+        count_hits: bool,
+        paginator_options: Optional[Mapping[str, Any]],
+        search_filters: Optional[Sequence[SearchFilter]],
+        date_from: Optional[datetime],
+        date_to: Optional[datetime],
+        max_hits: Optional[int] = None,
+    ) -> CursorResult:
 
         now = timezone.now()
         end = None
@@ -541,22 +561,22 @@ class PostgresSnubaQueryExecutor(AbstractQueryExecutor):
 
     def calculate_hits(
         self,
-        group_ids,
-        too_many_candidates,
-        sort_field,
-        projects,
-        retention_window_start,
-        group_queryset,
-        environments,
-        sort_by,
-        limit,
-        cursor,
-        count_hits,
-        paginator_options,
-        search_filters,
-        start,
-        end,
-    ):
+        group_ids: Sequence[int],
+        too_many_candidates: bool,
+        sort_field: str,
+        projects: Sequence[Project],
+        retention_window_start: Optional[datetime],
+        group_queryset: Query,
+        environments: Sequence[Environment],
+        sort_by: str,
+        limit: int,
+        cursor: Cursor,
+        count_hits: bool,
+        paginator_options: Mapping[str, Any],
+        search_filters: Sequence[SearchFilter],
+        start: datetime,
+        end: datetime,
+    ) -> Optional[int]:
         """
         This method should return an integer representing the number of hits (results) of your search.
         It will return 0 if hits were calculated and there are none.
@@ -718,7 +738,7 @@ class CdcPostgresSnubaQueryExecutor(PostgresSnubaQueryExecutor):
         search_filters: Sequence[SearchFilter],
         date_from: Optional[datetime],
         date_to: Optional[datetime],
-    ):
+    ) -> Tuple[datetime, datetime, datetime]:
         now = timezone.now()
         end = None
         end_params = [_f for _f in [date_to, get_search_filter(search_filters, "date", "<")] if _f]
@@ -748,7 +768,7 @@ class CdcPostgresSnubaQueryExecutor(PostgresSnubaQueryExecutor):
         search_filters: Sequence[SearchFilter],
         date_from: Optional[datetime],
         date_to: Optional[datetime],
-        max_hits=None,
+        max_hits: Optional[int] = None,
     ) -> CursorResult:
 
         if not validate_cdc_search_filters(search_filters):