diff --git a/src/sentry/tsdb/base.py b/src/sentry/tsdb/base.py index c9ce2416c65ad2..8575613cd7daea 100644 --- a/src/sentry/tsdb/base.py +++ b/src/sentry/tsdb/base.py @@ -14,6 +14,7 @@ ONE_DAY = ONE_HOUR * 24 TSDBKey = TypeVar("TSDBKey", str, int) +TSDBItem = TypeVar("TSDBItem", str, int) class IncrMultiOptions(TypedDict): @@ -661,13 +662,13 @@ def get_most_frequent_series( def get_frequency_series( self, model: TSDBModel, - items: Mapping[str, Sequence[str]], + items: Mapping[TSDBKey, Sequence[TSDBItem]], start: datetime, end: datetime | None = None, rollup: int | None = None, environment_id: int | None = None, tenant_ids: dict[str, str | int] | None = None, - ) -> dict[str, list[tuple[float, dict[str, float]]]]: + ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]: """ Retrieve the frequency of known items in a table over time. @@ -708,10 +709,10 @@ def merge_frequencies( self, model: TSDBModel, destination: str, - sources: list[str], + sources: Sequence[TSDBKey], timestamp: datetime | None = None, environment_ids: Iterable[int] | None = None, - ): + ) -> None: """ Transfer all frequency tables from the source keys to the destination key. diff --git a/src/sentry/tsdb/dummy.py b/src/sentry/tsdb/dummy.py index 35741979e36f4a..ff012c733a64f9 100644 --- a/src/sentry/tsdb/dummy.py +++ b/src/sentry/tsdb/dummy.py @@ -1,7 +1,7 @@ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime -from sentry.tsdb.base import BaseTSDB, TSDBKey, TSDBModel +from sentry.tsdb.base import BaseTSDB, TSDBItem, TSDBKey, TSDBModel class DummyTSDB(BaseTSDB): @@ -120,13 +120,14 @@ def get_most_frequent_series( def get_frequency_series( self, - model, - items: Mapping[str, Sequence[str]], - start, - end=None, - rollup=None, - environment_id=None, - ): + model: TSDBModel, + items: Mapping[TSDBKey, Sequence[TSDBItem]], + start: datetime, + end: datetime | None = None, + rollup: int | None = None, + environment_id: int | None = None, + tenant_ids: dict[str, str | int] | None = None, + ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]: self.validate_arguments([model], [environment_id]) rollup, series = self.get_optimal_rollup_series(start, end, rollup) @@ -153,7 +154,14 @@ def get_frequency_totals( results[key] = {member: 0.0 for member in members} return results - def merge_frequencies(self, model, destination, sources, timestamp=None, environment_ids=None): + def merge_frequencies( + self, + model: TSDBModel, + destination: str, + sources: Sequence[TSDBKey], + timestamp: datetime | None = None, + environment_ids: Iterable[int] | None = None, + ) -> None: environment_ids = list( (set(environment_ids) if environment_ids is not None else set()).union([None]) ) diff --git a/src/sentry/tsdb/inmemory.py b/src/sentry/tsdb/inmemory.py index 60ac8c0064eb20..f8dcb6b36683a5 100644 --- a/src/sentry/tsdb/inmemory.py +++ b/src/sentry/tsdb/inmemory.py @@ -1,10 +1,10 @@ from collections import Counter, defaultdict -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from django.utils import timezone -from sentry.tsdb.base import BaseTSDB, TSDBKey, TSDBModel +from sentry.tsdb.base import BaseTSDB, TSDBItem, TSDBKey, TSDBModel from sentry.utils.dates import to_datetime @@ -286,13 +286,14 @@ def get_most_frequent_series( def get_frequency_series( self, - model, - items: Mapping[str, Sequence[str]], - start, - end=None, - rollup=None, - environment_id=None, - ): + model: TSDBModel, + items: Mapping[TSDBKey, Sequence[TSDBItem]], + start: datetime, + end: datetime | None = None, + rollup: int | None = None, + environment_id: int | None = None, + tenant_ids: dict[str, str | int] | None = None, + ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]: self.validate_arguments([model], [environment_id]) rollup, series = self.get_optimal_rollup_series(start, end, rollup) @@ -330,7 +331,14 @@ def get_frequency_totals( return results - def merge_frequencies(self, model, destination, sources, timestamp=None, environment_ids=None): + def merge_frequencies( + self, + model: TSDBModel, + destination: str, + sources: Sequence[TSDBKey], + timestamp: datetime | None = None, + environment_ids: Iterable[int] | None = None, + ) -> None: environment_ids = (set(environment_ids) if environment_ids is not None else set()).union( [None] ) diff --git a/src/sentry/tsdb/redis.py b/src/sentry/tsdb/redis.py index 4fbf885e29921a..6a4474ac1ba355 100644 --- a/src/sentry/tsdb/redis.py +++ b/src/sentry/tsdb/redis.py @@ -15,7 +15,7 @@ from django.utils.encoding import force_bytes from redis.client import Script -from sentry.tsdb.base import BaseTSDB, IncrMultiOptions, TSDBKey, TSDBModel +from sentry.tsdb.base import BaseTSDB, IncrMultiOptions, TSDBItem, TSDBKey, TSDBModel from sentry.utils.dates import to_datetime from sentry.utils.redis import ( check_cluster_versions, @@ -915,13 +915,13 @@ def unpack_response(response: rb.Promise) -> dict[str, float]: def get_frequency_series( self, model: TSDBModel, - items: Mapping[str, Sequence[str]], + items: Mapping[TSDBKey, Sequence[TSDBItem]], start: datetime, end: datetime | None = None, rollup: int | None = None, environment_id: int | None = None, tenant_ids: dict[str, str | int] | None = None, - ) -> dict[str, list[tuple[float, dict[str, float]]]]: + ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]: self.validate_arguments([model], [environment_id]) if not self.enable_frequency_sketches: @@ -935,7 +935,7 @@ def get_frequency_series( # provided) with the original input values to compose the result. items = {k: list(members) for k, members in items.items()} - commands: dict[str, list[tuple[Script, list[str], list[str | int]]]] = {} + commands: dict[TSDBKey, list[tuple[Script, list[str], list[str | int]]]] = {} arguments = ["ESTIMATE"] + list(self.DEFAULT_SKETCH_PARAMETERS) for item_key, members in items.items(): @@ -949,7 +949,7 @@ def get_frequency_series( commands[item_key] = [(CountMinScript, ks, arguments + list(members))] - results: dict[str, list[tuple[float, dict[str, float]]]] = {} + results: dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]] = {} cluster, _ = self.get_cluster(environment_id) for _key, responses in cluster.execute_commands(commands).items(): @@ -992,10 +992,10 @@ def merge_frequencies( self, model: TSDBModel, destination: str, - sources: list[str], + sources: Sequence[TSDBKey], timestamp: datetime | None = None, environment_ids: Iterable[int] | None = None, - ): + ) -> None: ids = (set(environment_ids) if environment_ids is not None else set()).union([None]) self.validate_arguments([model], ids) @@ -1013,9 +1013,8 @@ def merge_frequencies( rollups.append((rollup, [to_datetime(item) for item in rollup_series])) for (cluster, durable), _ids in self.get_cluster_groups(ids): - exports: dict[str, list[tuple[Script, list[str], list[str]] | list[str]]] = defaultdict( - list - ) + exports: dict[TSDBKey, list[tuple[Script, list[str], list[str]] | list[str]]] + exports = defaultdict(list) for source in sources: for rollup, series in rollups: diff --git a/src/sentry/tsdb/snuba.py b/src/sentry/tsdb/snuba.py index 26ad02e9307ae1..d5af6c0c9a1ee7 100644 --- a/src/sentry/tsdb/snuba.py +++ b/src/sentry/tsdb/snuba.py @@ -28,7 +28,7 @@ from sentry.ingest.inbound_filters import FILTER_STAT_KEYS_TO_VALUES from sentry.issues.query import manual_group_on_time_aggregation from sentry.snuba.dataset import Dataset -from sentry.tsdb.base import BaseTSDB, TSDBKey, TSDBModel +from sentry.tsdb.base import BaseTSDB, TSDBItem, TSDBKey, TSDBModel from sentry.utils import outcomes, snuba from sentry.utils.dates import to_datetime from sentry.utils.snuba import ( @@ -884,14 +884,14 @@ def get_most_frequent_series( def get_frequency_series( self, - model, - items: Mapping[str, Sequence[str]], - start, - end=None, - rollup=None, - environment_id=None, - tenant_ids=None, - ): + model: TSDBModel, + items: Mapping[TSDBKey, Sequence[TSDBItem]], + start: datetime, + end: datetime | None = None, + rollup: int | None = None, + environment_id: int | None = None, + tenant_ids: dict[str, str | int] | None = None, + ) -> dict[TSDBKey, list[tuple[float, dict[TSDBItem, float]]]]: result = self.get_data( model, items,