ref: fix some types for sentry.api (#72641)

asottile-sentry · web-flow · commit 1a68ffaaa338 · 2024-06-13T09:43:06.000-04:00
fixes errors when BaseManager becomes typed

&lt;!-- Describe your PR here. --&gt;
diff --git a/src/sentry/api/authentication.py b/src/sentry/api/authentication.py
@@ -243,10 +243,13 @@ def authenticate_credentials(self, userid, password, request=None):
         if password:
             return None
 
+        key: ApiKeyReplica | ApiKey
         if SiloMode.get_current_mode() == SiloMode.REGION:
-            key = ApiKeyReplica.objects.filter(key=userid).last()
-            if key is None:
+            key_replica = ApiKeyReplica.objects.filter(key=userid).last()
+            if key_replica is None:
                 raise AuthenticationFailed("API key is not valid")
+            else:
+                key = key_replica
         else:
             try:
                 key = ApiKey.objects.get_from_cache(key=userid)
@@ -452,9 +455,9 @@ def accepts_auth(self, auth: list[bytes]) -> bool:
         return token_str.startswith(SENTRY_ORG_AUTH_TOKEN_PREFIX)
 
     def authenticate_token(self, request: Request, token_str: str) -> tuple[Any, Any]:
-        token = None
         token_hashed = hash_token(token_str)
 
+        token: OrgAuthTokenReplica | OrgAuthToken
         if SiloMode.get_current_mode() == SiloMode.REGION:
             try:
                 token = OrgAuthTokenReplica.objects.get(
diff --git a/src/sentry/api/endpoints/artifact_bundles.py b/src/sentry/api/endpoints/artifact_bundles.py
@@ -135,14 +135,14 @@ def delete(self, request: Request, project) -> Response:
         if bundle_id:
             error = None
 
-            project_artifact_bundles = ProjectArtifactBundle.objects.filter(
+            project_artifact_bundles_qs = ProjectArtifactBundle.objects.filter(
                 organization_id=project.organization_id,
                 artifact_bundle__bundle_id=bundle_id,
             ).select_related("artifact_bundle")
             # We group the bundles by their id, since we might have multiple bundles with the same bundle_id due to a
             # problem that was fixed in https://github.com/getsentry/sentry/pull/49836.
             grouped_bundles = defaultdict(list)
-            for project_artifact_bundle in project_artifact_bundles:
+            for project_artifact_bundle in project_artifact_bundles_qs:
                 grouped_bundles[project_artifact_bundle.artifact_bundle].append(
                     project_artifact_bundle
                 )
diff --git a/src/sentry/api/endpoints/artifact_lookup.py b/src/sentry/api/endpoints/artifact_lookup.py
@@ -62,9 +62,9 @@ def download_file(self, download_id, project: Project):
             )
             return HttpResponse({"Too many download requests"}, status=429)
 
-        file = None
+        file_m: ArtifactBundle | ReleaseFile | None = None
         if ty == "artifact_bundle":
-            file = (
+            file_m = (
                 ArtifactBundle.objects.filter(
                     id=ty_id,
                     projectartifactbundle__project_id=project.id,
@@ -76,16 +76,16 @@ def download_file(self, download_id, project: Project):
         elif ty == "release_file":
             # NOTE: `ReleaseFile` does have a `project_id`, but that seems to
             # be always empty, so using the `organization_id` instead.
-            file = (
+            file_m = (
                 ReleaseFile.objects.filter(id=ty_id, organization_id=project.organization.id)
                 .select_related("file")
                 .first()
             )
             metrics.incr("sourcemaps.download.release_file")
 
-        if file is None:
+        if file_m is None:
             raise Http404
-        file = file.file
+        file = file_m.file
 
         try:
             fp = file.getfile()
diff --git a/src/sentry/api/endpoints/group_current_release.py b/src/sentry/api/endpoints/group_current_release.py
@@ -24,13 +24,15 @@ def _get_current_release(self, group, environments):
             "release_id", flat=True
         )
 
-        release_envs = ReleaseEnvironment.objects.filter(
+        release_envs_qs = ReleaseEnvironment.objects.filter(
             release_id__in=release_projects,
             organization_id=group.project.organization_id,
         )
         if environments:
-            release_envs = release_envs.filter(environment_id__in=[env.id for env in environments])
-        release_envs = release_envs.order_by("-first_seen").values_list("release_id", flat=True)
+            release_envs_qs = release_envs_qs.filter(
+                environment_id__in=[env.id for env in environments]
+            )
+        release_envs = release_envs_qs.order_by("-first_seen").values_list("release_id", flat=True)
 
         group_releases = GroupRelease.objects.filter(
             group_id=group.id,
diff --git a/src/sentry/api/endpoints/group_details.py b/src/sentry/api/endpoints/group_details.py
@@ -230,9 +230,11 @@ def get(self, request: Request, group) -> Response:
                 data.update({"integrationIssues": integration_issues})
 
             if "sentryAppIssues" in expand:
-                external_issues = PlatformExternalIssue.objects.filter(group_id=group.id)
+                platform_external_issues = PlatformExternalIssue.objects.filter(group_id=group.id)
                 sentry_app_issues = serialize(
-                    list(external_issues), request, serializer=PlatformExternalIssueSerializer()
+                    list(platform_external_issues),
+                    request,
+                    serializer=PlatformExternalIssueSerializer(),
                 )
                 data.update({"sentryAppIssues": sentry_app_issues})
 
diff --git a/src/sentry/api/endpoints/organization_artifactbundle_assemble.py b/src/sentry/api/endpoints/organization_artifactbundle_assemble.py
@@ -60,9 +60,11 @@ def post(self, request: Request, organization) -> Response:
         if len(projects) == 0:
             return Response({"error": "You need to specify at least one project"}, status=400)
 
-        project_ids = Project.objects.filter(
-            organization=organization, status=ObjectStatus.ACTIVE, slug__in=projects
-        ).values_list("id", flat=True)
+        project_ids = list(
+            Project.objects.filter(
+                organization=organization, status=ObjectStatus.ACTIVE, slug__in=projects
+            ).values_list("id", flat=True)
+        )
         if len(project_ids) != len(projects):
             return Response({"error": "One or more projects are invalid"}, status=400)
 
diff --git a/src/sentry/api/endpoints/organization_codeowners_associations.py b/src/sentry/api/endpoints/organization_codeowners_associations.py
@@ -47,6 +47,7 @@ def get(self, request: Request, organization: Organization):
             )
         result = {}
         for pco in project_code_owners:
+            assert pco.raw is not None  # XXX: model field `raw` is nullable? seems wrong?
             associations, errors = validate_codeowners_associations(pco.raw, pco.project)
             result[pco.project.slug] = {"associations": associations, "errors": errors}
         return self.respond(result, status=status.HTTP_200_OK)
diff --git a/src/sentry/api/endpoints/organization_release_meta.py b/src/sentry/api/endpoints/organization_release_meta.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from collections import defaultdict
+from typing import TypedDict
 
 from rest_framework.request import Request
 from rest_framework.response import Response
@@ -15,6 +18,15 @@
 from sentry.models.releases.release_project import ReleaseProject
 
 
+class _ProjectDict(TypedDict):
+    id: int
+    slug: str | None
+    name: str
+    newGroups: int | None
+    platform: str | None
+    platforms: list[str]
+
+
 @region_silo_endpoint
 class OrganizationReleaseMetaEndpoint(OrganizationReleasesBaseEndpoint):
     publish_status = {
@@ -70,7 +82,7 @@ def get(self, request: Request, organization, version) -> Response:
             platforms_by_project[project_id].append(platform)
 
         # This must match what is returned from the `Release` serializer
-        projects = [
+        projects: list[_ProjectDict] = [
             {
                 "id": pr["project__id"],
                 "slug": pr["project__slug"],
diff --git a/src/sentry/api/endpoints/prompts_activity.py b/src/sentry/api/endpoints/prompts_activity.py
@@ -62,10 +62,10 @@ def get(self, request: Request, **kwargs) -> Response:
             condition = Q(feature=feature, **filters)
             conditions = condition if conditions is None else (conditions | condition)
 
-        result = PromptsActivity.objects.filter(conditions, user_id=request.user.id)
-        featuredata = {k.feature: k.data for k in result}
+        result_qs = PromptsActivity.objects.filter(conditions, user_id=request.user.id)
+        featuredata = {k.feature: k.data for k in result_qs}
         if len(features) == 1:
-            result = result.first()
+            result = result_qs.first()
             data = None if result is None else result.data
             return Response({"data": data, "features": featuredata})
         else:
diff --git a/src/sentry/api/helpers/teams.py b/src/sentry/api/helpers/teams.py
@@ -22,7 +22,7 @@ def get_teams(request, organization, teams=None):
     # do normal teams lookup based on request params
     requested_teams = set(request.GET.getlist("team", []) if teams is None else teams)
 
-    verified_ids = set()
+    verified_ids: set[int] = set()
 
     if "myteams" in requested_teams:
         requested_teams.remove("myteams")
diff --git a/src/sentry/api/serializers/models/filechange.py b/src/sentry/api/serializers/models/filechange.py
@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import Any
+
 from sentry.api.serializers import Serializer, register
 from sentry.api.serializers.models.commit import get_users_for_commits
 from sentry.models.commit import Commit
@@ -20,7 +24,7 @@ def get_attrs(self, item_list, user):
             ).values_list("id", "name")
         )
 
-        result = {}
+        result: dict[int, Any] = {}
         for item in item_list:
             commit = commits_by_id[item.commit_id]
             result[item] = {
diff --git a/src/sentry/api/serializers/models/release.py b/src/sentry/api/serializers/models/release.py
@@ -3,7 +3,7 @@
 import datetime
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from typing import Any, TypedDict, Union
+from typing import Any, NotRequired, TypedDict, Union
 
 from django.contrib.auth.models import AnonymousUser
 from django.core.cache import cache
@@ -20,6 +20,7 @@
 from sentry.models.release import Release, ReleaseStatus
 from sentry.models.releaseprojectenvironment import ReleaseProjectEnvironment
 from sentry.models.releases.release_project import ReleaseProject
+from sentry.release_health.base import ReleaseHealthOverview
 from sentry.services.hybrid_cloud.user.serial import serialize_generic_user
 from sentry.services.hybrid_cloud.user.service import user_service
 from sentry.utils import metrics
@@ -291,6 +292,17 @@ def get_users_for_authors(organization_id, authors, user=None) -> Mapping[str, A
     return results
 
 
+class _ProjectDict(TypedDict):
+    id: int
+    slug: str | None
+    name: str
+    new_groups: int | None
+    platform: str | None
+    platforms: list[str]
+    health_data: NotRequired[ReleaseHealthOverview | None]
+    has_health_data: NotRequired[bool]
+
+
 @register(Release)
 class ReleaseSerializer(Serializer):
     def __get_project_id_list(self, item_list):
@@ -491,7 +503,7 @@ def get_attrs(self, item_list, user, **kwargs):
             has_health_data = {}
 
         for pr in project_releases:
-            pr_rv = {
+            pr_rv: _ProjectDict = {
                 "id": pr["project__id"],
                 "slug": pr["project__slug"],
                 "name": pr["project__name"],
diff --git a/src/sentry/api/serializers/models/relocation.py b/src/sentry/api/serializers/models/relocation.py
@@ -2,8 +2,9 @@
 from collections.abc import Mapping, MutableMapping, Sequence
 from typing import Any
 
+from django.db.models import QuerySet
+
 from sentry.api.serializers import Serializer, register
-from sentry.db.models.manager.base import BaseManager
 from sentry.models.importchunk import BaseImportChunk, ControlImportChunkReplica, RegionImportChunk
 from sentry.models.relocation import Relocation
 from sentry.models.user import User
@@ -29,7 +30,7 @@ class RelocationMetadata:
     imported_org_ids: list[int]
 
 
-def get_all_imported_ids_of_model(chunks: BaseManager[BaseImportChunk]) -> list[int]:
+def get_all_imported_ids_of_model(chunks: QuerySet[BaseImportChunk]) -> list[int]:
     all_imported_ids = set()
     for chunk in chunks:
         all_imported_ids |= (
diff --git a/src/sentry/api/validators/project_codeowners.py b/src/sentry/api/validators/project_codeowners.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Collection, Mapping
+from collections.abc import Collection, Iterable, Mapping
 from typing import Any
 
 from django.db.models import Subquery
@@ -24,7 +24,7 @@ def validate_association_emails(
 
 def validate_association_actors(
     raw_items: Collection[str],
-    associations: Collection[ExternalActor],
+    associations: Iterable[ExternalActor],
 ) -> list[str]:
     raw_items_set = {str(item) for item in raw_items}
     # associations are ExternalActor objects
diff --git a/src/sentry/models/avatars/sentry_app_avatar.py b/src/sentry/models/avatars/sentry_app_avatar.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from collections.abc import Iterable
 from enum import Enum
 from typing import TYPE_CHECKING, ClassVar
 
@@ -25,7 +26,9 @@ def get_choices(cls):
 
 
 class SentryAppAvatarManager(BaseManager["SentryAppAvatar"]):
-    def get_by_apps_as_dict(self, sentry_apps: list[SentryApp]):
+    def get_by_apps_as_dict(
+        self, sentry_apps: Iterable[SentryApp]
+    ) -> dict[int, set[SentryAppAvatar]]:
         """
         Returns a dict mapping sentry_app_id (key) to List[SentryAppAvatar] (value)
         """
diff --git a/src/sentry/models/orgauthtoken.py b/src/sentry/models/orgauthtoken.py
@@ -149,7 +149,7 @@ def get_org_auth_token_id_from_auth(auth: object) -> int | None:
     return None
 
 
-def update_org_auth_token_last_used(auth: object, project_ids: list[int]):
+def update_org_auth_token_last_used(auth: object, project_ids: list[int]) -> None:
     org_auth_token_id = get_org_auth_token_id_from_auth(auth)
     organization_id = getattr(auth, "organization_id", None)
     if org_auth_token_id is None or organization_id is None:
diff --git a/src/sentry/models/release.py b/src/sentry/models/release.py
@@ -245,7 +245,7 @@ class Release(Model):
     # where they are "specialized" for a specific project.  The goal is to
     # later split up releases by project again.  This is for instance used
     # by the org release listing.
-    _for_project_id = None
+    _for_project_id: int | None = None
     # the user agent that set the release
     user_agent = models.TextField(null=True)
 

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,7 @@ def get(self, request: Request, organization: Organization):`
`47`	`47`	`)`
`48`	`48`	`result = {}`
`49`	`49`	`for pco in project_code_owners:`
	`50`	+ assert pco.raw is not None # XXX: model field `raw` is nullable? seems wrong?
`50`	`51`	`associations, errors = validate_codeowners_associations(pco.raw, pco.project)`
`51`	`52`	`result[pco.project.slug] = {"associations": associations, "errors": errors}`
`52`	`53`	`return self.respond(result, status=status.HTTP_200_OK)`