ref: fix more types for TSDB

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.

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])
         )

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]
         )

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:

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,